import torch

class PandemicDataset:
    def __init__(self, 
                 name:str,
                 group_names:list, 
                 N: int, 
                 t, 
                 *groups):
        """Class to hold all data for one training process.

        Args:
            group_names (list): Names of the groups, in which the population is split.
            N (int): Size of the population.
            t (np.array): Array of timesteps.
            *groups (np.array): Arrays of size data for each group for each timestep..
        """

        if torch.cuda.is_available():
            self.device_name = 'cuda'
        else:
            self.device_name = 'cpu'
        self.name = name
        self.N = N
        self.t_raw = torch.tensor(t, requires_grad=True, device=self.device_name)
        self.t_batch = self.t_raw.view(-1, 1).float()

        self.__group_dict = {}
        for i, name in enumerate(group_names):
            self.__group_dict.update({name : i})

        self.__group_names = group_names

        self.__groups = [torch.tensor(group, device=self.device_name) for group in groups]
        
        self.__mins = [torch.min(group) for group in self.__groups]
        self.__maxs = [torch.max(group) for group in self.__groups]
        self.__norms = [(self.__groups[i] - self.__mins[i]) / (self.__maxs[i] - self.__mins[i]) for i in range(len(groups))]

        self.number_groups = len(groups)

    def get_data(self):
        return self.__groups

    def get_group(self, name:str):
        return self.__groups[self.__group_dict[name]]
    
    def get_min(self, name:str):
        return self.__mins[self.__group_dict[name]]
    
    def get_max(self, name:str):
        return self.__maxs[self.__group_dict[name]]
    
    def get_norm(self, name:str):
        return self.__norms[self.__group_dict[name]]
    
    def get_group_names(self):
        return self.__group_names
    
