"""
TensorFlow 2.x reproduction of the EQL (Equation Learner) symbolic regression
architecture from:
    "Symbolic regression for scientific discovery: an application to wind speed forecasting"
    Abdellaoui & Mehrkanoon, arXiv:2102.10570
"""
import os, sys, time, copy
import numpy as np
from collections import OrderedDict
from tqdm import tqdm
from scipy.io import loadmat
from sklearn.metrics import mean_squared_error, mean_absolute_error
import tensorflow as tf
import sympy as sp
import h5py
import matplotlib
matplotlib.use('Agg')
from matplotlib import pyplot as plt

# ---------------------------------------------------------------------------
# 1. Primitive functions
# ---------------------------------------------------------------------------
class BaseFunction:
    def __init__(self, norm=1):
        self.norm = norm
    def sp(self, x): return None
    def tf(self, x):
        z = sp.symbols('z')
        return sp.utilities.lambdify(z, self.sp(z), 'tensorflow')(x)
    def np(self, x):
        z = sp.symbols('z')
        return sp.utilities.lambdify(z, self.sp(z), 'numpy')(x)
    def name(self, x): return str(self.sp)

class Constant(BaseFunction):
    def tf(self, x): return tf.ones_like(x)
    def sp(self, x): return 1
    def np(self, x): return np.ones_like

class Identity(BaseFunction):
    def tf(self, x): return tf.identity(x) / self.norm
    def sp(self, x): return x / self.norm
    def np(self, x): return np.array(x) / self.norm

class Square(BaseFunction):
    def tf(self, x): return tf.square(x) / self.norm
    def sp(self, x): return x ** 2 / self.norm
    def np(self, x): return np.square(x) / self.norm

class Sin(BaseFunction):
    def tf(self, x): return tf.sin(x * 4 * np.pi) / self.norm
    def sp(self, x): return sp.sin(x * 4 * sp.pi) / self.norm
    def np(self, x): return np.sin(x * 4 * np.pi) / self.norm

class Sigmoid(BaseFunction):
    def sp(self, x): return 1 / (1 + sp.exp(-20 * x)) / self.norm
    def np(self, x): return 1 / (1 + np.exp(-20 * x)) / self.norm
    def name(self, x): return "sigmoid(x)"

class BaseFunction2:
    def __init__(self, norm=1.0): self.norm = norm
    def sp(self, x, y): return None
    def tf(self, x, y):
        a, b = sp.symbols('a b')
        return sp.utilities.lambdify([a, b], self.sp(a, b), 'tensorflow')(x, y)
    def np(self, x, y):
        a, b = sp.symbols('a b')
        return sp.utilities.lambdify([a, b], self.sp(a, b), 'numpy')(x, y)
    def name(self, x, y): return str(self.sp)

class Product(BaseFunction2):
    def __init__(self, norm=0.1): super().__init__(norm=norm)
    def sp(self, x, y): return x * y / self.norm

def count_double(funcs):
    return sum(1 for f in funcs if isinstance(f, BaseFunction2))

def count_inputs(funcs):
    i = 0
    for f in funcs:
        if isinstance(f, BaseFunction): i += 1
        elif isinstance(f, BaseFunction2): i += 2
    return i

# ---------------------------------------------------------------------------
# 2. Symbolic network
# ---------------------------------------------------------------------------
class SymbolicLayer:
    def __init__(self, funcs=None, initial_weight=None, variable=False, init_stddev=0.1):
        if funcs is None:
            funcs = default_func
        self.initial_weight = initial_weight
        self.W = None
        self.built = False
        if self.initial_weight is not None:
            if not variable:
                self.W = tf.Variable(self.initial_weight)
            else:
                self.W = self.initial_weight
            self.built = True
        self.output = None
        self.init_stddev = init_stddev
        self.n_funcs = len(funcs)
        self.funcs = [f.tf for f in funcs]
        self.n_double = count_double(funcs)
        self.n_single = self.n_funcs - self.n_double
        self.out_dim = self.n_funcs + self.n_double

    def build(self, in_dim):
        self.W = tf.Variable(tf.random.normal(shape=[in_dim, self.out_dim], stddev=self.init_stddev))
        self.built = True

    def __call__(self, x):
        if not self.built:
            self.build(int(x.shape[1]))
        g = tf.matmul(x, self.W)
        self.output = []
        in_i, out_i = 0, 0
        while out_i < self.n_single:
            self.output.append(self.funcs[out_i](g[:, in_i]))
            in_i += 1
            out_i += 1
        while out_i < self.n_funcs:
            self.output.append(self.funcs[out_i](g[:, in_i], g[:, in_i + 1]))
            in_i += 2
            out_i += 1
        self.output = tf.stack(self.output, axis=1)
        return self.output

    def get_weight(self): return self.W
    def set_weight(self, weight): self.W = weight


class SymbolicNet:
    def __init__(self, symbolic_depth, funcs=None, initial_weights=None,
                 initial_bias=None, variable=False, init_stddev=0.1):
        self.depth = symbolic_depth
        self.funcs = funcs
        self.shape = (None, 1)
        if initial_weights is not None:
            self.symbolic_layers = [
                SymbolicLayer(funcs=funcs, initial_weight=initial_weights[i], variable=variable)
                for i in range(self.depth)
            ]
            if not variable:
                self.output_weight = tf.Variable(initial_weights[-1])
            else:
                self.output_weight = initial_weights[-1]
        else:
            if isinstance(init_stddev, list):
                self.symbolic_layers = [SymbolicLayer(funcs=funcs, init_stddev=init_stddev[i]) for i in range(self.depth)]
            else:
                self.symbolic_layers = [SymbolicLayer(funcs=funcs, init_stddev=init_stddev) for _ in range(self.depth)]
            self.output_weight = tf.Variable(tf.random.uniform(shape=(self.symbolic_layers[-1].n_funcs, 1)))

    def build(self, input_dim):
        in_dim = input_dim
        for i in range(self.depth):
            self.symbolic_layers[i].build(in_dim)
            in_dim = self.symbolic_layers[i].n_funcs

    def __call__(self, input):
        self.shape = (int(input.shape[1]), 1)
        h = input
        for i in range(self.depth):
            h = self.symbolic_layers[i](h)
        h = tf.matmul(h, self.output_weight)
        return h

    def get_weights(self):
        return [self.symbolic_layers[i].W for i in range(self.depth)] + [self.output_weight]

    def set_weights(self, weights):
        for i in range(self.depth):
            self.symbolic_layers[i].W = weights[i]
        self.output_weight = weights[self.depth]


# ---------------------------------------------------------------------------
# 3. Regularization
# ---------------------------------------------------------------------------
def l12_smooth(input_tensor, a=0.05):
    if isinstance(input_tensor, list):
        return sum(l12_smooth(t, a) for t in input_tensor)
    smooth_abs = tf.where(
        tf.abs(input_tensor) < a,
        tf.pow(input_tensor, 4) / (-8 * a ** 3) + tf.square(input_tensor) * 3 / 4 / a + 3 * a / 8,
        tf.abs(input_tensor)
    )
    return tf.reduce_sum(tf.sqrt(smooth_abs))

# ---------------------------------------------------------------------------
# 4. Pretty print
# ---------------------------------------------------------------------------
def apply_activation(W, funcs, n_double=0):
    W = sp.Matrix(W)
    if n_double == 0:
        for i in range(W.shape[0]):
            for j in range(W.shape[1]):
                W[i, j] = funcs[j](W[i, j])
    else:
        W_new = W.copy()
        out_size = len(funcs)
        for i in range(W.shape[0]):
            in_j, out_j = 0, 0
            while out_j < out_size - n_double:
                W_new[i, out_j] = funcs[out_j](W[i, in_j])
                in_j += 1
                out_j += 1
            while out_j < out_size:
                W_new[i, out_j] = funcs[out_j](W[i, in_j], W[i, in_j + 1])
                in_j += 2
                out_j += 1
        for _ in range(n_double):
            W_new.col_del(-1)
        W = W_new
    return W


def sym_pp(W_list, funcs, var_names, threshold=1e-3, n_double=0):
    vars = [sp.Symbol(v) if isinstance(v, str) else v for v in var_names]
    expr = sp.Matrix(vars).T
    for W in W_list:
        W = filter_mat(sp.Matrix(W), threshold=threshold)
        expr = expr * W
        expr = apply_activation(expr, funcs, n_double=n_double)
    return expr


def last_pp(eq, W):
    return eq * filter_mat(sp.Matrix(W))


def network(weights, funcs, var_names, threshold=1e-3):
    n_double = count_double(funcs)
    funcs = [f.sp for f in funcs]
    expr = sym_pp(weights[:-1], funcs, var_names, threshold=threshold, n_double=n_double)
    expr = last_pp(expr, weights[-1])
    expr = expr[0, 0]
    return expr


def filter_mat(mat, threshold=0.01):
    for i in range(mat.shape[0]):
        for j in range(mat.shape[1]):
            if abs(mat[i, j]) < threshold:
                mat[i, j] = 0
    return mat

# ---------------------------------------------------------------------------
# 5. Data helpers
# ---------------------------------------------------------------------------
def tensor_to_matrix(tensor):
    number_features = tensor.shape[0]
    number_cities = tensor.shape[1]
    len_dates = tensor.shape[2]
    matrix_for_scaling = np.ones((len_dates, number_cities * number_features))
    for i in range(number_cities):
        for j in range(len_dates):
            features_city = tensor[:, i, j]
            matrix_for_scaling[j, i * number_features:(i + 1) * number_features] = features_city
    return matrix_for_scaling


def get_x_right_format(phase, steps_ahead, feature):
    steps_ahead_to_step = {6: "1", 12: "2", 18: "3", 24: "4"}
    filename = "{}/step{}.mat".format(feature, steps_ahead_to_step[steps_ahead])
    data = loadmat('Denmark_data/{}'.format(filename))
    x = data["Xtr"] if phase == "train" else data["Xtest"]
    x = np.transpose(x, (0, 3, 2, 1))  # => (Dates, Features, Lags, Cities)
    x = np.transpose(x, (0, 1, 3, 2))  # => (Dates, Features, Cities, Lags)
    all_features_all_cities = x.shape[1] * x.shape[2]
    x_output = np.zeros((x.shape[0], 80))
    for i in range(x.shape[0]):
        temp_matrix = tensor_to_matrix(x[i])
        for j in range(temp_matrix.shape[0]):
            x_output[i, j * all_features_all_cities:(j + 1) * all_features_all_cities] = temp_matrix[j]
    return x_output


def generate_all_data(phase, steps_ahead, feature, city_index):
    steps_ahead_to_step = {6: "1", 12: "2", 18: "3", 24: "4"}
    filename = "{}/step{}.mat".format(feature, steps_ahead_to_step[steps_ahead])
    data = loadmat('Denmark_data/{}'.format(filename))
    if phase == "train":
        y_temp = data["Ytr"][:, city_index]
    elif phase == "test":
        y_temp = data["Ytest"][:, city_index]
    else:
        return None, None
    x = get_x_right_format(phase, steps_ahead, feature).astype("float32")
    y = y_temp.reshape(y_temp.shape[0], 1).astype('float32')
    return x, y


def generate_variable_list(alphabet_size, num_variables):
    lower_case = [chr(i + 97) for i in range(26)]
    variables = lower_case[:alphabet_size]
    prefix_index = 0
    letter_index = 0
    prefix = variables[0]
    for i in range(num_variables):
        if (i + 1) % alphabet_size == 0:
            variables.append(prefix + lower_case[letter_index])
            prefix_index += 1
            prefix = variables[prefix_index]
            letter_index = 0
        else:
            variables.append(prefix + lower_case[letter_index])
            letter_index += 1
    return variables


# ---------------------------------------------------------------------------
# 6. Utils
# ---------------------------------------------------------------------------
def create_experiment_folder(experiment_number):
    os.makedirs("ExperimentsSR/Experiment{}".format(experiment_number), exist_ok=True)


def get_experiment_number():
    if not os.path.isdir("ExperimentsSR"):
        return 1
    folders = os.listdir("ExperimentsSR/")
    nums = []
    import re
    for f in folders:
        n = re.findall(r'\d+', f)
        if n:
            nums.append(int(n[0]))
    return max(nums) + 1 if nums else 1


def create_summary_file(experiment_number):
    open("ExperimentsSR/Experiment{}/summary_experiment{}.txt".format(experiment_number, experiment_number), "w").close()


def get_folders_started():
    os.makedirs("ExperimentsSR", exist_ok=True)
    experiment_number = get_experiment_number()
    create_experiment_folder(experiment_number)
    create_summary_file(experiment_number)
    return experiment_number


def record_base_info(experiment_number, **config):
    with open("ExperimentsSR/Experiment{}/summary_experiment{}.txt".format(experiment_number, experiment_number), "a+") as f:
        for k, v in config.items():
            f.write("{}: {}\n".format(k, v))


def append_text_to_summary(experiment_number, text):
    with open("ExperimentsSR/Experiment{}/summary_experiment{}.txt".format(experiment_number, experiment_number), "a+") as f:
        f.write(text)


def plot_train_vs_validation(experiment_number, num_epochs, train_loss, validation_loss, phase):
    x = range(len(train_loss))
    plt.plot(x, train_loss, label="Training")
    plt.plot(x, validation_loss, label="Validation")
    plt.ylabel("MSE")
    plt.xlabel("Epochs")
    plt.title(phase)
    plt.legend()
    plt.savefig("ExperimentsSR/Experiment{}/train_vs_validation_{}.jpg".format(experiment_number, phase))
    plt.clf()


def plot_histogram(experiment_number, weights, phase, type_weights, a):
    plt.hist(weights)
    plt.title(phase + ", " + type_weights + ", a:{}".format(a))
    plt.savefig("ExperimentsSR/Experiment{}/hist_{}_{}.jpg".format(experiment_number, phase, type_weights))
    plt.clf()


def plot_descaled_real_vs_prediction(experiment_number, y_real, y_predicted, y_min, y_max):
    yp = y_predicted.numpy() if hasattr(y_predicted, 'numpy') else np.array(y_predicted)
    yr = y_real.numpy() if hasattr(y_real, 'numpy') else np.array(y_real)
    y_predicted_rescaled = yp * (y_max - y_min) + y_min
    y_test_rescaled = yr * (y_max - y_min) + y_min
    mse = mean_squared_error(y_predicted_rescaled, y_test_rescaled)
    mae = mean_absolute_error(y_predicted_rescaled, y_test_rescaled)
    plt.figure(figsize=(10, 8))
    plt.plot(y_test_rescaled[:2000], label="Real")
    plt.plot(y_predicted_rescaled[:2000], label="Prediction from formula")
    plt.legend()
    plt.title("MAE: {:.2e}, MSE: {:.2e}".format(mae, mse))
    plt.savefig("ExperimentsSR/Experiment{}/real_vs_prediction_phase1.jpg".format(experiment_number))
    plt.clf()
    return mae


# ---------------------------------------------------------------------------
# 7. Weights I/O
# ---------------------------------------------------------------------------
def save_weights(weights, experiment_number, phase, best_val_loss):
    if phase == "phase1":
        filename = "ExperimentsSR/Experiment{}/nt_val_weights_{:.2e}.hdf5".format(experiment_number, best_val_loss)
    elif phase == "phase2":
        filename = "ExperimentsSR/Experiment{}/phase2_val_weights_{:.2e}.hdf5".format(experiment_number, best_val_loss)
    with h5py.File(filename, "w") as f:
        for i in range(len(weights)):
            w = weights[i]
            if hasattr(w, 'numpy'):
                w = w.numpy()
            f.create_dataset('dataset{}'.format(i), data=w)


def load_weights(filename):
    weights = []
    with h5py.File(filename, "r") as f:
        for i in range(3):
            weights.append(f.get('dataset{}'.format(i))[()])
    return weights