import torch
import torch.nn as nn
import numpy as np
from .tft import sub_nn

try:
    import lightning.pytorch as pl
    from .base_v2 import Base
    OLD_PL = False
except:
    import pytorch_lightning as pl
    OLD_PL = True
    from .base import Base
from typing import List, Union
from ..data_structure.utils import beauty_string
from .utils import  get_scope

class Diffusion(Base):
    handle_multivariate = False
    handle_future_covariates = True
    handle_categorical_variables = True
    handle_quantile_loss = False
    description = get_scope(handle_multivariate,handle_future_covariates,handle_categorical_variables,handle_quantile_loss)
    def __init__(self, 
                 d_model: int,
                 out_channels: int,
                 past_steps: int,
                 future_steps: int, 
                 past_channels: int,
                 future_channels: int,
                 embs: List[int],

                 learn_var:bool, 
                 cosine_alpha: bool,
                 diffusion_steps: int,
                 beta: float,
                 gamma:float,

                 #for subnet
                 n_layers_RNN: int,
                 d_head: int,
                 n_head: int,
                 dropout_rate: float,
                 activation: str,
                 subnet:int,
                 perc_subnet_learning_for_step:float,

                 persistence_weight:float=0.0,
                 loss_type: str='l1',
                 quantiles:List[float]=[],
                 optim:Union[str,None]=None,
                 optim_config:Union[dict,None]=None,
                 scheduler_config:Union[dict,None]=None,
                 **kwargs)->None:
        """Denoising Diffusion Probabilistic Model

        Args:
            d_model (int): 
            out_channels (int): number of target variables
            past_steps (int): size of past window
            future_steps (int): size of future window to be predicted
            past_channels (int): number of variables available for the past context
            future_channels (int): number of variables known in the future, available for forecasting
            embs (list[int]): categorical variables dimensions for embeddings
            learn_var (bool): Flag to make the model train the posterior variance (if True) or use the variance of posterior distribution 
            cosine_alpha (bool): Flag for the generation of alphas and betas
            diffusion_steps (int): number of noising steps for the initial sample
            beta (float): starting variable to generate the diffusion perturbations. Ignored if cosine_alpha == True
            gamma (float): trade_off variable to balance loss over noise prediction and NegativeLikelihood/KL_Divergence.
            n_layers_RNN (int): param for subnet
            d_head (int): param for subnet
            n_head (int): param for subnet
            dropout_rate (float): param for subnet
            activation (str): param for subnet
            subnet (int): =1 for attention subnet, =2 for linear subnet. Others can be added(wait for Black Friday for discounts)
            perc_subnet_learning_for_step (float): percentage to choose how many subnet has to be trained for every batch. Decrease this value if the loss blows up.
            persistence_weight (float, optional): Defaults to 0.0.
            loss_type (str, optional): Defaults to 'l1'.
            quantiles (List[float], optional): Only [] accepted. Defaults to [].
            optim (Union[str,None], optional): Defaults to None.
            optim_config (Union[dict,None], optional): Defaults to None.
            scheduler_config (Union[dict,None], optional): Defaults to None.
        """
        
        super().__init__(**kwargs)
        self.save_hyperparameters(logger=False)

        self.dropout = dropout_rate
        self.persistence_weight = persistence_weight 
        self.loss_type = loss_type
        self.optim = optim
        self.optim_config = optim_config
        self.scheduler_config = scheduler_config

        #* HANDLING LOSSES 
        # With respect to other models, here quantiles are not used
        # Here we define the loss used for noise predicted and actual noise_loss
        # Losses for distribution are defined as functions below.
        
        # trade off for noise loss and distribution loss 
        self.gamma = gamma
        ## can not handle quantile
        if len(quantiles)>0:
            quantiles = []
        assert len(quantiles) ==0
        self.mul = 1
        self.use_quantiles = False
        if self.loss_type == 'mse':
            self.loss = nn.MSELoss()
        else:
            self.loss = nn.L1Loss()
        
        #* >>>>>>>>>>>>> canonical data parameters
        # dimension of the model, number of variables and sequence length info
        self.d_model = d_model
        self.past_steps = past_steps
        self.future_steps = future_steps
        self.past_channels = past_channels
        self.future_channels = future_channels
        self.output_channels = out_channels

        #* >>>>>>>>>>>>> specific model parameters
        # if we want to learn also the variance, instead of using the standard posterior variance of Diffusion NN
        self.learn_var = learn_var
        
        # number of noising steps
        self.T = diffusion_steps 

        # distribution weigths to avoid less trained subnet 
        self.multinomial_step_weights = np.ones(diffusion_steps) 
        
        # % of all subnets trained every batch of every epoch
        # this percentage is controlled by the parameter 'perc_subnet_learning_for_step': 
        # - decrease or increase according to the efficiency of your machine
        self.simultaneous_steps = max(int(diffusion_steps*perc_subnet_learning_for_step), 1) 
        

        #* >>>>>>>>>>>>> specific diffusion setup
        self.s = (100*self.T)**(-1)  # offset variable to avoid problems with computations near 0
        # value found by try and error

        # betas and cumulative products of alphas are the main values for the diffusion model, both in (0,1) at each step t:
        # - betas_t -> variance added at t-th step
        # - alphas_t = 1 - betas_t
        # - alphas_cumprod_t -> accumulation of alphas up to step t. 
        #  - - It can be considered as the remaining signal of the starting input at t-th step!!

        # according to the flag below we can choose how to generate them!
        if cosine_alpha:
            # COSINE ALPHA Computation
            # aux_perc = 0.05
            # avoid_comp_err_norm = self.T*(1+aux_perc) # enlarging self.T to avoid errors in computations using cos^2
            # the t-th cumulative product of alphas is the 'forgetting' schedule of the inital sample after t diffusion step
            # in this procedure we use the function below to produce all the cumulative products of alphas
            f_cos_t = [(np.cos( (t/self.T +self.s)/(1+self.s) * np.pi/2 ))**2 for t in range(self.T)]

            self.alphas_cumprod = np.append(1-self.s, f_cos_t[1:]/f_cos_t[0]) # computed as scaled cumulative product of alphas f_cos_t[1:]/f_cos_t[0]
            self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1]) # auxiliar vector to use the same index to access alpha_cumprod_t and alpha_cumprod_{t-1} 
            self.alphas = self.alphas_cumprod * (self.alphas_cumprod_prev)**(-1)
            self.betas = 1 - self.alphas

        else:
            # STANDARD ALPHA Computation
            # beta is considered constant in [0,1) for all time steps. Good values near 0.03
            # Unlike before, here we generate all needed values starting from betas
            self.betas = np.array([beta]*self.T) 
            self.alphas = 1 - self.betas
            self.alphas_cumprod = np.cumprod(self.alphas)
            self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1]) # auxiliar vector to use the same index to access alpha_cumprod_t and alpha_cumprod_{t-1}

        # values for posterior distribution, id est the target distribution of each subnet
        # All these values will be casted to tensors during computations using the function _extract_into_tensor
        self.posterior_mean_coef1 = self.betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
        self.posterior_mean_coef2 = (1.0 - self.alphas_cumprod_prev) * np.sqrt(self.alphas) / (1.0 - self.alphas_cumprod)
        self.posterior_variance = np.append(self.s, self.betas[1:] * (1.0 - self.alphas_cumprod_prev[1:]) / (1.0 - self.alphas_cumprod[1:]))
        self.posterior_log_variance = np.log(self.posterior_variance)

        #* >>>>>>>>>>>>> LAYERS
        # for other numerical variables in the past
        self.aux_past_channels = past_channels - out_channels
        self.linear_aux_past = nn.ModuleList([nn.Linear(1, d_model) for _ in range(self.aux_past_channels)])

        # for numerical variables in the future
        self.aux_fut_channels = future_channels
        self.linear_aux_fut = nn.ModuleList([nn.Linear(1, d_model) for _ in range(self.aux_fut_channels)])
        
        # embedding categorical for both past and future (ASSUMING BOTH AVAILABLE OR NO ONE)
        self.seq_len = past_steps + future_steps
        self.emb_cat_var = sub_nn.embedding_cat_variables(self.seq_len, future_steps, d_model, embs, self.device)

        # diffusion sub nets, one subnet for each step
        if subnet == 1:
            self.sub_nets = nn.ModuleList([
                SubNet1(self.aux_past_channels, self.aux_fut_channels, learn_var, out_channels, d_model, d_head, n_head, activation, dropout_rate) for _ in range(diffusion_steps)
            ])
        elif subnet == 2:
            self.sub_nets = nn.ModuleList([
                SubNet2(self.aux_past_channels, self.aux_fut_channels, learn_var, past_steps, future_steps, out_channels, d_model, activation, dropout_rate) for _ in range(diffusion_steps)
            ])
        elif subnet ==3 :
            aux_num_available = self.aux_past_channels>0 or self.aux_fut_channels>0 # if we have numerical vars, use it
            self.sub_nets = nn.ModuleList([
                SubNet3(learn_var, aux_num_available, out_channels, d_model, future_steps, n_layers_RNN, d_head, n_head, dropout_rate) for _ in range(diffusion_steps)
            ])
        else:
            raise ValueError("Wrong number for Subnet. Not yet implemented!")

    def forward(self, batch:dict)-> float:
        """training process of the diffusion network

        Args:
            batch (dict): variables loaded

        Returns:
            float: total loss about the prediction of the noises over all subnets extracted
        """

        # LOADING TARGET VARIABLES
        y_to_be_pred = batch['y'].to(self.device)

        # LOADING AUTOREGRESSIVE CONTEXT OF TARGET VARIABLES
        num_past = batch['x_num_past'].to(self.device)
        idx_target = batch['idx_target'][0]
        y_past = num_past[:,:,idx_target]

        # LOADING EMBEDDING CATEGORICAL VARIABLES
        emb_cat_past, emb_cat_fut = self.cat_categorical_vars(batch)
        emb_cat_past = torch.mean(emb_cat_past, dim = 2)
        emb_cat_fut = torch.mean(emb_cat_fut, dim = 2)

        ### LOADING PAST AND FUTURE NUMERICAL VARIABLES
        # load in the model auxiliar numerical variables

        if self.aux_past_channels>0: # if we have more numerical variables about past
            aux_num_past = self.remove_var(num_past, idx_target, 2) # remove the autoregressive variable
            assert self.aux_past_channels == aux_num_past.size(2),  beauty_string(f"{self.aux_past_channels} LAYERS FOR PAST VARS AND {aux_num_past.size(2)} VARS",'section',True) # to check if we are using the expected number of variables about past
            # concat all embedded vars and mean of them
            aux_emb_num_past = torch.Tensor().to(self.device)
            for i, layer in enumerate(self.linear_aux_past):
                aux_emb_past = layer(aux_num_past[:,:,[i]]).unsqueeze(2)
                aux_emb_num_past = torch.cat((aux_emb_num_past, aux_emb_past), dim=2)
            aux_emb_num_past = torch.mean(aux_emb_num_past, dim = 2)
        else: 
            aux_emb_num_past = None # non available vars
            
        if self.aux_fut_channels>0: # if we have more numerical variables about future
            # AUX means AUXILIARY variables
            aux_num_fut = batch['x_num_future'].to(self.device)
            assert self.aux_fut_channels == aux_num_fut.size(2), beauty_string(f"{self.aux_fut_channels} LAYERS FOR PAST VARS AND {aux_num_fut.size(2)} VARS",'section',True)  # to check if we are using the expected number of variables about fut
            # concat all embedded vars and mean of them
            aux_emb_num_fut = torch.Tensor().to(self.device)
            for j, layer in enumerate(self.linear_aux_fut):
                aux_emb_fut = layer(aux_num_fut[:,:,[j]]).unsqueeze(2)
                aux_emb_num_fut = torch.cat((aux_emb_num_fut, aux_emb_fut), dim=2)
            aux_emb_num_fut = torch.mean(aux_emb_num_fut, dim = 2)
        else:
            aux_emb_num_fut = None # non available vars

        ### actual DIFFUSION process ----------------------------------------------

        ##* CHOOSE THE t SUBNET
        # We have T subnets: [0, 1, ..., T-1].
        values = list(range(self.T))
        
        ## Probabilistic way to choose the subnet properly
        # avoid exploding step_weights going on with trainings
        self.improving_weight_during_training()
        # normalizing weights for a more stable subnet training
        t_wei = self.multinomial_step_weights/np.sum(self.multinomial_step_weights)
        # extract times t
        drawn_t = np.random.choice(values, size=self.simultaneous_steps, replace=False, p=t_wei) # type: ignore
        
        if 0 not in drawn_t: drawn_t = np.append(drawn_t, 0)
        # update weights
        non_draw_val = np.delete(values, drawn_t) # type: ignore
        self.multinomial_step_weights[non_draw_val] += 1 # increase weights of non-extracted subnet

        # init negative loss for the first step
        tot_loss = -1
        for t in drawn_t:
            # LOADING THE SUBNET
            sub_net = self.sub_nets[t]

            # Get y and noise it
            y_noised, true_mean, true_log_var, actual_noise = self.q_sample(y_to_be_pred, t)

            # compute the output from that network using the sample with noises
            # output composed of: noise predicted and, if learn_var=True, vector for variances
            if self.learn_var:
                #predict the noise!
                eps_pred, var_aux_out = sub_net(y_noised, y_past, emb_cat_past, emb_cat_fut, aux_emb_num_past, aux_emb_num_fut)

                # compute posterior variance of NN (using interpolation)
                pre_var_t = self._extract_into_tensor(self.betas, t, eps_pred.shape)
                post_var_t = self._extract_into_tensor(self.posterior_variance, t, eps_pred.shape)
                post_sigma = torch.exp( var_aux_out*torch.log(pre_var_t) + (1-var_aux_out)*torch.log(post_var_t) ) # variance, not log_var
            else:
                eps_pred = sub_net(y_noised, y_past, emb_cat_past, emb_cat_fut, aux_emb_num_past, aux_emb_num_fut)
                post_sigma = self._extract_into_tensor(self.posterior_variance, t, eps_pred.shape)

            # posterior mean assuming the predicted noise is the actual one
            out_mean = self._extract_into_tensor(np.sqrt(1/self.alphas), t, eps_pred.shape) * ( y_noised - self._extract_into_tensor(self.betas/np.sqrt(1-self.alphas_cumprod), t, eps_pred.shape) * eps_pred )
            
            # # At the first timestep return the negative likelihood,
            if t==0:
                # post_var =  self._extract_into_tensor(self.posterior_variance, t, y_to_be_pred.shape)
                neg_likelihoods = -torch.log(self.gaussian_likelihood(y_to_be_pred, out_mean, post_sigma)) #! (values to be predicted, mean of values predicted, variance)
                distribution_loss = torch.mean(neg_likelihoods)

            # # otherwise return KL( q(x_{t-1}|x_t, x_0) || p(x_{t-1}|x_t) )
            else:
                # COMPUTE LOSS between TRUE eps and DRAWN eps_pred
                kl_divergence = self.normal_kl(true_mean, true_log_var, out_mean, torch.log(post_sigma)) # (true mean, true log var, mean of values predicted, log var predicted)
                distribution_loss = torch.mean(kl_divergence)

            # always compute the loss about the straight prediction of the noise
            noise_loss = self.loss(eps_pred, actual_noise)

            # if tot_loss == -1:
            #     beauty_string(f'NOISE LOSS: {noise_loss.item()}','info',True)
            #     beauty_string(f'ACTUAL NOISE: {actual_noise[0].min()}, {actual_noise[0].max()}, {actual_noise[0].mean()}, {actual_noise[0].var()}','info',True)
            #     beauty_string(f'PREDICTED NOISE: {eps_pred[0].min()}, {eps_pred[0].max()}, {eps_pred[0].mean()}, {eps_pred[0].var()}','info',True)

            noise_loss += self.gamma*distribution_loss # add, scaled according to gamma, the distribution_loss

            # update the total loss
            if tot_loss==-1:
                tot_loss = noise_loss
            else:
                tot_loss += noise_loss
        return tot_loss

    def training_step(self, batch, batch_idx):
        # the training loss is already computed in the forward method
        loss_eps = self(batch)
        return loss_eps
        
    def inference(self, batch:dict) -> torch.Tensor:
        """Inference process to forecast future y

        Args:
            batch (dict): Keys checked ['x_num_past, 'idx_target', 'x_num_future', 'x_cat_past', 'x_cat_future']

        Returns:
            torch.Tensor: generated sequence [batch_size, future_steps, num_var]
        """
        # LOADING AUTOREGRESSIVE CONTEXT OF TARGET VARIABLES
        num_past = batch['x_num_past'].to(self.device)
        batch_size = num_past.shape[0]
        idx_target = batch['idx_target'][0]
        y_past = num_past[:,:,idx_target]        

        # LOADING EMBEDDING CATEGORICAL VARIABLES
        emb_cat_past, emb_cat_fut = self.cat_categorical_vars(batch)
        emb_cat_past = torch.mean(emb_cat_past, dim = 2)
        emb_cat_fut = torch.mean(emb_cat_fut, dim = 2)

        ### LOADING PAST AND FUTURE NUMERICAL VARIABLES
        # this check is done simultaneously 
        # because in the model we use auxiliar numerical variables 
        # only if we have both them in the past and in the future

        ### LOADING PAST AND FUTURE NUMERICAL VARIABLES
        # load in the model auxiliar numerical variables

        if self.aux_past_channels>0: # if we have more numerical variables about past
            aux_num_past = self.remove_var(num_past, idx_target, 2) # remove the autoregressive variable
            assert self.aux_past_channels == aux_num_past.size(2),  beauty_string(f"{self.aux_past_channels} LAYERS FOR PAST VARS AND {aux_num_past.size(2)} VARS",'section',True) # to check if we are using the expected number of variables about past
            # concat all embedded vars and mean of them
            aux_emb_num_past = torch.Tensor().to(self.device)
            for i, layer in enumerate(self.linear_aux_past):
                aux_emb_past = layer(aux_num_past[:,:,[i]]).unsqueeze(2)
                aux_emb_num_past = torch.cat((aux_emb_num_past, aux_emb_past), dim=2)
            aux_emb_num_past = torch.mean(aux_emb_num_past, dim = 2)
        else: 
            aux_emb_num_past = None # non available vars
            
        if self.aux_fut_channels>0: # if we have more numerical variables about future
            # AUX means AUXILIARY variables
            aux_num_fut = batch['x_num_future'].to(self.device)
            assert self.aux_fut_channels == aux_num_fut.size(2), beauty_string(f"{self.aux_fut_channels} LAYERS FOR PAST VARS AND {aux_num_fut.size(2)} VARS",'section',True)  # to check if we are using the expected number of variables about fut
            # concat all embedded vars and mean of them
            aux_emb_num_fut = torch.Tensor().to(self.device)
            for j, layer in enumerate(self.linear_aux_fut):
                aux_emb_fut = layer(aux_num_fut[:,:,[j]]).unsqueeze(2)
                aux_emb_num_fut = torch.cat((aux_emb_num_fut, aux_emb_fut), dim=2)
            aux_emb_num_fut = torch.mean(aux_emb_num_fut, dim = 2)
        else:
            aux_emb_num_fut = None # non available vars
        
        # DIFFUSION INFERENCE
        # import pdb; pdb.set_trace() # can use also torch.normal(0, 1, size=y_noised.shape)
        y_noised = torch.randn((batch_size, self.future_steps, self.output_channels)).to(self.device)
        # pass the white noise in sub nets
        for t in range(self.T-1, -1, -1): # INVERSE cycle over all subnets, but not the last one
            sub_net = self.sub_nets[t] # load the subnet

            ## CHECK THE NUMBER OF PARAMS
            #   model_parameters = filter(lambda p: p.requires_grad, model.parameters())
            #   params = sum([np.prod(p.size()) for p in model_parameters]) -> 13K
            if self.learn_var:
                eps_pred, var_aux_out = sub_net(y_noised, y_past, emb_cat_past, emb_cat_fut, aux_emb_num_past, aux_emb_num_fut)
                # interpolazion of variance
                pre_var_t = self._extract_into_tensor(self.betas, t, eps_pred.shape)
                post_var_t = self._extract_into_tensor(self.posterior_variance, t, eps_pred.shape)
                post_sigma = torch.exp(var_aux_out*torch.log(pre_var_t) + (1-var_aux_out)*torch.log(post_var_t))
            else:
                eps_pred = sub_net(y_noised, y_past, emb_cat_past, emb_cat_fut, aux_emb_num_past, aux_emb_num_fut)
                post_sigma = self._extract_into_tensor(self.posterior_variance, t, eps_pred.shape)

            # Sample x_{t-1} from the model at the given timestep.
            # y_noised = self._extract_into_tensor(1/np.sqrt(self.alphas), t, y_noised.shape)*( y_noised - self._extract_into_tensor(np.sqrt(self.betas), t, eps_pred.shape)*eps_pred )
            y_noised = 1/torch.sqrt(1-post_sigma)*(y_noised - torch.sqrt(post_sigma)*eps_pred)

            # if t>0 :
            #     noise = torch.rand_like(y_noised).to(self.device)
            #     y_noised = y_noised + torch.sqrt(post_sigma)*noise        
        
        out = y_noised.view(-1, self.future_steps, self.output_channels, 1)
        return out

    # for validation extract the output from the self.inference method
    def validation_step(self, batch, batch_idx):
        out = self.inference(batch)
        loss = self.compute_loss(batch,out)
        return loss

    def can_be_compiled(self):
        return False

    # function to concat embedded categorical variables
    def cat_categorical_vars(self, batch:dict):
        """Extracting categorical context about past and future

        Args:
            batch (dict): Keys checked -> ['x_cat_past', 'x_cat_future']

        Returns:
            List[torch.Tensor, torch.Tensor]: cat_emb_past, cat_emb_fut
        """
        # GET AVAILABLE CATEGORICAL CONTEXT
        if 'x_cat_past' in batch.keys() and 'x_cat_future' in batch.keys(): # if we have both
            # HERE WE ASSUME SAME NUMBER AND KIND OF VARIABLES IN PAST AND FUTURE
            cat_past = batch['x_cat_past'].to(self.device)
            cat_fut = batch['x_cat_future'].to(self.device)
            cat_full = torch.cat((cat_past, cat_fut), dim = 1)
            # EMB CATEGORICAL VARIABLES AND THEN SPLIT IN PAST AND FUTURE
            emb_cat_full = self.emb_cat_var(cat_full,self.device)
        else:
            emb_cat_full = self.emb_cat_var(batch['x_num_past'].shape[0],self.device)

        # CONCAT THEM, according to self.emb_cat_var usage  
        cat_full = torch.cat((cat_past, cat_fut), dim = 1)
        # actual embedding
        emb_cat_full = self.emb_cat_var(cat_full,self.device)
        # split past and future categorical embedded variables
        cat_emb_past = emb_cat_full[:,:self.past_steps,:,:]
        cat_emb_fut = emb_cat_full[:,-self.future_steps:,:,:]

        return cat_emb_past, cat_emb_fut

    #function to extract from batch['x_num_past'] all variables except the one autoregressive
    def remove_var(self, tensor: torch.Tensor, indexes_to_exclude: list, dimension: int)-> torch.Tensor:
        """Function to remove variables from tensors in chosen dimension and position 

        Args:
            tensor (torch.Tensor): starting tensor
            indexes_to_exclude (list): index of the chosen dimension we want t oexclude
            dimension (int): dimension of the tensor on which we want to work (not list od dims!!)

        Returns:
            torch.Tensor: new tensor without the chosen variables
        """

        remaining_idx = torch.tensor([i for i in range(tensor.size(dimension)) if i not in indexes_to_exclude]).to(tensor.device)
        # Select the desired sub-tensor
        extracted_subtensors = torch.index_select(tensor, dim=dimension, index=remaining_idx)
        
        return extracted_subtensors
    
    def improving_weight_during_training(self):
        """
        Each time we sample from multinomial we subtract the minimum for more precise sampling, 
        avoiding great learning differences among subnets

        This lead to more stable inference also in early training, mainly for common context embedding.

        For probabilistic reason, weights has to be >0, so we subtract min-1
        """
        self.multinomial_step_weights -= (self.multinomial_step_weights.min()-1)
        return
    
    ### >>>>>>>>>>>>> AUXILIARY MODEL FUNCS
    def q_sample(self, x_start: torch.Tensor, t: int)-> List[torch.Tensor]:
        """Diffuse x_start for t diffusion steps.

        In other words, sample from q(x_t | x_0).

        Also, compute the mean and variance of the diffusion posterior:

            q(x_{t-1} | x_t, x_0)

        Posterior mean and variance are the ones to be predicted

        Args:
            x_start (torch.Tensor): values to be predicted
            t (int): diffusion step

        Returns:
            List[torch.Tensor, torch.Tensor, torch.Tensor]: q_sample, posterior mean, posterior log variance and the actual noise
        """
        # noise from normal distribution
        noise = torch.randn_like(x_start)

        # direct diffusion at t-th step
        q_sample = self._extract_into_tensor(np.sqrt(self.alphas_cumprod), t, x_start.shape) * x_start + self._extract_into_tensor(np.sqrt(1 - self.alphas_cumprod), t, x_start.shape) * noise

        # compute meean and variance
        q_mean = self._extract_into_tensor(self.posterior_mean_coef1, t, q_sample.shape) * x_start + self._extract_into_tensor(self.posterior_mean_coef2, t, q_sample.shape) * q_sample
        q_log_var = self._extract_into_tensor( self.posterior_log_variance, t, q_sample.shape )

        # return, the sample, its posterior mean and log_variance, the noise used
        return [q_sample, q_mean, q_log_var, noise]

    def normal_kl(self, mean1, logvar1, mean2, logvar2):
        """
        Compute the KL divergence between two gaussians. Also called relative entropy.
        KL divergence of P from Q is the expected excess surprise from using Q as a model when the actual distribution is P.
        KL(P||Q) = P*log(P/Q) or -P*log(Q/P)

        # In the context of machine learning, KL(P||Q) is often called the 'information gain' 
        # achieved if P would be used instead of Q which is currently used.

        Shapes are automatically broadcasted, so batches can be compared to
        scalars, among other use cases.
        """
        # -1/2 + log(sigma2/sigma1) + sigma1^2/2sigma2^2 + (mu1-mu2)^2/2sigma2^2
        return 0.5 * (
            -1.0
            + logvar2 - logvar1
            + torch.exp(logvar1 - logvar2)
            + ((mean1 - mean2) ** 2) * torch.exp(-logvar2)
        )
    
    def gaussian_likelihood(self, x, mean, var):
        term1 = 1.0 / torch.sqrt(2 * np.pi * var)
        term2 = torch.exp(-0.5 * ((x - mean)**2 / var))
        likelihood = term1 * term2
        return likelihood

    def gaussian_log_likelihood(self, x, mean, var):
        term1 = -0.5 * ((x - mean) / torch.sqrt(var))**2
        term2 = -0.5 * torch.log(2 * torch.tensor(np.pi) * var)
        log_likelihood = term1 + term2
        return log_likelihood

    def _extract_into_tensor(self, arr, timesteps, broadcast_shape):
        """
        Extract values from a 1-D numpy array for a batch of indices.

        :param arr: the 1-D numpy array.
        :param timesteps: a tensor of indices into the array to extract.
        :param broadcast_shape: a larger shape of K dimensions with the batch
                                dimension equal to the length of timesteps.
        :return: a tensor of shape 'broadcast_shape' where the shape has K dims.
        """
        ten = torch.tensor(arr[timesteps])
        return ten.expand(broadcast_shape).to(self.device)

### >>>>>>>>>>>>>  SUB NET 
class SubNet1(nn.Module):
    def __init__(self, aux_past_ch, aux_fut_ch, learn_var:bool, output_channel:int, d_model:int, d_head:int, n_head:int, activation:str, dropout_rate:float) -> None:
        """ -> SUBNET of the DIFFUSION MODEL (DDPM)

        It starts with an autoregressive LSTM Network computation of epsilon, then subtracted to 'y_noised' tensor. This is always possible!
        Now we have an approximation of our 'eps_hat', that at the end will pass in a residual connection with its embedded version 'emb_eps_hat'.

        'emb_eps_hat' will be update with respect to available info about categorical values of our serie:
        Through an ATTENTION Network we compare past categorical with future categorical to update the embedded noise predicted.

        Also, if we have values about auxiliary numerical variables both in past and future, the changes of these variables will be fetched 
        by another ATTENTION Network.

        The goal is ensure valuable computations for 'eps' always, and then updating things if we have enough data.
        Both attentions uses { Q = *_future, K = *_past, V = y_past } using as much as possible context variables for better updates.

        Args:
            learn_var (bool): set if the network has to learn the optim variance of each step
            output_channel (int): number of variables to be predicted 
            future_steps (int): number of step in the future, so the number of timesstep to be predicted
            d_model (int): hidden dimension of the model
            num_layers_RNN (int): number of layers for autoregressive prediction
            d_head (int): number of heads for Attention Networks
            n_head (int): hidden dimension of heads for Attention Networks
            dropout_rate (float): 
        """
        super().__init__()
        self.aux_past_channels = aux_past_ch
        self.aux_fut_channels = aux_fut_ch

        self.learn_var = learn_var
        activation_fun = eval(activation)

        self.y_noised_linear = nn.Linear(output_channel, d_model)
        self.y_past_linear = nn.Linear(output_channel, d_model)

        self.past_sequential = nn.Sequential(
            nn.Linear(d_model*3, d_model*2) if self.aux_past_channels>0 else nn.Linear(d_model*2, d_model*2),
            activation_fun(),
            nn.Linear(d_model*2, d_model)
        )
        
        self.fut_sequential = nn.Sequential(
            nn.Linear(d_model*3, d_model*2) if self.aux_fut_channels>0 else nn.Linear(d_model*2, d_model*2),
            activation_fun(),
            nn.Linear(d_model*2, d_model)
        )

        self.y_sequential = nn.Sequential(
            nn.Linear(d_model*2, d_model),
            activation_fun(),
            nn.Linear(d_model, d_model)
        )

        self.attention = sub_nn.InterpretableMultiHead(d_model, d_head, n_head)

        # if learn_var == True, we want to predict an additional variable for he variance
        # just an intermediate dimension for linears
        hidden_size = int(d_model/3)
        self.eps_out_sequential = nn.Sequential(
            nn.Linear(d_model, hidden_size),
            activation_fun(),
            nn.Linear(hidden_size, output_channel)
        )

        self.var_out_sequential = nn.Sequential(
            nn.Linear(output_channel, hidden_size),
            nn.Linear(hidden_size, d_model),
            activation_fun(),
            nn.Linear(d_model, d_model),
            activation_fun(),
            nn.Linear(d_model, hidden_size),
            nn.Linear(hidden_size, output_channel)
        )


    def forward(self, y_noised:torch.Tensor, y_past:torch.Tensor,
                cat_past:torch.Tensor, cat_fut:torch.Tensor, 
                num_past:Union[torch.Tensor,None] = None, num_fut:Union[torch.Tensor,None] = None):
        """'DIFFUSION SUBNET
        Args:
            y_noised (torch.Tensor): [B, future_step, num_var]
            y_past (torch.Tensor): [B, past_step, num_var]
            cat_past (torch.Tensor, optional): [B, past_step, d_model]. Defaults to None.
            cat_fut (torch.Tensor, optional): [B, future_step, d_model]. Defaults to None.
            num_past (torch.Tensor, optional): [B, past_step, d_model]. Defaults to None.
            num_fut (torch.Tensor, optional): [B, future_step, d_model]. Defaults to None.

        Returns:
            torch.Tensor: predicted noise [B, future_step, num_var]. According to 'learn_var' param in initialization, the subnet returns another tensor of same size about the variance 
        """
        emb_y_noised = self.y_noised_linear(y_noised.float())
        emb_y_past = self.y_past_linear(y_past)
        
        # LIN FOR PAST
        past = [emb_y_past, cat_past]
        if self.aux_past_channels>0:
            past.append(num_past)
        past_seq_input = torch.cat(past, dim=2) # type: ignore
        past_seq = self.past_sequential(past_seq_input) # -> [B, future_step, d_model]

        # LIN FOR FUT
        fut = [emb_y_noised, cat_fut]
        if self.aux_fut_channels>0:
            fut.append(num_fut)
        fut_seq_input = torch.cat(fut, dim=2) # type: ignore
        fut_seq = self.fut_sequential(fut_seq_input) # -> [B, future_step, d_model]
        # ATTENTION
        attention = self.attention(fut_seq, past_seq, emb_y_past)
        # OUTPUT
        eps_out = self.eps_out_sequential(attention)
        # if LEARN_VAR
        if self.learn_var:
            var_out = eps_out.detach()
            var_out = self.var_out_sequential(var_out)
            return eps_out, var_out

        return eps_out
    
class SubNet2(nn.Module):
    def __init__(self, aux_past_ch, aux_fut_ch, learn_var:bool, past_steps, future_steps, output_channel:int, d_model:int, activation:str, dropout_rate:float):
        super().__init__()
        self.aux_past_channels = aux_past_ch
        self.aux_fut_channels = aux_fut_ch
        self.learn_var = learn_var
        # in_size changes wrt numerical vars
        in_size = ( past_steps*(2+bool(aux_past_ch)) + future_steps*(2 + bool(aux_fut_ch)) ) * d_model
        out_size = output_channel * future_steps

        activation_fun = eval(activation)

        self.y_noised_linear = nn.Linear(output_channel, d_model)
        self.y_past_linear = nn.Linear(output_channel, d_model)

        hidden_size = int( (output_channel + d_model)/2 )
        self.eps_out_sequential = nn.Sequential(
            nn.Linear(in_size, hidden_size),
            nn.Dropout(dropout_rate),
            nn.Linear(hidden_size, d_model),
            activation_fun(),
            nn.Dropout(dropout_rate),
            nn.Linear(d_model, hidden_size),
            activation_fun(),
            nn.Dropout(dropout_rate),
            nn.Linear(hidden_size, out_size)
        )
        
        self.var_out_sequential = nn.Sequential(
            nn.Linear(in_size, hidden_size),
            nn.Dropout(dropout_rate),
            nn.Linear(hidden_size, hidden_size),
            activation_fun(),
            nn.Dropout(dropout_rate),
            nn.Linear(hidden_size, out_size)
        )

    def forward(self, y_noised:torch.Tensor, y_past:torch.Tensor,
                cat_past:torch.Tensor, cat_fut:torch.Tensor, 
                num_past:Union[torch.Tensor,None] = None, num_fut:Union[torch.Tensor,None] = None):
        
        B, fut_step, n_var = y_noised.shape
        emb_y_noised = self.y_noised_linear(y_noised.float()).view(B, -1)
        emb_y_past = self.y_past_linear(y_past).view(B, -1)

        # concat auroregressive variables and categorical ones that are always available
        full_concat = torch.cat((emb_y_noised, emb_y_past, cat_past.view(B, -1), cat_fut.view(B, -1)), dim=1)
        # concat numerical vars when available
        if num_past is not None:
            assert self.aux_past_channels>0 # check with flag in subnet init
            full_concat = torch.cat((full_concat, num_past.view(B, -1)), dim = 1)
        if num_fut is not None:
            assert self.aux_fut_channels>0 # check with flag in subnet init
            full_concat = torch.cat((full_concat, num_fut.view(B, -1)), dim = 1)

        eps_out = self.eps_out_sequential(full_concat).view(B, fut_step, n_var)
        if self.learn_var:
            var_out = self.var_out_sequential(full_concat.detach()).view(B, fut_step, n_var)
            return eps_out, var_out
        return eps_out

class SubNet3(nn.Module):
    def __init__(self, learn_var, flag_aux_num, num_var, d_model, pred_step, num_layers, d_head, n_head, dropout):
        super().__init__()
        self.learn_var = learn_var
        self.flag_aux_num = flag_aux_num
        
        # Autoregressive with RNN (y NOT embedded as inpute)
        self.y_d_model = nn.Linear(num_var, d_model)
        self.rnn = sub_nn.LSTM_Model(num_var, d_model, pred_step, num_layers, dropout)
        self.eps_pred_grn = sub_nn.GRN(d_model, dropout)

        #categorical
        self.cat_MHA = sub_nn.InterpretableMultiHead(d_model, d_head, n_head)
        self.cat_grn = sub_nn.GRN(d_model, dropout)
        self.cat_res_conn = sub_nn.ResidualConnection(d_model, dropout)

        #numerical
        if flag_aux_num:
            self.num_MHA = sub_nn.InterpretableMultiHead(d_model, d_head, n_head)
            self.num_grn = sub_nn.GRN(d_model, dropout)
            self.num_res_conn = sub_nn.ResidualConnection(d_model, dropout)
        
        # EPS PREDICTION
        self.eps_final_grn = sub_nn.GRN(d_model, dropout)
        self.eps_out_linear = nn.Linear(d_model, num_var)

        if learn_var:
            self.emb_eps_pred = nn.Linear(num_var, d_model)
            self.var_att = sub_nn.InterpretableMultiHead(d_model, d_head, n_head)
            self.var_grn = sub_nn.GRN(d_model, dropout)
            self.var_out = nn.Linear(d_model, num_var)

    def forward(self, y_noised:torch.Tensor, y_past:torch.Tensor, 
                cat_past:torch.Tensor, cat_fut:torch.Tensor, 
                num_past:Union[torch.Tensor,None] = None, num_fut:Union[torch.Tensor,None] = None):

        # Autoregressive
        emb_y_past = self.y_d_model(y_past)
        pred_y_fut = self.rnn(emb_y_past)
        #re-embedding future
        emb_pred_y_fut = self.y_d_model(pred_y_fut)
        emb_y_noised = self.y_d_model(y_noised.float())

        eps_pred = self.eps_pred_grn(emb_pred_y_fut - emb_y_noised, using_norm=False)

        # Categorical contribute
        cat_att = self.cat_MHA(cat_fut, cat_past, emb_y_past)
        cat_att = self.cat_grn(cat_att, using_norm=False)
        eps_pred = self.cat_res_conn(cat_att, eps_pred, using_norm=False)

        # Numerical contribute
        if self.flag_aux_num:
            if num_past is None:
                num_past = torch.ones_like(cat_past)
            if num_fut is None:
                num_fut = torch.ones_like(cat_fut)
            num_att = self.num_MHA(num_fut, cat_past, emb_y_past)
            num_att = self.num_grn(num_att, using_norm=False)
            eps_pred = self.cat_res_conn(num_att, eps_pred, using_norm=False)

        eps_pred = self.eps_final_grn(eps_pred, using_norm=False)
        eps_pred = self.eps_out_linear(eps_pred)

        if self.learn_var:
            emb_eps_pred = self.emb_eps_pred(eps_pred.detach())
            emb_eps_pred = self.var_att(emb_y_noised.detach(), emb_pred_y_fut.detach(), emb_eps_pred)
            emb_var_pred = self.var_grn(emb_eps_pred, using_norm=False)
            var_pred = self.var_out(emb_var_pred)
            return eps_pred, var_pred
        return eps_pred
            