"""In this file, a neural network is developed to fit the dataset generated by Tutor.
Depending on the observation space and action space, the tutor model can/has to be
adjusted.

The Junior model returns a one-hot encoded output, based on the number of actions.

Credit: The junior is a more general approach of the original code, see
@https://github.com/AsprinChina/L2RPN_NIPS_2020_a_PPO_Solution
"""
import logging
from pathlib import Path
from typing import Union, Optional, Tuple, List

import grid2op.Environment
import numpy as np
from sklearn.base import BaseEstimator
from sklearn.preprocessing import StandardScaler

from curriculumagent.junior.junior_student_pytorch import JuniorParamTorch, JuniorTorch
from curriculumagent.junior.junior_student_tensorflow import JuniorParamTF, JuniorTF


class Junior:
    """
    Create the overall Junior which can train either a torch or a tensorflow model.
    """

    def __init__(self,
                 action_space_file: Union[Path, List[Path], str],
                 run_with_tf: bool = True,
                 config: Union[JuniorParamTF, JuniorParamTorch] = {},
                 seed: Optional[int] = None,
                 scaler: Optional[StandardScaler] = None, #ToDo ask about scaler
                 run_nni: bool = False
                 ):
        """Constructor of the Junior with either Pytorch Lightning or with Keras-Tensorflow.
        The junior can receive hyperparameters in the config to create a more flexible model.

        Furhter, it is possible to create a hyperparameter search with tune ore with NNI.

        Except setting all variables, the init additionally requires the size of the train set and optionally
        the number of epochs.

        Note:
            Pass epochs, learning_rate, batchsize, layer_size in the config.

        Args:
            run_with_tf: Should the agent run with tensorflow. If set to False, torch is selected instead.
            action_space_file: Action Space file that was used for the Tutor training. This is needed to extract the
            correct shape of the Junior model.
            config: Dictionary containing the correct input for the hyperparameters.
            seed: Optional Seed to reproduce results.
            run_nni: Whether NNI is used. If True, then a specific callback is added.

        Returns:
            None.
        """

        # Set self actions to a list to iterate for later
        if run_with_tf:
            self.junior_model = JuniorTF(action_space_file=action_space_file,
                                         config=config,
                                         seed=seed,
                                         run_nni=run_nni)
        else:
            self.junior_model = JuniorTorch(action_space_file=action_space_file,
                                            config=config,
                                            seed=seed,
                                            run_nni=run_nni)
        self.run_with_tf = run_with_tf
        self.actions = None
        if isinstance(action_space_file, Path):
            assert action_space_file.is_file()
            self.actions = np.load(str(Path(action_space_file)))

        elif isinstance(action_space_file, list):
            for act_path in action_space_file:
                assert act_path.is_file()
            self.actions = np.concatenate([np.load(str(act_path)) for act_path in action_space_file], axis=0)

        self.scaler = scaler if isinstance(scaler, BaseEstimator) else None

    def train(self,
              run_name: str,
              dataset_path: Path,
              target_model_path: Path,
              dataset_name: str = "junior_dataset",
              epochs: int = 1000,
              **kwargs
              ) -> dict:
        """Loads the dataset and then trains the JuniorModel with either Tensorflow or torch.


        Args:
            run_name: The name of the training run.
            dataset_path: Path to the dataset files.
            target_model_path: Path, where to save the model.
            action_space_file: Optional action space file of the tutor.
            dataset_name: The name of the dataset in {dataset_name}_train.npz.
            epochs: The number of epochs to train.
            kwargs: Additional arguments for either the Torch or TF junior

        Returns:
            Training history with at least the keys accuracy and val_accuracy

        """
        if not target_model_path.is_dir():
            logging.warning(f"{target_model_path} does not exists yet. Create directory")
            target_model_path.mkdir(parents=True, exist_ok=True)

        ckpt_dir = target_model_path / f"ckpt-{run_name}"

        if self.run_with_tf:
            s_train, a_train, s_validate, a_validate, s_test, a_test = load_dataset(dataset_path, dataset_name)

            if self.scaler:
                s_train = self.scaler.transform(s_train)
                s_validate = self.scaler.transform(s_validate)
            # Get maximum number of actions:

            history = self.junior_model.train(
                x_train=s_train,
                y_train=a_train,
                x_validate=s_validate,
                y_validate=a_validate,
                log_dir=None,
                ckpt_dir=ckpt_dir,
                epochs=epochs,
                **kwargs
            )
            # Save model
            self.junior_model.model.save(target_model_path)

        else:
            history = self.junior_model.train(
                dataset_path=dataset_path,
                target_model_path=ckpt_dir,
                dataset_name=dataset_name,
                epochs=epochs,
                scaler=self.scaler,
                **kwargs)

        return history

    def test(self,
             dataset_path: Path,
             checkpoint_path: Optional[Path] = None,
             dataset_name: str = "junior_dataset"
             ) -> dict:
        """Test the given checkpoint against the test data set.

        Args:
            checkpoint_path: The checkpoint file to use for the conversion/input.
            dataset_path: Path to the dataset used to train the checkpoint.
            dataset_name: The name of the dataset in {dataset_name}_test.npz.

        Returns:
            Dictionary with accuracy values achieved on the testing dataset.

        """
        logging.info(f"A total of {len(self.actions)} are assumed, based on the action_space_file input.")

        if self.run_with_tf:
            _, _, _, _, s_test, a_test = load_dataset(dataset_path, dataset_name)

            if self.scaler:
                s_test = self.scaler.transform(s_test)

            accuracy = self.junior_model.test(x=s_test, y=a_test, save_path=checkpoint_path)

        else:

            if self.scaler:
                accuracy = self.junior_model.test(dataset_path=dataset_path, dataset_name=dataset_name, scaler=self.scaler, checkpoint_path=checkpoint_path)
            else:
                accuracy = self.junior_model.test(dataset_path=dataset_path, dataset_name=dataset_name, checkpoint_path=checkpoint_path)
        return accuracy


def load_dataset(
        dataset_path: Union[str, Path], dataset_name: str
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Function for loading the datasets of the Tutor run.

    Note:
        All datasets (training, validation and testing) have the same name but different endings.
        As example:
        dataset_name = data
        training: data_train.npz
        validation: data_val.npz
        test: data_test.npz

    Args:
        dataset_path: Path where to find the Tutor results.
        dataset_name: Name of the tutor results.

    Returns:
        Tuple, containing vectors (numpy arrays) with training, val and test data.

    """
    train_path = Path(dataset_path) / f"{dataset_name}_train.npz"
    train_data = np.load(train_path)

    validation_path = Path(dataset_path) / f"{dataset_name}_val.npz"
    validation_data = np.load(validation_path)

    test_path = Path(dataset_path) / f"{dataset_name}_test.npz"
    test_data = np.load(test_path)

    return (
        train_data["s_train"],
        train_data["a_train"],
        validation_data["s_validate"],
        validation_data["a_validate"],
        test_data["s_test"],
        test_data["a_test"],
    )
