import logging
import cvnn
import cvnn.layers as layers
import cvnn.dataset as dp
from cvnn.cvnn_model import CvnnModel
from cvnn.data_analysis import MonteCarloAnalyzer
from cvnn.layers import Dense
from cvnn.utils import transform_to_real, randomize
from tensorflow.keras.losses import categorical_crossentropy
import pandas as pd
import copy
import sys
from openpyxl import load_workbook, Workbook
from openpyxl.worksheet.table import Table
import os
from tqdm import tqdm
import numpy as np
from pdb import set_trace
from time import sleep

logger = logging.getLogger(cvnn.__name__)


class MonteCarlo:

    def __init__(self):
        self.models = []
        self.pandas_full_data = pd.DataFrame()
        self.confusion_matrix = []
        self.monte_carlo_analyzer = MonteCarloAnalyzer()  # All at None

    def add_model(self, model):
        self.models.append(model)

    def run(self, x, y, data_summary='', polar=False, do_conf_mat=True, validation_split=0.2, validation_data=None,
            iterations=100, epochs=10, batch_size=100,
            shuffle=False, debug=False, display_freq=1, checkpoints=False):
        x, y = randomize(x, y)
        # Reset data frame
        self.pandas_full_data = pd.DataFrame()
        if do_conf_mat:
            for i in range(len(self.models)):
                self.confusion_matrix.append({"name": "model_name", "matrix": pd.DataFrame()})
        self.save_summary_of_run(self._run_summary(iterations, epochs, batch_size, shuffle),
                                 data_summary)
        if not debug:
            pbar = tqdm(total=iterations)
        for it in range(iterations):
            if debug:
                logger.info("Iteration {}/{}".format(it + 1, iterations))
            else:
                pbar.update()
            if shuffle:  # shuffle all data at each iteration
                x, y = randomize(x, y)
            for i, model in enumerate(self.models):
                val_data_fit = None
                if model.is_complex():
                    x_fit = x
                    if validation_data is not None:
                        val_data_fit = validation_data
                else:
                    x_fit = transform_to_real(x, polar=polar)
                    if validation_data is not None:
                        val_data_fit = (transform_to_real(validation_data[0], polar=polar), validation_data[1])
                test_model = copy.deepcopy(model)
                test_model.fit(x_fit, y, validation_split=validation_split, validation_data=val_data_fit,
                               epochs=epochs, batch_size=batch_size,
                               verbose=debug, display_freq=display_freq,
                               save_csv_history=True)  # Must have save_csv_history to do the montecarlo results latter
                self.pandas_full_data = pd.concat([self.pandas_full_data,
                                                   test_model.plotter.get_full_pandas_dataframe()], sort=False)
                if do_conf_mat:
                    dataset = dp.Dataset(x_fit, y, ratio=1 - validation_split)
                    self.confusion_matrix[i]["name"] = test_model.name
                    set_trace()
                    self.confusion_matrix[i]["matrix"] = pd.concat((self.confusion_matrix[i]["matrix"],
                                                                    test_model.get_confusion_matrix(dataset.x_test,
                                                                                                    dataset.y_test)))
            if checkpoints:
                # Save checkpoint in case Monte Carlo stops in the middle
                self.pandas_full_data.to_csv(self.monte_carlo_analyzer.path / "run_data.csv", index=False)
        if not debug:
            pbar.close()
        self.pandas_full_data = self.pandas_full_data.reset_index(drop=True)
        conf_mat = None
        if do_conf_mat:
            conf_mat = self.confusion_matrix
        self.monte_carlo_analyzer.set_df(self.pandas_full_data, conf_mat)

    @staticmethod
    def _run_summary(iterations, epochs, batch_size, shuffle):
        ret_str = "Monte Carlo run\n"
        ret_str += "\tIterations: {}\n".format(iterations)
        ret_str += "\tepochs: {}\n".format(epochs)
        ret_str += "\tbatch_size: {}\n".format(batch_size)
        if shuffle:
            ret_str += "\tShuffle data at each iteration\n"
        else:
            ret_str += "\tData is not shuffled at each iteration\n"
        return ret_str

    def save_summary_of_run(self, run_summary, data_summary):
        with open(str(self.monte_carlo_analyzer.path / "run_summary.txt"), "w") as file:
            file.write(run_summary)
            file.write(data_summary)
            file.write("Models:\n")
            for model in self.models:
                file.write(model.summary())


class RealVsComplex(MonteCarlo):

    def __init__(self, complex_model, capacity_equivalent=True, equiv_technique='ratio'):
        super().__init__()
        # add models
        self.add_model(complex_model)
        self.add_model(complex_model.get_real_equivalent(capacity_equivalent=capacity_equivalent,
                                                         equiv_technique=equiv_technique, name="real_network"))


# ====================================
#     Excel logging
# ====================================
def run_montecarlo(models, dataset, open_dataset=None, iterations=500,
                   epochs=150, batch_size=100, display_freq=1,
                   debug=False, polar=False, do_all=True):
    if open_dataset:
        dataset = dp.OpenDataset(open_dataset)  # Warning, open_dataset overwrites dataset

    # Monte Carlo
    monte_carlo = MonteCarlo()
    for model in models:
        monte_carlo.add_model(model)
    if not open_dataset:
        dataset.save_data(monte_carlo.monte_carlo_analyzer.path)
    monte_carlo.run(dataset.x, dataset.y, iterations=iterations,
                    epochs=epochs, batch_size=batch_size, display_freq=display_freq,
                    shuffle=False, debug=debug, data_summary=dataset.summary(), polar=polar)
    if do_all:
        monte_carlo.monte_carlo_analyzer.do_all()

    # Save data to remember later what I did.
    _save_montecarlo_log(path=str(monte_carlo.monte_carlo_analyzer.path),
                         models_names=[str(model.name) for model in models],
                         dataset_name=dataset.dataset_name,
                         num_classes=str(dataset.y.shape[1]),
                         polar_mode='Yes' if polar == 'Apple' else 'No',
                         dataset_size=str(dataset.x.shape[0]),
                         features_size=str(dataset.x.shape[1]), epochs=epochs, batch_size=batch_size
                         )
    return str(monte_carlo.monte_carlo_analyzer.path / "run_data.csv")


def run_gaussian_dataset_montecarlo(iterations=1000, m=10000, n=128, param_list=None,
                                    epochs=150, batch_size=100, display_freq=1, optimizer='Adam',
                                    shape_raw=None, activation='cart_relu', debug=False, polar=False, do_all=True,
                                    dropout=None):
    """
    This function is used to compare CVNN vs RVNN performance over statistical non-circular data.
        1. Generates a complex-valued gaussian correlated noise with the characteristics given by the inputs.
        2. It then runs a monte carlo simulation of several iterations of both CVNN and an equivalent RVNN model.
        3. Saves several files into ./log/montecarlo/date/of/run/
            3.1. run_summary.txt: Summary of the run models and data
            3.2. run_data.csv: Full information of performance of iteration of each model at each epoch
            3.3. complex_network_statistical_result.csv: Statistical results of all iterations of CVNN per epoch
            3.4. complex_network_statistical_result.csv: Statistical results of all iterations of RVNN per epoch
            3.5. (Optional) `plot/` folder with the corresponding plots generated by MonteCarloAnalyzer.do_all()
    :param iterations: Number of iterations to be done for each model
    :param m: Total size of the dataset (number of examples)
    :param n: Number of features / input vector
    :param param_list: A list of len = number of classes.
        Each element of the list is another list of len = 3 with values: [correlation_coeff, sigma_x, sigma_y]
        Example for dataset type A of paper https://arxiv.org/abs/2009.08340:
            param_list = [
                [0.5, 1, 1],
                [-0.5, 1, 1]
            ]
        Default: None will default to the example.
    :param epochs: Number of epochs for each iteration
    :param batch_size: Batch size at each iteration
    :param display_freq: Frequency in terms of epochs of when to do a checkpoint.
    :param optimizer: Optimizer to be used. Keras optimizers are not allowed.
            Can be either cvnn.optimizers.Optimizer or a string listed in opt_dispatcher.
    :param shape_raw: List of sizes of each hidden layer.
        For example [64] will generate a CVNN with one hidden layer of size 64.
        Default None will default to example.
    :param activation: Activation function to be used at each hidden layer
    :param debug:
    :param polar: Boolean weather the RVNN should receive real and imaginary part (False) or amplitude and phase (True)
    :param do_all: If true (default) it creates a `plot/` folder with the plots generated by MonteCarloAnalyzer.do_all()
    :param dropout: (float) Dropout to be used at each hidden layer. If None it will not use any dropout.
    :return: (string) Full path to the run_data.csv generated file.
        It can be used by cvnn.data_analysis.SeveralMonteCarloComparison to compare several runs.
    """
    # Get parameters
    if param_list is None:
        param_list = [
            [0.5, 1, 1],
            [-0.5, 1, 1]
        ]
    dataset = dp.CorrelatedGaussianCoeffCorrel(m, n, param_list, debug=False)
    return mlp_run_real_comparison_montecarlo(dataset, None, iterations, epochs, batch_size, display_freq, optimizer,
                                              shape_raw, activation, debug, polar, do_all, dropout=dropout)


def mlp_run_real_comparison_montecarlo(dataset, open_dataset=None, iterations=1000,
                                       epochs=150, batch_size=100, display_freq=1, optimizer='Adam',
                                       shape_raw=None, activation='cart_relu', debug=False, polar=False, do_all=True,
                                       dropout=0.5, validation_split=0.2, validation_data=None,
                                       capacity_equivalent=True, equiv_technique='ratio', do_conf_mat=True):
    if shape_raw is None:
        shape_raw = [64]
    if open_dataset:
        dataset = dp.OpenDataset(open_dataset)  # Warning, open_dataset overwrites dataset
    # Create complex network
    input_size = dataset.x.shape[1]  # Size of input
    output_size = dataset.y.shape[1]  # Size of output
    if not len(shape_raw) > 0:
        logger.error("Shape raw was empty")
        sys.exit(-1)
    layers.ComplexLayer.last_layer_output_dtype = None
    layers.ComplexLayer.last_layer_output_size = None
    if len(shape_raw) == 0:
        logger.warning("No hidden layers are used. activation and dropout will be ignored")
        shape = [
            Dense(input_size=input_size, output_size=output_size, activation='softmax_real',
                  input_dtype=np.complex64, dropout=None)
        ]
    else:  # len(shape_raw) > 0:
        shape = [Dense(input_size=input_size, output_size=shape_raw[0], activation=activation,
                       input_dtype=np.complex64, dropout=dropout)]
        for i in range(1, len(shape_raw)):
            shape.append(Dense(output_size=shape_raw[i], activation=activation, dropout=dropout))
        shape.append(Dense(output_size=output_size, activation='softmax_real', dropout=None))

    complex_network = CvnnModel(name="complex_network", shape=shape, loss_fun=categorical_crossentropy,
                                optimizer=optimizer, verbose=False, tensorboard=False)

    # Monte Carlo
    monte_carlo = RealVsComplex(complex_network,
                                capacity_equivalent=capacity_equivalent, equiv_technique=equiv_technique)
    sleep(1)  # I have error if not because not enough time passed since creation of models to be in diff folders
    monte_carlo.run(dataset.x, dataset.y, iterations=iterations,
                    epochs=epochs, batch_size=batch_size, display_freq=display_freq,
                    shuffle=False, debug=debug, data_summary=dataset.summary(), polar=polar,
                    validation_split=validation_split, validation_data=validation_data, do_conf_mat=do_conf_mat)
    if do_all:
        monte_carlo.monte_carlo_analyzer.do_all()

    # Save data to remember later what I did.
    max_epoch = monte_carlo.pandas_full_data['epoch'].max()
    epoch_filter = monte_carlo.pandas_full_data['epoch'] == max_epoch
    complex_filter = monte_carlo.pandas_full_data['network'] == "complex network"
    real_filter = monte_carlo.pandas_full_data['network'] == "real network"
    complex_last_epochs = monte_carlo.pandas_full_data[epoch_filter & complex_filter]
    real_last_epochs = monte_carlo.pandas_full_data[epoch_filter & real_filter]
    complex_median = complex_last_epochs['test accuracy'].median()
    real_median = real_last_epochs['test accuracy'].median()
    complex_median_train = complex_last_epochs['train accuracy'].median()
    real_median_train = real_last_epochs['train accuracy'].median()
    _save_rvnn_vs_cvnn_montecarlo_log(path=str(monte_carlo.monte_carlo_analyzer.path),
                                      dataset_name=dataset.dataset_name,
                                      optimizer=optimizer, loss=categorical_crossentropy,
                                      hl=str(len(shape_raw)), shape=str(shape_raw),
                                      dropout=str(dropout), num_classes=str(dataset.y.shape[1]),
                                      polar_mode='Yes' if polar else 'No',
                                      activation=activation,
                                      dataset_size=str(dataset.x.shape[0]),
                                      feature_size=str(dataset.x.shape[1]), epochs=epochs, batch_size=batch_size,
                                      winner='CVNN' if complex_median > real_median else 'RVNN',
                                      complex_median=complex_median, real_median=real_median,
                                      complex_median_train=complex_median_train, real_median_train=real_median_train,
                                      complex_iqr=complex_last_epochs['test accuracy'].quantile(.75)
                                                  - complex_last_epochs['test accuracy'].quantile(.25),
                                      real_iqr=real_last_epochs['test accuracy'].quantile(.75)
                                               - real_last_epochs['test accuracy'].quantile(.25)
                                      )
    return str(monte_carlo.monte_carlo_analyzer.path / "run_data.csv")


# ====================================
#     Excel logging
# ====================================
def _create_excel_file(fieldnames, row_data, filename=None, percentage_cols=None):
    if filename is None:
        filename = './log/montecarlo_summary.xlsx'
    file_exists = os.path.isfile(filename)
    if file_exists:
        wb = load_workbook(filename)
        ws = wb.worksheets[0]
        del ws.tables["Table1"]
    else:
        wb = Workbook()
        ws = wb.worksheets[0]
        ws.append(fieldnames)
    ws.append(row_data)
    # TODO: What if len(row_data) is longer than the dictionary? It corresponds with excel's column names?
    tab = Table(displayName="Table1", ref="A1:" + str(chr(64 + len(row_data))) + str(ws.max_row))
    if percentage_cols is not None:
        for col in percentage_cols:
            ws[col + str(ws.max_row)].number_format = '0.00%'
    ws.add_table(tab)
    wb.save(filename)


def _save_rvnn_vs_cvnn_montecarlo_log(path, dataset_name, hl, shape, dropout, num_classes, polar_mode,
                                      activation, optimizer, loss,
                                      dataset_size, feature_size, epochs, batch_size, winner,
                                      complex_median, real_median, complex_iqr, real_iqr,
                                      complex_median_train, real_median_train,
                                      comments='', filename=None):
    fieldnames = ['dataset', '# Classes', "Dataset Size", 'Feature Size', "Polar Mode", "Optimizer", "Loss",
                  'HL', 'Shape', 'Dropout', "Activation Function", 'epochs', 'batch size',
                  "Winner", "CVNN median", "RVNN median", 'CVNN IQR', 'RVNN IQR',
                  "CVNN train median", "RVNN train median",
                  'path', "cvnn version", "Comments"
                  ]
    row_data = [dataset_name, num_classes, dataset_size, feature_size, polar_mode,  # Dataset information
                optimizer, str(loss), hl, shape, dropout, activation, epochs, batch_size,  # Model information
                winner, complex_median, real_median, complex_iqr, real_iqr,  # Preliminary results
                complex_median_train, real_median_train,
                path, cvnn.__version__, comments  # Library information
                ]
    percentage_cols = ['O', 'P', 'Q', 'R', 'S', 'T']
    _create_excel_file(fieldnames, row_data, filename, percentage_cols=percentage_cols)


def _save_montecarlo_log(path, dataset_name, models_names, num_classes, polar_mode, dataset_size,
                         features_size, epochs, batch_size, filename=None):
    fieldnames = [
        'dataset', '# Classes', "Dataset Size", 'Feature Size',  # Dataset information
        'models', 'epochs', 'batch size', "Polar Mode",  # Models information
        'path', "cvnn version"  # Library information
    ]
    row_data = [
        dataset_name, num_classes, dataset_size, features_size,
        '-'.join(models_names), epochs, batch_size, polar_mode,
        path, cvnn.__version__
    ]
    _create_excel_file(fieldnames, row_data, filename)


if __name__ == "__main__":
    # Base case with one hidden layer size 64 and dropout 0.5
    run_gaussian_dataset_montecarlo(iterations=10, dropout=None)
