from bolsonaro.data.dataset_parameters import DatasetParameters
from bolsonaro.data.dataset_loader import DatasetLoader
from bolsonaro.models.model_raw_results import ModelRawResults
from bolsonaro.models.model_factory import ModelFactory
from bolsonaro.visualization.plotter import Plotter

import argparse
import pathlib
from dotenv import find_dotenv, load_dotenv
import os


if __name__ == "__main__":
    # get environment variables in .env
    load_dotenv(find_dotenv('.env.example'))

    DEFAULT_RESULTS_DIR = os.environ["project_dir"] + os.sep + 'results'
    DEFAULT_MODELS_DIR = os.environ["project_dir"] + os.sep + 'models'
    DEFAULT_EXPERIMENT_IDS = None

    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--results_dir', nargs='?', type=str, default=DEFAULT_RESULTS_DIR, help='The output directory of the results.')
    parser.add_argument('--models_dir', nargs='?', type=str, default=DEFAULT_MODELS_DIR, help='The output directory of the trained models.')
    parser.add_argument('--experiment_ids', nargs='+', type=int, default=DEFAULT_EXPERIMENT_IDS, help='Compute the results of the specified experiment id(s)')
    args = parser.parse_args()

    # Create recursively the results dir tree
    pathlib.Path(args.results_dir).mkdir(parents=True, exist_ok=True)

    """
    Use specified list of experiments ids if availabe.
    Otherwise, list all existing experiment ids from
    the specified models directory.
    """
    experiments_ids = [str(experiment_id) for experiment_id in args.experiment_ids] \
        if args.experiment_ids is not None \
        else os.listdir(args.models_dir)

    """
    Raise an error if there's no experiments ids found both
    in parameter or in models directory.
    """
    if experiments_ids is None or len(experiments_ids) == 0:
        raise ValueError("No experiment id was found or specified.")

    # Compute the plots for each experiment id
    for experiment_id in experiments_ids:
        experiment_id_path = args.models_dir + os.sep + experiment_id # models/{experiment_id}
        # Create recursively the tree results/{experiment_id}
        pathlib.Path(args.results_dir + os.sep + experiment_id).mkdir(parents=True, exist_ok=True)
        experiment_seed_root_path = experiment_id_path + os.sep + 'seeds' # models/{experiment_id}/seeds

        """
        Dictionaries to temporarly store the scalar results with the following structure:
        {seed_1: [score_1, ..., score_m], ... seed_n: [score_1, ..., score_k]}
        TODO: to complete to retreive more results
        """
        experiment_train_scores = dict()
        experiment_dev_scores = dict()
        experiment_test_scores = dict()

        experiment_weights = dict()

        # Used to check if all losses were computed using the same metric (it should be the case)
        experiment_score_metrics = list()

        # For each seed results stored in models/{experiment_id}/seeds
        for seed in os.listdir(experiment_seed_root_path):
            experiment_seed_path = experiment_seed_root_path + os.sep + seed # models/{experiment_id}/seeds/{seed}
            dataset_parameters = DatasetParameters.load(experiment_seed_path, experiment_id) # Load the dataset parameters of this experiment, with this specific seed
            dataset = DatasetLoader.load(dataset_parameters) # Load the dataset using the previously loaded dataset parameters
            extracted_forest_size_root_path = experiment_seed_path + os.sep + 'extracted_forest_size' # models/{experiment_id}/seeds/{seed}/extracted_forest_size

            # {{seed}:[]}
            experiment_train_scores[seed] = list()
            experiment_dev_scores[seed] = list()
            experiment_test_scores[seed] = list()

            experiment_weights[seed] = list()

            # List the forest sizes in models/{experiment_id}/seeds/{seed}/extracted_forest_size
            extracted_forest_sizes = os.listdir(extracted_forest_size_root_path)
            for extracted_forest_size in extracted_forest_sizes:
                # models/{experiment_id}/seeds/{seed}/extracted_forest_size/{extracted_forest_size}
                extracted_forest_size_path = extracted_forest_size_root_path + os.sep + extracted_forest_size
                # Load models/{experiment_id}/seeds/{seed}/extracted_forest_size/{extracted_forest_size}/model_raw_results.pickle file
                model_raw_results = ModelRawResults.load(extracted_forest_size_path)
                # Load [...]/model_parameters.json file and build the model using these parameters and the weights and forest from model_raw_results.pickle
                model = ModelFactory.load(dataset.task, extracted_forest_size_path, experiment_id, model_raw_results)
                # Save temporarly some raw results (TODO: to complete to retreive more results)
                # Save the scores
                experiment_train_scores[seed].append(model_raw_results.train_score)
                experiment_dev_scores[seed].append(model_raw_results.dev_score)
                experiment_test_scores[seed].append(model_raw_results.test_score)
                # Save the weights
                experiment_weights[seed].append(model_raw_results.weights)
                # Save the metric
                experiment_score_metrics.append(model_raw_results.score_metric)

        if len(set(experiment_score_metrics)) > 1:
            raise ValueError("The metrics used to compute the dev score aren't the same everytime")

        """
        Example of plot that just plots the losses computed
        on the train, dev and test subsets using a trained
        model, with the CI, and depending on the extracted
        forest size.
        """
        Plotter.plot_losses(
            file_path=args.results_dir + os.sep + experiment_id + os.sep + 'losses.png',
            all_experiment_scores=[experiment_train_scores, experiment_dev_scores, experiment_test_scores],
            x_value=extracted_forest_sizes,
            xlabel='Number of trees extracted',
            ylabel=experiment_score_metrics[0],
            all_labels=['train', 'dev', 'test'],
            title='Loss values of the trained model'
        )

        """
        TODO:
        For each dataset:
        Stage 1) A figure for the selection of the best base forest model hyperparameters (best vs default/random hyperparams)
        Stage 2) A figure for the selection of the best dataset normalization method
        Stage 3) A figure for the selection of the best combination of dataset: normalization vs D normalization vs weights normalization
        Stage 4) A figure for the selection of the most relevant subsets combination: train,dev vs train+dev,train+dev vs train,train+dev
        Stage 5) A figure for the selection of the best extracted forest size?
        Stage 6) A figure to finally compare the perf of our approach using the previous selected parameters vs the baseline vs other papers

        Stage 3)
        In all axis:
        - untrained forest
        - trained base forest (straight line cause it doesn't depend on the number of extracted trees)

        Axis 1:
        - test with forest on train+dev and OMP on train+dev
        - test with forest on train+dev and OMP on train+dev with dataset normalization
        - test with forest on train+dev and OMP on train+dev with dataset normalization + D normalization
        - test with forest on train+dev and OMP on train+dev with dataset normalization + weights normalization
        - test with forest on train+dev and OMP on train+dev with dataset normalization + D normalization + weights normalization

        Axis 2:
        - test with forest on train and OMP on dev
        - test with forest on train and OMP on dev with dataset normalization
        - test with forest on train and OMP on dev with dataset normalization + D normalization
        - test with forest on train and OMP on dev with dataset normalization + weights normalization
        - test with forest on train and OMP on dev with dataset normalization + D normalization + weights normalization

        Axis 3:
        - test with forest on train and OMP train+dev
        - test with forest on train and OMP train+dev with dataset normalization
        - test with forest on train and OMP train+dev with dataset normalization + D normalization
        - test with forest on train and OMP train+dev with dataset normalization + weights normalization
        - test with forest on train and OMP train+dev with dataset normalization + D normalization + weights normalization

        IMPORTANT: Same seeds used in all axis.
        """

        # Plot the density of the weights
        Plotter.weight_density(
            file_path=args.results_dir + os.sep + experiment_id + os.sep + 'density_weight.png',
            all_experiment_weights=experiment_weights
        )