# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/datasets_phenom.ipynb (unless otherwise specified).

__all__ = ['datasets_phenom', 'datasets_phenom', 'datasets_phenom', 'datasets_phenom']

# Cell
from andi_test.models_phenom import models_phenom

import inspect
import numpy as np
import pandas as pd

import os
import warnings

# Cell
class datasets_phenom():
    def __init__(self):
            ''' Constructor of the class '''
            self._get_models()

    def _get_models(self):
        '''Loading subclass of models'''

        available_models = inspect.getmembers(models_phenom(), inspect.ismethod)
        available_models = available_models[1:] # we need this to get rid of the init
        self.avail_models_name = [x[0] for x in available_models]
        self.avail_models_func = [x[1] for x in available_models]

# Cell
class datasets_phenom(datasets_phenom):

    def create_dataset(self,
                       T = None,
                       N_model = None,
                       dics = False,
                       path = '',
                       save = False, load = False):

        self.T = T
        self.N_model = N_model
        self.path = path
        self.dics = dics

        '''Managing dictionaries'''
        # If the input is a single dictionary, transform it to list
        if isinstance(self.dics, dict): self.dics = [self.dics]
        # if dics is False, we select trajectories from all models with default values
        if self.dics is False: self.dics = [{'model': model} for model in self.avail_models_name]


        '''Managing folders of the datasets'''
        self.save = save
        self.load = load
        if self.save or self.load:
            if self.load:
                self.save = False
            if not os.path.exists(self.path) and self.load:
                raise FileNotFoundError('The directory from where you want to load the dataset does not exist')
            if not os.path.exists(self.path) and self.save:
                os.makedirs(self.path)


        '''Create trajectories'''
        trajs, labels = self._create_trajectories()

        return trajs, labels

# Cell
class datasets_phenom(datasets_phenom):

    def _create_trajectories(self):

        for dic in self.dics:

            dataset_idx, df = self._inspect_dic(dic)

            # If the dataset does not yet exists
            if dataset_idx is False:
                # Retrive name and function of diffusion model
                model_f = self.avail_models_func[self.avail_models_name.index(dic['model'])]
                # Create dictionary with only arguments
                dic_args = dict(dic); dic_args.pop('model')

                trajs, labels = model_f(**dic_args)

                # Save the trajectories if asked
                if self.save:
                    self._save_trajectories(trajs = trajs,
                                            labels = labels,
                                            dic = dic,
                                            df = df,
                                            dataset_idx = dataset_idx,
                                            path = self.path)
            else:
                trajs, labels = self._load_trajectories(model_name = dic['model'],
                                                        dataset_idx = dataset_idx,
                                                        path = self.path)

            # Stack dataset
            try:
                data_t = np.hstack((data_t, trajs))
                data_l = np.hstack((data_l, labels))
            except:
                data_t = trajs
                data_l = labels

        return data_t, data_l

    def _save_trajectories(self, trajs, labels, dic, df, dataset_idx, path):

        file_name = path+dic['model']+'_'+str(df.shape[0])+'.npy'

        # Save information in CSV handler
        df = df.append(dic, ignore_index = True)
        df.to_csv(path+dic['model']+'.csv')

        # Save trajectories and labels
        data = np.stack((trajs, labels))
        np.save(file_name, data)

    def _load_trajectories(self, model_name, dataset_idx, path):

        file_name = path+model_name+'_'+str(dataset_idx)+'.npy'
        data = np.load(file_name)
        return data[0], data[1]


# Cell
class datasets_phenom(datasets_phenom):

    def _inspect_dic(self, dic):
        '''Checks the information of the dictionaries and managesloading/saving information.'''

        # Add time and number of trajectories information
        if self.N_model is not None:
            dic['N'] = self.N_model
        if self.T is not None:
            dic['T'] = self.T

        # Check if CSV with information of dataset exists. If not, create it
        model_m = dic['model']
        model_f = self.avail_models_func[self.avail_models_name.index(model_m)]
        # Check arguments and defaults from model's function
        args = inspect.getfullargspec(model_f).args[1:]
        defaults = inspect.getfullargspec(model_f).defaults
        try:
            df = pd.read_csv(self.path+model_m+'.csv', index_col=0)
        except:
            # convert to dataframe and add model
            df = pd.DataFrame(columns = args+['model'])

        # Assign missing keys in dic with default values
        for arg, default in zip(args, defaults):
            if arg not in dic.keys():
                dic[arg] = default

        # Check if updated keys of dic equal keys of csv.
        if set(list(df.keys())) != set(list(dic.keys())):
            raise ValueError('Input model dictionary does not match model´s properties')

        # Check if the dataset already exists:
        df_conditions = df.copy()
        df_conditions = df_conditions.where(pd.notnull(df_conditions), None) # Need in case of empty elements because deafults are None
        for key in dic:
            # We need to transform it to str to do a fair comparison between matrices (e.g. transition matrix, Ds, alphas,...)
            df_conditions = df_conditions.loc[(df_conditions[key].astype(str) == str(dic[key]))]
            if len(df_conditions.index) == 0:
                break

        # If dataset exists
        if len(df_conditions.index) > 0:
            # if the dataset exists and save was True, do not save but load
            if self.save:
                wrn_str = f'The dataset you want to save already exists (file: {model_m}_{df_conditions.index[0]}.npy). Switching to Load mode.'
                warnings.warn(wrn_str)
                dataset_idx = df_conditions.index[0]
            elif self.load:
                dataset_idx = df_conditions.index[0]
            else:
                dataset_idx = False

        # If dataset does no exists
        else:
            if self.load:
                raise ValueError('The dataset you want to load does not exist.')
            else: # If the dataset does not exist, append empty string.
                # This allows to mix saving and loading
                dataset_idx = False

        return dataset_idx, df