# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/03_models/00_layers.ipynb.

# %% auto 0
__all__ = ['BatchNorm_1D_Stateful', 'SeqLinear', 'Normalizer1D', 'AR_Model', 'NormalizedModel']

# %% ../../nbs/03_models/00_layers.ipynb 2
from ..data import *
from fastai.basics import *

from fastai.callback.progress import * #import activate learning progress bar
from torch.nn import Parameter

# %% ../../nbs/03_models/00_layers.ipynb 6
class BatchNorm_1D_Stateful(nn.Module):
    '''Batchnorm for stateful models. Stores batch statistics for for every timestep seperately to mitigate transient effects.'''
    __constants__ = ['track_running_stats', 'momentum', 'eps', 'weight', 'bias',
                     'running_mean', 'running_var', 'num_batches_tracked']

    def __init__(self, hidden_size, seq_len=None,stateful=False,batch_first=True, eps=1e-7, momentum=0.1, affine=True,
                 track_running_stats=True):#num_features
        super().__init__()
        channel_d=hidden_size
        self.seq_len = seq_len
        self.stateful = stateful
        self.batch_first = batch_first
        self.eps = eps
        self.momentum = momentum
        self.affine = affine
        self.track_running_stats = track_running_stats
        self.axes=(1,)
        if self.affine:
            self.weight = Parameter(torch.Tensor(channel_d))
            self.bias = Parameter(torch.Tensor(channel_d))
            self.register_parameter('weight', self.weight)
            self.register_parameter('bias', self.bias)
        else:
            self.register_parameter('weight', None)
            self.register_parameter('bias', None)
        if self.track_running_stats:
            if seq_len is not None:
                self.register_buffer('running_mean', torch.zeros(seq_len,channel_d))
                self.register_buffer('running_var', torch.ones(seq_len,channel_d))
            self.register_buffer('num_batches_tracked', torch.tensor(0, dtype=torch.long))
        else:
            self.register_parameter('running_mean', None)
            self.register_parameter('running_var', None)
            self.register_parameter('num_batches_tracked', None)
        self.reset_parameters()
        self.reset_state()
        
    def reset_state(self):
        self.seq_idx = 0

    def reset_parameters(self):
        if self.track_running_stats and self.seq_len is not None:
            self.running_mean.zero_()
            self.running_var.fill_(1)
            self.num_batches_tracked.zero_()
        if self.affine:
            self.weight.data.fill_(1.0)
            self.bias.data.fill_(0.0)
            
    def forward(self, input, BN_start=None):
        if input.dim() != 3:
            raise ValueError('expected 3D input (got {}D input)'.format(input.dim()))
        if self.batch_first: input = input.transpose(0,1)
            
        input_t,n_batch,hidden_size =input.size()
        
        if self.track_running_stats and self.seq_len is None:
            self.seq_len = input_t
            self.register_buffer('running_mean', torch.zeros((input_t,hidden_size),device=input.device))
            self.register_buffer('running_var', torch.ones((input_t,hidden_size),device=input.device))
            
        if BN_start is None:
            if self.stateful:
                BN_start = self.seq_idx
            else:
                BN_start = 0

        exponential_average_factor = 0.0
        if self.training and self.track_running_stats:
            if self.num_batches_tracked is not None:
                self.num_batches_tracked += 1
                if self.momentum is None:  # use cumulative moving average
                    exponential_average_factor = 1.0 / float(self.num_batches_tracked)
                else:  # use exponential moving average
                    exponential_average_factor = self.momentum

        BN_stop = BN_start+input_t
        self.seq_idx = BN_stop #new starting point for next forward call 
        
        
        if self.training:
            mean = input.mean(1)
            var = input.var(1, unbiased=False)# use biased var in train
            
            if self.seq_len-BN_start > 0: #frame has to be in statistics window for updates
                with torch.no_grad():
                    self.running_mean[BN_start:BN_stop] = exponential_average_factor * mean[:self.seq_len-BN_start] + (1 - exponential_average_factor) * self.running_mean[BN_start:BN_stop]
                    self.running_var[BN_start:BN_stop] = exponential_average_factor * var[:self.seq_len-BN_start] * n_batch / (n_batch - 1)  + (1 - exponential_average_factor) * self.running_var[BN_start:BN_stop] # update running_var with unbiased var
        else:        
            mean = self.running_mean[BN_start:BN_stop]
            var = self.running_var[BN_start:BN_stop]
            
            #if elements outside of the statistics are requested, append the last element repeatedly
#             import pdb;pdb.set_trace()
            if BN_stop >= self.seq_len:
                cat_len = input_t - max(self.seq_len-BN_start,0)  # min(BN_stop-self.seq_len,self.seq_len) 
                mean = torch.cat((mean,self.running_mean[-1:].repeat(cat_len,1)))
                var = torch.cat((var,self.running_var[-1:].repeat(cat_len,1)))

        output = (input - mean[:, None, :]) / (torch.sqrt(var[:, None, :] + self.eps))
        if self.affine:
            output = output * self.weight[None, None, :] + self.bias[None, None, :]#[None, :, None, None]
            

        if self.batch_first: output = output.transpose(0,1)
            
        return output

# %% ../../nbs/03_models/00_layers.ipynb 8
class SeqLinear(nn.Module):
    
    def __init__(self,input_size,output_size,hidden_size=100,hidden_layer=1,act=Mish,batch_first=True):
        super().__init__()
        self.batch_first = batch_first
        def conv_act(inp,out): return nn.Sequential(nn.Conv1d(inp,out,1),act())
        
        if hidden_layer < 1:
            self.lin = nn.Conv1d(input_size,output_size,1)
        else:
            self.lin = nn.Sequential(conv_act(input_size,hidden_size),
                                     *[conv_act(hidden_size,hidden_size) for _ in range(hidden_layer-1)],
                                    nn.Conv1d(hidden_size,output_size,1))
            
    def forward(self, x):
        if not self.batch_first: x = x.transpose(0,1)
        out = self.lin(x.transpose(1,2)).transpose(1,2)
        
        if not self.batch_first: out = out.transpose(0,1)
        return out

# %% ../../nbs/03_models/00_layers.ipynb 10
class Normalizer1D(nn.Module):
    _epsilon = 1e-16

    def __init__(self, mean, std):
        super(Normalizer1D, self).__init__()
        self.register_buffer('std', std.clone().detach() + self._epsilon)
        self.register_buffer('mean', mean.clone().detach())

    def normalize(self, x):
        if x.device != self.mean.device:
            self.mean = self.mean.to(x.device)
            self.std = self.std.to(x.device)
        return (x-self.mean)/self.std

    def unnormalize(self, x):
        if x.device != self.mean.device:
            self.mean = self.mean.to(x.device)
            self.std = self.std.to(x.device)
        return x*self.std + self.mean

# %% ../../nbs/03_models/00_layers.ipynb 11
class AR_Model(nn.Module):
    '''
    Autoregressive model container which work autoregressively if the sequence y is not provided, otherwise it works as a normal model.
    This way it can be trained either with teacher forcing or with autoregression 
    '''
    def __init__(self,model,ar=True,stateful=False,model_has_state=False,return_state=False,out_sz=None):
        super().__init__()
        store_attr()
        if return_state and not model_has_state: raise ValueError('return_state=True requires model_has_state=True')
        self.norm = None
        self.y_init = None
    
    def init_normalize(self, batch,axes = [0,1]):
        x = batch[1]
        mean = x.mean(axes, keepdim=True)
        std = x.std(axes, keepdim=True)
        self.init_normalize_values(mean,std)
        
    def init_normalize_values(self, mean, std):
        self.norm = Normalizer1D(mean,std)
        
    def forward(self, u,y=None,h_init=None,ar=None):
        if ar is None: ar = self.ar

        if ar: #autoregressive mode
            y_e = []
            
            y_next = self.y_init if self.y_init is not None else torch.zeros(u.shape[0],1,self.out_sz).to(u.device)

            #two loops in the if clause to avoid the if inside the loop
            if self.model_has_state:
                h0 = h_init
                for u_in in u.split(1,dim=1):
                    x = torch.cat((u_in, y_next), dim=2)
                    y_next,h0 = self.model(x,h0)
                    y_e.append(y_next)
            else:
                for u_in in u.split(1,dim=1):
                    x = torch.cat((u_in, y_next), dim=2)
                    y_next = self.model(x)
                    y_e.append(y_next)
                    
            y_e = torch.cat(y_e,dim=1)
            
        else: #teacherforcing mode
            if y is None: raise ValueError('y must be provided in teacher forcing mode')

            x = torch.cat([u,y],dim=2)

            if self.model_has_state:
                y_e,h0 = self.model(x,h_init)
            else:
                y_e = self.model(x)
                
        if self.stateful:
            self.y_init = to_detach(y_e[:,-1:], cpu=False, gather=False) 

        if self.norm is not None: y_e = self.norm.unnormalize(y_e)
        return y_e if not self.return_state else (y_e,h0)
    
    def reset_state(self):
        self.y_init = None

# %% ../../nbs/03_models/00_layers.ipynb 14
from ..datasets.core import  extract_mean_std_from_dls
class NormalizedModel(nn.Module):
    def __init__(self, model, mean, std):
        super().__init__()
        self.model = model
        self.register_buffer('mean', torch.clone(mean))
        self.register_buffer('std', torch.clone(std))

    @classmethod
    def from_dls(cls,model,dls):
        mean, std = extract_mean_std_from_dls(dls)
        return cls(model,mean,std)

    def forward(self, xb):
        # Normalize x and then pass it through the model
        xb = (xb - self.mean) / self.std
        return self.model(xb)
