import pickle
import datetime
import numpy as np

class Trial:
    """
    Third level: Trial. A trial has all the information, metadata and results of a run of the experiment.
    """
    def __init__(self, items:dict):
        items['time'] = datetime.datetime.now().strftime("%d-%m-%Y %H:%M")
        self.items = items
    def get_keys(self):
        return self.items.keys()
    def get_by_keys(self,keys:list[str]):
        for k in keys:
            if not k in self.items:
                raise Exception(f"Error: Key {k} is not present in trial")
        return [self.items[k] for k in keys]
    def __str__(self):
        return f'{self.items}'
    def __repr__(self):
        return self.__str__()

class Setting:
    """
    Second Level: Setting. A setting describe a configuration of hiperparemters in the experiment and has many trials.
    """
    def __init__(self, conf:dict):
        self.conf = conf
        self.trials = []
    def match(self, other_conf:dict):
        if self.conf.keys() != other_conf.keys():
            return False

        for k in self.conf.keys():
            v1,v2 = self.conf[k], other_conf[k]
            if isinstance(v1, np.ndarray) or isinstance(v2, np.ndarray):
                if not np.array_equal(v1,v2):
                    return False
            else:
                if v1!=v2:
                    return False
        return True

    def insert(self, trial):
        if len(self.trials):
            if self.trials[0].get_keys() != trial.get_keys():
                key_list = list(self.trials[0].get_keys())
                raise Exception(f"The trial you are trying to insert has different keys than the previuos ones: {key_list}")
        self.trials.append( trial )
    def get_trials(self, keys:list[str]|None = None):
        if keys is None: 
            return [t.items for t in self.trials]
        return [ t.get_by_keys(keys) for t in self.trials ]
    def flush(self):
        self.trials = []
    def __str__(self):
        return f'{self.conf}'
    def __repr__(self):
        return self.__str__()

class GlauxExp:
    """
    First level: Experiment. The experiment has many runts, sorted and stored by setting. To access a setting you must provide a configuration dictionary.
    """
    def __init__(self, filename:str, description:str|None=None):
        self.filename = filename
        try:
            with open(filename,'rb') as file:
                data = pickle.load(file)
                self.description = data['description']
                self.settings    = data['settings']
        except FileNotFoundError:
            self.description = description
            self.settings = []

    def save(self):
        with open(self.filename,'wb') as file:
            pickle.dump( {'settings':self.settings,
                          'description':self.description}, file ) 

    def __search(self, conf:dict):
        for s in self.settings:
            if s.match(conf):  return s
        return None

    def insert_trial(self, settings:dict, trial:dict):
        assert type(settings)==dict, "settings parameter must be a dict"
        assert type(trial)==dict, "trial parameter must be a dict"

        trial = Trial(trial)
        new_setting = self.__search(settings)
        if new_setting is None:   
            new_setting = Setting(settings)
            self.settings.append( new_setting )
        new_setting.insert( trial )
        self.save()

    def get_trials(self, setting_id:int, keys:list[str]|None = None ):
        assert type(setting_id)==int, "setting_id parameter must be Int"
        assert type(keys)==list or keys is None, "keys parameter must be a list of str or None"
        assert setting_id<len(self.settings),'setting_id must be a valid id for this experiment'

        return self.settings[setting_id].get_trials(keys)
    
    def flush(self, setting_id):
        assert type(setting_id)==int, "setting_id parameter must be Int"
        assert setting_id<len(self.settings),'setting_id must be a valid id for this experiment'

        cmd = input('write "ok" to confirm flush')
        if cmd == 'ok':
            self.settings[setting_id].flush()
            self.save()
            print('flushed')
        else:
            print('not flushed')
    
    def flush_all(self):
        cmd = input('write "ok" to confirm flush')
        if cmd == 'ok':
            self.settings = []
            self.save()
            print('flushed')
        else:
            print('not flushed')

    def __str__(self):
        msg = self.description + '\n'
        for idx, setting in enumerate(self.settings):
            msg += f'id ({idx}) | n_trials {len(setting.trials)} | settings:{setting}\n'
        return msg
    def __repr__(self):
        return self.__str__()

    