import logging
import os
import json
import tensorflow as tf
import pandas as pd
import numpy as np
from tqdm import tqdm
import datetime
from pdb import set_trace
from time import sleep
from openpyxl import load_workbook, Workbook
from openpyxl.worksheet.table import Table
from tensorflow.keras.losses import categorical_crossentropy
from tensorflow.keras import Model
# Own modules
import cvnn
import cvnn.layers as layers
import cvnn.dataset as dp
from cvnn.data_analysis import MonteCarloAnalyzer, Plotter, get_confusion_matrix
from cvnn.layers import ComplexDense, ComplexDropout
from cvnn.utils import transform_to_real, randomize
from cvnn.real_equiv_tools import get_real_equivalent
from cvnn.utils import median_error
# typing
from pathlib import Path
from typing import Union, Optional, List, Tuple
from cvnn.activations import t_activation
from tensorflow import data
from typing import Type

logger = logging.getLogger(cvnn.__name__)

t_path = Union[str, Path]


class MonteCarlo:

    def __init__(self):
        """
        Class that allows the statistical comparison of several models on the same dataset
        """
        self.models = []
        self.pandas_full_data = pd.DataFrame()
        self.monte_carlo_analyzer = MonteCarloAnalyzer()  # All at None
        self.output_config = {
            'plot_all': False,
            'confusion_matrix': False,
            'excel_summary': True,
            'debug': False,
            'summary_of_run': True,
            'tensorboard': False,
            'save_weights': False,
            'safety_checkpoints': False
        }

    def add_model(self, model: Type[Model]):
        """
        Adds a cvnn.CvnnModel to the list to then compare between them
        """
        self.models.append(model)

    def run(self, x, y, data_summary: str = '',
            real_cast_modes: Optional[Union[str, List[Optional[str]], Tuple[Optional[str]]]] = None,
            validation_split: float = 0.2,
            validation_data: Optional[Union[Tuple[np.ndarray, np.ndarray], data.Dataset]] = None,  # TODO: Add the tuple of validation data details.
            test_data: Optional[Union[Tuple[np.ndarray, np.ndarray], data.Dataset]] = None,
            iterations: int = 100, epochs: int = 10, batch_size: int = 100,
            shuffle: bool = True, debug: bool = False, display_freq: int = 1):
        """
        This function is used to compare all models added with `self.add_model` method.
        Runs the iteration dataset (x, y).
        1. It then runs a monte carlo simulation of several iterations of both CVNN and an equivalent RVNN model.
        2. Saves several files into ./log/montecarlo/date/of/run/
            2.1. run_summary.txt: Summary of the run models and data
            2.2. run_data.csv: Full information of performance of iteration of each model at each epoch
            2.3. <model.name>_network_statistical_result.csv: Statistical results of all iterations of CVNN per epoch
            2.4. (Optional with parameter plot_all)
                `plot/` folder with the corresponding plots generated by MonteCarloAnalyzer.do_all()
        :param x: Input data. It could be:
            - A Numpy array (or array-like), or a list of arrays (in case the model has multiple inputs).
            - A TensorFlow tensor, or a list of tensors (in case the model has multiple inputs).
            - A tf.data dataset. Should return a tuple (inputs, targets). Preferred data type (less overhead).
        :param y: Labels/Target data. Like the input data x, it could be either Numpy array(s) or TensorFlow tensor(s).
            If f x is a dataset then y will be ignored (default None)
        :param data_summary:  (String) Dataset name to keep track of it
        :param real_cast_modes: mode parameter used by cvnn.utils.transform_to_real to be used when the model to
            train is real-valued. One of the following:
            - String with the mode listed in cvnn.utils.transform_to_real to be used by all the real-valued models to
                cast complex data to real.
            - List or Tuple of strings: Same size of self.models. mode on how to cast complex data to real for each
                model in self.model.
                real_cast_modes[i] will indicate how to cast data for self.models[i] (ignored when model is complex).
        :param validation_split: Float between 0 and 1.
            Percentage of the input data to be used as test set (the rest will be use as train set)
            Default: 0.0 (No validation set).
            This input is ignored if validation_data is given.
        :param validation_data: Data on which to evaluate the loss and any model metrics at the end of each epoch.
            The model will not be trained on this data. This parameter takes precedence over validation_split.
            It can be:
                - tuple (x_val, y_val) of Numpy arrays or tensors. Preferred data type (less overhead).
                - A tf.data dataset.
        :param test_data: Data on which to evaluate the loss and any model metrics at the end of a model training.
            The model will not be trained on this data.
            If test data is not None (default) it will generate a file called `test_results.csv` with the
                statistical results from the test data.
            It can be:
                - tuple (x_val, y_val) of Numpy arrays or tensors. Preferred data type (less overhead).
                - A tf.data dataset.
        :param iterations: Number of iterations to be done for each model
        :param epochs: Number of epochs for each iteration
        :param batch_size: Batch size at each iteration
        :param display_freq: Integer (Default 1). Only relevant if validation data is provided.
            Frequency on terms of epochs before running the validation.
        :param shuffle: (Boolean) Whether to shuffle the training data before each epoch.
        :param debug:
        :return: (string) Full path to the run_data.csv generated file.
            It can be used by cvnn.data_analysis.SeveralMonteCarloComparison to compare several runs.
        """
        self.output_config['debug'] = debug
        test_data_cols = None
        if test_data is not None:
            test_data_cols = ['network'] + [n.get_config()['name'] for n in self.models[0].metrics]
        real_cast_modes = self._check_real_cast_modes(real_cast_modes)
        confusion_matrix, pbar, test_results = self._beginning_callback(iterations, epochs, batch_size,
                                                                        shuffle, data_summary, test_data_cols)
        x, y = randomize(x, y)
        w_save = []                     # TODO: Find a better method
        for model in self.models:       # ATTENTION: This will make all models have the SAME weights, not ideal
            w_save.append(model.get_weights())     # Save model weight
        # np.save(self.monte_carlo_analyzer.path / "initial_debug_weights.npy", np.array(w_save))     # TODO
        for it in range(iterations):
            if debug:
                logger.info("Iteration {}/{}".format(it + 1, iterations))
            if shuffle:  # shuffle all data at each iteration
                x, y = randomize(x, y)
            for i, model in enumerate(self.models):
                x_fit, val_data_fit, test_data_fit = self._get_fit_dataset(model.inputs[0].dtype.is_complex, x,
                                                                           validation_data, test_data,
                                                                           real_cast_modes[i])
                model.set_weights(w_save[i])
                temp_path = self.monte_carlo_analyzer.path / f"run/iteration{it}_model{i}_{model.name}"
                os.makedirs(temp_path, exist_ok=True)
                callbacks = None
                if self.output_config['tensorboard']:
                    tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=temp_path / 'tensorboard',
                                                                          histogram_freq=1)
                    callbacks = [tensorboard_callback]
                run_result = model.fit(x_fit, y, validation_split=validation_split, validation_data=val_data_fit,
                                       epochs=epochs, batch_size=batch_size,
                                       verbose=debug, validation_freq=display_freq,
                                       callbacks=callbacks)
                test_results = self._inner_callback(model, validation_data, confusion_matrix, real_cast_modes, i, run_result,
                                                    test_results, test_data_fit, temp_path)
            self._outer_callback(pbar)
        return self._end_callback(x, y, iterations, data_summary, real_cast_modes, epochs, batch_size,
                                  confusion_matrix, test_results, pbar, w_save)

    def _check_real_cast_modes(self, real_cast_modes):
        # TODO: I can check the real models input size corresponds to the real_cast_mode. And change it with a warning?
        if real_cast_modes is None:
            real_cast_modes = "real_imag"
        if isinstance(real_cast_modes, str):
            real_cast_modes = [real_cast_modes for _ in self.models]
        # I suppose then real_cast_modes is a list or tuple. Not checked TODO
        assert len(real_cast_modes) == len(self.models), "Size of real_cast_modes should be equal to the total models"
        return real_cast_modes

    @staticmethod
    def _get_fit_dataset(is_complex: bool, x, validation_data, test_data, polar):
        # TODO: This does not work with tf.Dataset.
        val_data_fit = None
        test_data_fit = None
        if is_complex:
            x_fit = x
            val_data_fit = validation_data
            test_data_fit = test_data
        else:
            x_fit = transform_to_real(x, mode=polar)
            if validation_data is not None:
                val_data_fit = (transform_to_real(validation_data[0], mode=polar), validation_data[1])
            if test_data is not None:
                test_data_fit = (transform_to_real(test_data[0], mode=polar), test_data[1])
        return x_fit, val_data_fit, test_data_fit

    # Callbacks
    def _beginning_callback(self, iterations, epochs, batch_size, shuffle, data_summary, test_data_cols):
        confusion_matrix = None
        pbar = None
        # Reset data frame
        self.pandas_full_data = pd.DataFrame()
        if not self.output_config['debug']:
            pbar = tqdm(total=iterations)
        if self.output_config['confusion_matrix']:
            confusion_matrix = []
            for mdl in self.models:
                confusion_matrix.append({"name": mdl.name, "matrix": pd.DataFrame()})
        if self.output_config['summary_of_run']:
            self._save_summary_of_run(self._run_summary(iterations, epochs, batch_size, shuffle), data_summary)
        test_results = None
        if test_data_cols is not None:
            # TODO: Consider making it a tuple so it is unmutable
            test_results = pd.DataFrame(columns=test_data_cols)
        return confusion_matrix, pbar, test_results

    def _end_callback(self, x, y, iterations, data_summary, polar, epochs, batch_size,
                      confusion_matrix, test_results, pbar, w_save):
        if not self.output_config['debug']:
            pbar.close()
        self.pandas_full_data = self.pandas_full_data.reset_index(drop=True)
        self.monte_carlo_analyzer.set_df(self.pandas_full_data)
        if self.output_config['save_weights']:
            np.save(self.monte_carlo_analyzer.path / "initial_weights.npy", np.array(w_save))
        if self.output_config['excel_summary']:
            try:  # TODO: Think this better
                num_classes = str(y.shape[1])
            except IndexError:
                num_classes = max(y) - min(y)
            self._save_montecarlo_log(iterations=iterations,
                                      dataset_name=data_summary,
                                      num_classes=num_classes, polar_mode='Yes' if polar else 'No',
                                      dataset_size=str(x.shape[0]), features_size=str(x.shape[1:]),
                                      epochs=epochs, batch_size=batch_size
                                      )
        if self.output_config['confusion_matrix']:
            if confusion_matrix is not None:
                for model_cm in confusion_matrix:
                    # If the first prediction does not predict a given class, the order will be wrong, so I sort it.
                    cm = model_cm['matrix']
                    cols = cm.columns.tolist()
                    strs = list(filter(lambda x: type(x) == str, cols))
                    ints = list(filter(lambda x: type(x) == int, cols))
                    ints.sort()
                    strs.sort()
                    cm_sorted = cm.fillna(0)[ints + strs]  # Sorted confusion matrix
                    model_cm['matrix'] = cm_sorted.groupby(cm_sorted.index).mean()
                    model_cm['matrix'].to_csv(self.monte_carlo_analyzer.path / (model_cm['name'] + "_confusion_matrix.csv"))
        if test_results is not None:
            test_results.groupby('network').describe().to_csv(self.monte_carlo_analyzer.path / ("test_results.csv"))
        if self.output_config['plot_all']:
            return self.monte_carlo_analyzer.do_all()

    def _inner_callback(self, model, validation_data, confusion_matrix, polar, model_index,
                        run_result, test_results, test_data_fit, temp_path):
        # TODO: Must have save_csv_history to do the montecarlo results latter
        # Save all results
        plotter = Plotter(path=temp_path, data_results_dict=run_result.history, model_name=model.name)
        self.pandas_full_data = pd.concat([self.pandas_full_data, plotter.get_full_pandas_dataframe()], sort=False)
        if self.output_config['confusion_matrix']:
            if validation_data is not None:  # TODO: Haven't yet done all cases here!
                if model.inputs[0].dtype.is_complex:
                    x_test, y_test = validation_data
                else:
                    x_test, y_test = (transform_to_real(validation_data[0], mode=polar), validation_data[1])
                try:
                    confusion_matrix[model_index]["matrix"] = pd.concat((confusion_matrix[model_index]["matrix"],
                                                                         get_confusion_matrix(model.predict(x_test),
                                                                                              y_test)))
                except ValueError:
                    logger.warning("ValueError: Could not do confusion matrix. No objects to concatenate.")
                    # TODO: I think confusion matrix stopped working.
            else:
                print("Confusion matrix only available for validation_data")
        if self.output_config['save_weights']:
            # model.save_weights(temp_path / "final_weights")
            np.save(temp_path / "final_weights.npy", model.get_weights())
        if test_results is not None:
            tmp_result = [model.name] + model.evaluate(x=test_data_fit[0], y=test_data_fit[1], verbose=0)
            cols = ['network'] + [n.get_config()['name'] for n in self.models[0].metrics]
            test_results = test_results.append(pd.DataFrame([tmp_result], columns=cols), ignore_index=True)
        return test_results

    def _outer_callback(self, pbar):
        if not self.output_config['debug']:
            pbar.update()
        if self.output_config['safety_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)

    # Saver functions
    def _save_montecarlo_log(self, iterations, dataset_name,  num_classes, polar_mode, dataset_size,
                             features_size, epochs, batch_size):
        fieldnames = [
            'iterations',
            'dataset', '# Classes', "Dataset Size", 'Feature Size',     # Dataset information
            'models', 'epochs', 'batch size', "Polar Mode",             # Models information
            'path', "cvnn version"                                      # Library information
        ]
        row_data = [
            iterations,
            dataset_name, num_classes, dataset_size, features_size,
            '-'.join([str(model.name) for model in self.models]), epochs, batch_size, polar_mode,
            str(self.monte_carlo_analyzer.path), cvnn.__version__
        ]
        _create_excel_file(fieldnames, row_data, './log/monte_carlo_summary.xlsx')

    @staticmethod
    def _run_summary(iterations: int, epochs: int, batch_size: int, shuffle: bool) -> str:
        ret_str = "Monte Carlo run\n"
        ret_str += f"\tIterations: {iterations}\n"
        ret_str += f"\tepochs: {epochs}\n"
        ret_str += f"\tbatch_size: {batch_size}\n"
        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):
        """
        Saves 2 files:
            - run_summary.txt: A user-friendly resume of the monte carlo run.
            - models_details.json: A full serialized version of the models.
                Contains info that lacks in the txt file like the loss or optimizer.
        """
        with open(str(self.monte_carlo_analyzer.path / "run_summary.txt"), "w") as file:
            file.write(run_summary)
            file.write(data_summary + "\n")
            file.write("Models:\n")
            for model in self.models:
                model.summary(print_fn=lambda x: file.write(x + '\n'))
        json_dict = {}
        for i, model in enumerate(self.models):
            json_dict[str(i)] = {
                'name': model.name,
                'loss': tf.keras.losses.serialize(model.loss),
                'optimizer': tf.keras.optimizers.serialize(model.optimizer),
                'layers': [tf.keras.layers.serialize(layer) for layer in model.layers]
            }
        with open(self.monte_carlo_analyzer.path / 'models_details.json', 'w') as fp:
            json.dump(str(json_dict), fp)


class RealVsComplex(MonteCarlo):
    """
    Inherits from MonteCarlo. Compares a complex model with it's real equivalent.

    Example usage:
    ```
    # Assume you already have complex data 'x' with its labels 'y'... and a Cvnn model.

    montecarlo = RealVsComplex(complex_model)
    montecarlo.run(x, y)
    ```
    """

    def __init__(self, complex_model: Type[Model], capacity_equivalent: bool = True, equiv_technique: str = 'ratio'):
        """
        :param complex_model: Complex keras model (ex: sequential)
        :param capacity_equivalent: An equivalent model can be equivalent in terms of layer neurons or
                        trainable parameters (capacity equivalent according to: https://arxiv.org/abs/1811.12351)
            - True, it creates a capacity-equivalent model in terms of trainable parameters
            - False, it will double all layer size (except the last one if classifier=True)
        :param equiv_technique: Used to define the strategy of the capacity equivalent model.
            This parameter is ignored if capacity_equivalent=False
            - 'ratio': neurons_real_valued_layer[i] = r * neurons_complex_valued_layer[i], 'r' constant for all 'i'
            - 'alternate': Method described in https://arxiv.org/abs/1811.12351 where one alternates between
                    multiplying by 2 or 1. Special case on the middle is treated as a compromise between the two.
        """
        super().__init__()
        # add models
        self.add_model(complex_model)
        self.add_model(get_real_equivalent(complex_model, capacity_equivalent=capacity_equivalent,
                                           equiv_technique=equiv_technique, name="real_network"))

    def _save_montecarlo_log(self, iterations, dataset_name, num_classes, polar_mode, dataset_size,
                             features_size, epochs, batch_size):
        max_epoch = self.pandas_full_data['epoch'].max()
        epoch_filter = self.pandas_full_data['epoch'] == max_epoch
        complex_filter = self.pandas_full_data['network'] == self.models[0].name
        real_filter = self.pandas_full_data['network'] == self.models[1].name
        complex_last_epochs = self.pandas_full_data[epoch_filter & complex_filter]
        real_last_epochs = self.pandas_full_data[epoch_filter & real_filter]
        complex_median_train = complex_last_epochs['accuracy'].median()
        real_median_train = real_last_epochs['accuracy'].median()
        complex_median = complex_last_epochs['val_accuracy'].median()
        real_median = real_last_epochs['val_accuracy'].median()
        complex_err = median_error(complex_last_epochs['accuracy'].quantile(.75),
                                   complex_last_epochs['accuracy'].quantile(.25), iterations)
        real_err = median_error(real_last_epochs['val_accuracy'].quantile(.75),
                                real_last_epochs['val_accuracy'].quantile(.25), iterations)
        fieldnames = ['iterations', 'dataset', '# Classes', "Dataset Size", 'Feature Size', "Polar Mode",
                      "Optimizer", "Loss",
                      'epochs', 'batch size',
                      "Winner", "CVNN val median", "RVNN val median", 'CVNN err', 'RVNN err',
                      "CVNN train median", "RVNN train median",
                      'path', "cvnn version"
                      ]
        row_data = [iterations, dataset_name, num_classes, dataset_size, features_size, polar_mode,
                    # Dataset information
                    str(tf.keras.losses.serialize(self.models[0].loss)),
                    str(tf.keras.optimizers.serialize(self.models[0].optimizer)),
                    epochs, batch_size,  # Model information
                    'CVNN' if complex_median > real_median else 'RVNN',
                    complex_median, real_median, complex_err, real_err,  # Preliminary results
                    complex_median_train, real_median_train,
                    str(self.monte_carlo_analyzer.path), cvnn.__version__
                    ]
        percentage_cols = ['P', 'Q', 'L', 'M']
        _create_excel_file(fieldnames, row_data, './log/rvnn_vs_cvnn_monte_carlo_summary.xlsx',
                           percentage_cols=percentage_cols)


# ====================================
#   Monte Carlo simulation methods
# ====================================
def run_montecarlo(models: List[Model], dataset: cvnn.dataset.Dataset, open_dataset: Optional[t_path] = None,
                   iterations: int = 30,
                   epochs: int = 300, batch_size: int = 100, display_freq: int = 1,
                   validation_split: float = 0.2,
                   validation_data: Optional[Union[Tuple, data.Dataset]] = None,  # TODO: Add vallidation data tuple details
                   debug: bool = False, do_conf_mat: bool = False, do_all: bool = True, tensorboard: bool = False,
                   polar: Optional[Union[str, List[Optional[str]], Tuple[Optional[str]]]] = None,
                   plot_data: bool = False) -> str:
    """
    This function is used to compare different neural networks performance.
    1. Runs simulation and compares them.
    2. Saves several files into ./log/montecarlo/date/of/run/
        2.1. run_summary.txt: Summary of the run models and data
        2.2. run_data.csv: Full information of performance of iteration of each model at each epoch
        2.3. <model_name>_statistical_result.csv: Statistical results of all iterations of each model per epoch
        2.4. (Optional) `plot/` folder with the corresponding plots generated by MonteCarloAnalyzer.do_all()

    :param models: List of cvnn.CvnnModel to be compared.
    :param dataset: cvnn.dataset.Dataset with the dataset to be used on the training
    :param open_dataset: (None)
        If dataset is saved inside a folder and must be opened, path of the Dataset to be opened. Else None (default)
    :param iterations: Number of iterations to be done for each model
    :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 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 validation_split: Float between 0 and 1.
            Percentage of the input data to be used as test set (the rest will be use as train set)
            Default: 0.0 (No validation set).
            This input is ignored if validation_data is given.
        :param validation_data: Data on which to evaluate the loss and any model metrics at the end of each epoch.
            The model will not be trained on this data. This parameter takes precedence over validation_split.
            It can be:
                - tuple (x_val, y_val) of Numpy arrays or tensors. Preferred data type (less overhead).
                - A tf.data dataset.
    :param do_conf_mat: Generate a confusion matrix based on results.
    :return: (string) Full path to the run_data.csv generated file.
        It can be used by cvnn.data_analysis.SeveralMonteCarloComparison to compare several runs.
    """
    if open_dataset:
        dataset = dp.OpenDataset(open_dataset)  # Warning, open_dataset overwrites dataset

    # Monte Carlo
    monte_carlo = MonteCarlo()
    for model in models:
        # model.training_param_summary()
        monte_carlo.add_model(model)
    if not open_dataset:
        dataset.save_data(monte_carlo.monte_carlo_analyzer.path)
    monte_carlo.output_config['excel_summary'] = False
    monte_carlo.output_config['tensorboard'] = tensorboard
    monte_carlo.output_config['confusion_matrix'] = do_conf_mat
    monte_carlo.output_config['plot_all'] = do_all
    if plot_data:
        dataset.plot_data(overlapped=True, showfig=False, save_path=monte_carlo.monte_carlo_analyzer.path,
                          library='matplotlib')
    monte_carlo.run(dataset.x, dataset.y, iterations=iterations,
                    validation_split=validation_split, validation_data=validation_data,
                    epochs=epochs, batch_size=batch_size, display_freq=display_freq,
                    shuffle=False, debug=debug, data_summary=dataset.summary(), real_cast_modes=polar)

    # Save data to remember later what I did.
    _save_montecarlo_log(iterations=iterations,
                         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=str(polar),
                         dataset_size=str(dataset.x.shape[0]),
                         features_size=str(dataset.x.shape[1]), epochs=epochs, batch_size=batch_size
                         # filename='./log/run_data.csv'
                         )
    return str("./log/run_data.csv")


def run_gaussian_dataset_montecarlo(iterations: int = 30, m: int = 10000, n: int = 128, param_list=None,
                                    epochs: int = 300, batch_size: int = 100, display_freq: int = 1,
                                    optimizer='sgd', validation_split: float = 0.2,      # TODO: Add typing here
                                    shape_raw: List[int] = None, activation: t_activation = 'cart_relu',
                                    debug: bool = False, do_all: bool = True, tensorboard: bool = False,
                                    polar: Optional[Union[str, List[Optional[str]], Tuple[Optional[str]]]] = None,
                                    capacity_equivalent: bool = True, equiv_technique: str = 'ratio',
                                    dropout: Optional[float] = None, models: Optional[List[Model]] = None,
                                    plot_data: bool = True) -> str:
    """
    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. real_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 validation_split: float between 0 and 1. Fraction of the training data to be used as validation data.
        The model will set apart this fraction of the training data, will not train on it, and will evaluate the loss
            and any model metrics on this data at the end of each epoch.
        The validation data is selected from the last samples in the x and y data provided, before shuffling.
        This argument is not supported when x is a dataset, generator or keras.utils.Sequence instance.
    :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 tensorboard: If True, it will generate tensorboard outputs to check training values.
    :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.
    :param models: List of models to be compared.
    :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.3, 1, 1],
            [-0.3, 1, 1]
        ]
    dataset = dp.CorrelatedGaussianCoeffCorrel(m, n, param_list, debug=False)
    print("Database loaded...")
    if models is not None:
        return run_montecarlo(models=models, dataset=dataset, open_dataset=None,
                              iterations=iterations, epochs=epochs, batch_size=batch_size, display_freq=display_freq,
                              validation_split=validation_split, validation_data=None,
                              debug=debug, polar=polar, do_all=do_all, tensorboard=tensorboard, do_conf_mat=False,
                              plot_data=plot_data)
    else:
        return mlp_run_real_comparison_montecarlo(dataset=dataset, open_dataset=None, iterations=iterations,
                                                  epochs=epochs, batch_size=batch_size, display_freq=display_freq,
                                                  optimizer=optimizer, shape_raw=shape_raw, activation=activation,
                                                  debug=debug, polar=polar, do_all=do_all,
                                                  tensorboard=tensorboard,
                                                  capacity_equivalent=capacity_equivalent,
                                                  equiv_technique=equiv_technique,
                                                  dropout=dropout, validation_split=validation_split,
                                                  plot_data=plot_data)


def mlp_run_real_comparison_montecarlo(dataset: cvnn.dataset.Dataset, open_dataset: Optional[t_path] = None,
                                       iterations: int = 30,
                                       epochs: int = 300, batch_size: int = 100, display_freq: int = 1,
                                       optimizer='adam',     # TODO: Typing
                                       shape_raw=None, activation: t_activation = 'cart_relu',
                                       debug:  bool = False, do_all: bool = True,
                                       polar: Optional[Union[str, List[Optional[str]], Tuple[Optional[str]]]] = None,
                                       dropout: float = 0.5, validation_split: float = 0.2,
                                       validation_data: Optional[Union[Tuple, data.Dataset]] = None,    # TODO: Add typing of tuple
                                       capacity_equivalent: bool = True, equiv_technique: str = 'ratio',
                                       shuffle: bool = False, tensorboard: bool = False,
                                       plot_data: bool = True) -> str:
    """
    This function is used to compare CVNN vs RVNN performance over any dataset.
    1. Automatically creates two Multi-Layer Perceptrons (MLP), one complex and one real.
    2. Runs simulation and compares them.
    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. real_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 dataset: cvnn.dataset.Dataset with the dataset to be used on the training
    :param open_dataset: (None)
        If dataset is saved inside a folder and must be opened, path of the Dataset to be opened. Else None (default)
    :param iterations: Number of iterations to be done for each model
    :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.
    :param validation_split: Float between 0 and 1.
            Percentage of the input data to be used as test set (the rest will be use as train set)
            Default: 0.0 (No validation set).
            This input is ignored if validation_data is given.
        :param validation_data: Data on which to evaluate the loss and any model metrics at the end of each epoch.
            The model will not be trained on this data. This parameter takes precedence over validation_split.
            It can be:
                - tuple (x_val, y_val) of Numpy arrays or tensors. Preferred data type (less overhead).
                - A tf.data dataset.
    :param capacity_equivalent: An equivalent model can be equivalent in terms of layer neurons or
                        trainable parameters (capacity equivalent according to: https://arxiv.org/abs/1811.12351)
            - True, it creates a capacity-equivalent model in terms of trainable parameters
            - False, it will double all layer size (except the last one if classifier=True)
    :param equiv_technique: Used to define the strategy of the capacity equivalent model.
        This parameter is ignored if capacity_equivalent=False
        - 'ratio': neurons_real_valued_layer[i] = r * neurons_complex_valued_layer[i], 'r' constant for all 'i'
        - 'alternate': Method described in https://arxiv.org/abs/1811.12351 where one alternates between
                multiplying by 2 or 1. Special case on the middle is treated as a compromise between the two.
    :param shuffle: TODO
    :return: (string) Full path to the run_data.csv generated file.
        It can be used by cvnn.data_analysis.SeveralMonteCarloComparison to compare several runs.
    """
    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
    shape = [
        layers.ComplexInput(input_shape=input_size)
    ]
    if len(shape_raw) == 0:
        logger.warning("No hidden layers are used. activation and dropout will be ignored")
        shape.append(
            ComplexDense(units=output_size, activation='softmax_real', input_dtype=np.complex64)
        )
    else:  # len(shape_raw) > 0:
        for s in shape_raw:
            shape.append(ComplexDense(units=s, activation=activation))   # Add dropout!
            if dropout is not None:
                shape.append(ComplexDropout(rate=dropout))
        shape.append(ComplexDense(units=output_size, activation='softmax_real'))

    complex_network = tf.keras.Sequential(shape, name="complex_network")
    complex_network.compile(optimizer=optimizer, loss=tf.keras.losses.CategoricalCrossentropy(), metrics=['accuracy'])

    # Monte Carlo
    monte_carlo = RealVsComplex(complex_network,
                                capacity_equivalent=capacity_equivalent, equiv_technique=equiv_technique)
    monte_carlo.output_config['tensorboard'] = tensorboard
    # monte_carlo.output_config['confusion_matrix'] = do_conf_mat
    monte_carlo.output_config['plot_all'] = do_all
    monte_carlo.output_config['excel_summary'] = False
    if plot_data:
        dataset.plot_data(overlapped=True, showfig=False, save_path=monte_carlo.monte_carlo_analyzer.path,
                          library='matplotlib')
    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=shuffle, debug=debug, data_summary=dataset.summary(), real_cast_modes=polar,
                    validation_split=validation_split, validation_data=validation_data)

    # 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_train = complex_last_epochs['accuracy'].median()
    real_median_train = real_last_epochs['accuracy'].median()
    try:
        complex_median = complex_last_epochs['val_accuracy'].median()
        real_median = real_last_epochs['val_accuracy'].median()
        complex_err = median_error(complex_last_epochs['val_accuracy'].quantile(.75),
                                   complex_last_epochs['val_accuracy'].quantile(.25), iterations)
        real_err = median_error(real_last_epochs['val_accuracy'].quantile(.75),
                                real_last_epochs['val_accuracy'].quantile(.25), iterations)
        winner = 'CVNN' if complex_median > real_median else 'RVNN'
    except KeyError:
        complex_median = None
        real_median = None
        complex_err = median_error(complex_last_epochs['accuracy'].quantile(.75),
                                   complex_last_epochs['accuracy'].quantile(.25), iterations)
        real_err = median_error(real_last_epochs['accuracy'].quantile(.75),
                                real_last_epochs['accuracy'].quantile(.25), iterations)
        if complex_median_train > real_median_train:
            winner = 'CVNN'
        elif complex_median_train == real_median_train:
            winner = None
        else:
            winner = 'RVNN'
    _save_rvnn_vs_cvnn_montecarlo_log(
        iterations=iterations,
        path=str(monte_carlo.monte_carlo_analyzer.path),
        dataset_name=dataset.dataset_name,
        optimizer=str(complex_network.optimizer.__class__),
        loss=str(complex_network.loss.__class__),
        hl=str(len(shape_raw)), shape=str(shape_raw),
        dropout=str(dropout), num_classes=str(dataset.y.shape[1]),
        polar_mode=str(polar),
        activation=activation,
        dataset_size=str(dataset.x.shape[0]), feature_size=str(dataset.x.shape[1]),
        epochs=epochs, batch_size=batch_size,
        winner=winner,
        complex_median=complex_median, real_median=real_median,
        complex_median_train=complex_median_train, real_median_train=real_median_train,
        complex_err=complex_err, real_err=real_err,
        filename='./log/mlp_montecarlo_summary.xlsx'
    )
    return str(monte_carlo.monte_carlo_analyzer.path / "run_data.csv")


# ====================================
#     Excel logging
# ====================================
def _create_excel_file(fieldnames: List[str], row_data: List, filename: Optional[t_path] = None,
                       percentage_cols: Optional[List[str]] = 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(iterations, 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_err, real_err,
                                      complex_median_train, real_median_train,
                                      comments='', filename=None):
    fieldnames = ['iterations', 'dataset', '# Classes', "Dataset Size", 'Feature Size', "Polar Mode", "Optimizer",
                  "Loss",
                  'HL', 'Shape', 'Dropout', "Activation Function", 'epochs', 'batch size',
                  "Winner", "CVNN median", "RVNN median", 'CVNN err', 'RVNN err',
                  "CVNN train median", "RVNN train median",
                  'path', "cvnn version", "Comments"
                  ]
    row_data = [iterations, 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_err, real_err,  # Preliminary results
                complex_median_train, real_median_train,
                path, cvnn.__version__, comments  # Library information
                ]
    percentage_cols = ['P', 'Q', 'R', 'S', 'T', 'U']
    _create_excel_file(fieldnames, row_data, filename, percentage_cols=percentage_cols)


def _save_montecarlo_log(iterations, path, dataset_name, models_names, num_classes, polar_mode, dataset_size,
                         features_size, epochs, batch_size, filename=None):
    fieldnames = [
        'iterations',
        'dataset', '# Classes', "Dataset Size", 'Feature Size',     # Dataset information
        'models', 'epochs', 'batch size', "Polar Mode",             # Models information
        'path', "cvnn version"                                      # Library information
    ]
    row_data = [
        iterations,
        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=0.5)
