import numpy as np
import pandas as pd

import hex3_chen_old as hex3_chen
from g_process import GPModel
from optimal_u_constrained.hex3_gen_u_optim import load_data


def measurements_to_array(u, meas, d_order, meas_config, meas_set):
    # Distubances
    if isinstance(meas, pd.DataFrame):
        meas = meas.to_dict('records')

    if isinstance(u, pd.DataFrame):
        u_u0 = u['u0']
        u_u1 = u['u1']
    else:
        u_u0 = u[0]
        u_u1 = u[1]

    # Ensure keys keep order
    # if d_order is None:
    #     d_order = list(d[0].keys())
    #     # Remove the heat loss. Otherwise, random noise will be applied to 0 values, degrades performance.
    #     d_order.remove('h1')
    #     d_order.remove('h2')
    #     d_order.remove('h3')
    #     d_order.remove('Ts')

    d_keys = d_order.copy()
    # d_keys.remove('T1')
    # d_keys.remove('T2')
    # d_keys.remove('T3')
    toggle = 'alpha1' in d_keys

    if toggle:
        d_keys.remove('alpha1')
        d_keys.remove('alpha2')

    # Inputs
    max_temp = 135
    if isinstance(meas, np.ndarray):

        if toggle:
            meas = meas[:-2]

        meas_copy = meas.tolist()
        if meas_set < 0:
            if meas_config == 'txv only':
                for idx, key in enumerate(['T1', 'T2', 'T3']):
                    del meas_copy[hex3_chen.meas_sets[meas_set].index(key)-idx]
            meas_copy = np.array(meas_copy)

            if meas_config in ('both', 'txv only'):
                actual_t1v = max_temp - meas[hex3_chen.meas_sets[meas_set].index('T1')]
                actual_t2v = max_temp - meas[hex3_chen.meas_sets[meas_set].index('T2')]
                actual_t3v = max_temp - meas[hex3_chen.meas_sets[meas_set].index('T3')]

                meas_copy = np.append(meas_copy, [ actual_t1v ])
                meas_copy = np.append(meas_copy, [ actual_t2v ])
                meas_copy = np.append(meas_copy, [ actual_t3v ])
        #
        # meas_copy = meas.copy()
        # # Remove T1, T2, T3 if selected
        # if meas_config == 'txv only':
        #     for idx, key in enumerate(['T1', 'T2', 'T3']):
        #         del meas_copy[hex3_chen.meas_sets[meas_set].index(key) - idx]

        if toggle:
            # Move
            X = np.array([np.append(np.array([u_u0, u_u1]), meas_copy)])
        else:
            X = np.array([meas_copy])

    else:
        raise Exception('Unsupported code path')
    # X = [*([u_u0, u_u1] * toggle), *[meas[k] for k in d_keys]]
    # if meas_set < 0:
    #     X.append(max_temp - meas['T1'])
    #     X.append(max_temp - meas['T2'])
    #     X.append(max_temp - meas['T3'])
    #     # X = np.append(meas, [hex3_chen.Ti_max - meas['T2']])
    #     # X = np.append(meas, [hex3_chen.Ti_max - meas['T3']])
    #     # print(meas)
    # X = np.array([X], dtype=np.float)
    return X


def disturbance_to_array(u, d, d_order):
    # Distubances
    if isinstance(d, pd.DataFrame):
        d = d.to_dict('records')

    # Optimal u values

    # This is u_chen
    if isinstance(u, pd.DataFrame):
        u_u0 = u['u0']
        u_u1 = u['u1']
    else:
        u_u0 = u[0]
        u_u1 = u[1]

    # Ensure keys keep order
    if d_order is None:
        d_order = list(d[0].keys())
        # Remove the heat loss. Otherwise, random noise will be applied to 0 values, degrades performance.
        d_order.remove('h1')
        d_order.remove('h2')
        d_order.remove('h3')
        d_order.remove('Ts')

    d_keys = d_order.copy()
    toggle = 'alpha1' in d_keys

    if toggle:
        d_keys.remove('alpha1')
        d_keys.remove('alpha2')

    # Inputs
    X = [*([u_u0, u_u1] * toggle), *[d[k] for k in d_keys]]
    X = np.array([X], dtype=np.float)
    return X


def preprocess(u: pd.DataFrame, d: [pd.DataFrame, np.ndarray], d_order=None, u_input=None, add_noise=False):
    # Distubances
    if isinstance(d, pd.DataFrame):
        d = d.to_dict('records')

    # Optimal u values

    # This is u_chen
    if isinstance(u, pd.DataFrame):
        u_u0 = u['uc0']
        u_u1 = u['uc1']
    else:
        u_u0 = u[:, 0]
        u_u1 = u[:, 1]

    X = []

    # Ensure keys keep order
    if d_order is None:
        d_order = list(d[0].keys())
        # Remove the heat loss. Otherwise, random noise will be applied to 0 values, degrades performance.
        d_order.remove('h1')
        d_order.remove('h2')
        d_order.remove('h3')
        d_order.remove('Ts')

    # Random valve inputs
    if isinstance(u_input, pd.DataFrame):
        u_in0 = u_input['ur0']
        u_in1 = u_input['ur1']
    elif isinstance(u_input, np.ndarray):
        u_in0 = u_input[:, 0]
        u_in1 = u_input[:, 1]
    else:
        raise Exception('This should not be reached')

    d_keys = d_order.copy()
    toggle = 'alpha1' in d_keys

    if toggle:
        d_keys.remove('alpha1')
        d_keys.remove('alpha2')
    # Inputs

    np.random.seed(1)
    for idx, (u0, u1, d_row) in enumerate(zip(u_in0, u_in1, d)):
        x = [*([u0, u1] * toggle), *[d_row[k] + add_noise * np.random.normal(0, 1) for k in d_keys]]
        # x.extend([d_row['T1v'], d_row['T2v'], d_row['T3v']])
        X.append(x)

    X = np.array(X)
    # Optimal u values
    Y = np.array(list(zip(u_u0, u_u1)))

    return X, Y, d_order


def generate_meas_set_data(mset, u, d, meas_config):
    d = d.to_dict('records')

    if isinstance(u, pd.DataFrame):
        u = np.array(list(zip(u['ur0'], u['ur1'])))

    results = []

    for idx, (u_, d_) in enumerate(zip(u, d)):
        # print(u_)
        # print(d_)
        result = hex3_chen.output_meas(mset, u_, d_)
        if not result['success']:
            print('Failed, small step in u')
            print(u_)

        res_dict = {}
        for idx, key in enumerate(hex3_chen.meas_sets[mset]):
            res_dict[key] = result['y'][idx]

        if meas_config in ('both', 'txv only'):
            res_dict['T1v'] = hex3_chen.Ti_max - result['x'][hex3_chen.x_vars.index('T1')]
            res_dict['T2v'] = hex3_chen.Ti_max - result['x'][hex3_chen.x_vars.index('T2')]
            res_dict['T3v'] = hex3_chen.Ti_max - result['x'][hex3_chen.x_vars.index('T3')]

        if meas_config == 'txv only':
            for key in ['T1', 'T2', 'T3']:
                del res_dict[key]

        results.append(res_dict)
    return np.array(results, dtype=dict)


class InputConstrainedController:
    def __init__(self, training_database, meas_set=None, noise=False, meas_config='txv only'):
        train_data = load_data(training_database)

        if meas_set is not None:
            # hex3_chen.meas_sets[2].exted(['T1', ...])
            d_keys = hex3_chen.meas_sets[meas_set].copy()

            if meas_config == 'txv only':
                d_keys.remove('T1')
                d_keys.remove('T2')
                d_keys.remove('T3')
            if meas_config in ('txv only', 'both'):
                d_keys.extend(['T1v', 'T2v', 'T3v'])

            ur = train_data['u_rand']
            uc = train_data['u_chen']
            u1 = np.array(list(zip(ur['ur0'], ur['ur1'])))
            u2 = np.array(list(zip(uc['uc0'], uc['uc1'])))
            # u3 = u1
            diff = u2 - u1
            split = int(0.70 * len(u1))
            u3 = np.concatenate((u1[:split] + diff[:split] * 0.98, u1[split:] + diff[split:] * 0.5))

            disturbances_train = generate_meas_set_data(meas_set, u3, train_data['d'], meas_config=meas_config)

        else:
            raise Exception('All distrubances not supported for this controller.')
            # d_keys = None
            # disturbances_train = train_data['d']
            # u3 = train_data['u_rand']

        X_train, Y_train, key_order = preprocess(train_data['u_chen'], disturbances_train, d_order=d_keys,
                                                 u_input=u3, add_noise=noise)
        # print(X_train[0])
        # print(Y_train[0])

        self.outputs = len(Y_train[0])
        self.real_cv = np.empty((0, 2), int)
        self.t_samples = len(Y_train)

        self.meas_config = meas_config
        self.meas_set = meas_set
        self.d_keys = key_order

        self.gp = GPModel()
        self.gp.start()
        self.gp.train(X_train, Y_train, num_restarts=1, normalize_y=True)

    def __str__(self):
        return 'uc2'

    def plot_init(self, axes):
        lines = []
        rlines = []
        for idx, ax in enumerate(axes):
            ax.set_ylabel(f'Valve_{idx + 1} [-]')
            ax.set_xlabel('Time [-]')
            ax.set_ylim([0.25, 0.35])
            ax.set_xlim([0, 1])
            line, = ax.plot([], [], label='Predicted')
            lines.append(line)
            rline, = ax.plot([], [], '--', label='Optimal')
            rlines.append(rline)
            ax.legend()
        return lines, rlines

    def gen_real_cv(self, u, u_opt, d):
        self.real_cv = np.append(self.real_cv, np.array([u_opt]), axis=0)

    def plot_prediction(self, axes, lines, time, variables, real_lines):
        # print(variables)
        for idx, (ax, line) in enumerate(zip(axes, lines)):
            ylim = ax.get_ylim()
            # print(ylim)
            if not (ylim[0] <= variables[-1][idx] <= ylim[1]):
                ax.set_ylim([np.min(variables[:, idx]) - 0.05, np.max(variables[:, idx]) + 0.05])
            ylim = ax.get_ylim()
            if not (ylim[0] <= self.real_cv[-1][idx] <= ylim[1]):
                ax.set_ylim([np.min(self.real_cv[:, idx]) - 0.05, np.max(self.real_cv[:, idx]) + 0.05])

            ax.set_xlim([0, len(time)])
            line.set_xdata(time)
            line.set_ydata(variables[:, idx])
            real_lines[idx].set_xdata(time)
            real_lines[idx].set_ydata(self.real_cv[:, idx])

    def predict(self, u, measurement):
        # print('.-'*30)
        # print(u, measurement)
        try:
            if self.meas_set is not None:
                X = measurements_to_array(u, measurement, self.d_keys,
                                          meas_set=self.meas_set, meas_config=self.meas_config)
            else:
                X = disturbance_to_array(u, measurement, self.d_keys)
                # print(X)
        except:
            self.close()
            raise

        u_old = u
        alpha = 0.8
        u, u_var, _ = self.gp.test(X)  # print(f'Before {u} | grad {gradients}')

        u = u[0]

        C11u, C12u, C13u = u_old[0], u_old[0], u_old[0]
        C21u, C22u, C23u, C24u = u_old[1], u_old[1], u_old[1], u_old[1]

        N2 = np.array([0.0000, 1])
        N3 = np.array([1, 0.0000])
        N4 = np.array([-0.7071, 0.7071])

        C11u = C11u - 0.004 * (hex3_chen.Ti_max - measurement[hex3_chen.meas_sets[self.meas_set].index('T1')])
        C12u = C12u + 0.002 * (hex3_chen.Ti_max - measurement[hex3_chen.meas_sets[self.meas_set].index('T3')])
        C13u = u[0]  # C13u - 0.001 * N3 @ u

        C21u = C21u - 0.004 * (hex3_chen.Ti_max - measurement[hex3_chen.meas_sets[self.meas_set].index('T2')])
        C22u = C22u + 0.002 * (hex3_chen.Ti_max - measurement[hex3_chen.meas_sets[self.meas_set].index('T3')])
        C23u = u[1]  # C23u - 0.001 * N4 @ u

        # print(C11u)
        # print(C12u)
        # print(C13u)
        u1 = np.max((C11u, np.min((C12u, C13u))))
        u2 = np.max((C21u, np.min((C22u, C23u))))
        u = np.array([u1, u2])

        return u, u

    def __del__(self):
        self.close()

    def close(self):
        self.gp.exit()


if __name__ == '__main__':
    # train = f'..\\gradient\\datasets\\train_gradient{500}.csv'
    # controller = GradientController(train)
    # controller.close()
    # print('Controller closed')
    pass