"""Model initialization and architecture."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/source/01_models.ipynb.

# %% auto #0
__all__ = ['init_cnn', 'reparameterize', 'output_size_conv', 'output_size_after_n_conv', 'output_size_convt',
           'output_size_after_n_convt', 'View', 'VAEConv1d', 'sample_from_mix_gaussian', 'DilatedCausalConv1d',
           'ResidualBlock', 'VAEWaveNet']

# %% ../nbs/source/01_models.ipynb #65d993bb
from fastai.vision.all import *
from tqdm.auto import tqdm, trange

# %% ../nbs/source/01_models.ipynb #bd88bd96
import torch.nn.init as init

def init_cnn(m: nn.Module  # Module to initialize
            ):
    """Initialize module weights with kaiming normal in fan_out mode and bias to 0"""
    f = init.kaiming_normal_
    f = partial(f, mode='fan_out', nonlinearity='relu')

    if isinstance(m, (nn.Linear,  nn.Conv1d, nn.ConvTranspose1d)):
        f(m.weight)
        if m.bias is not None:  m.bias.data.zero_()

# %% ../nbs/source/01_models.ipynb #cddeae86
def reparameterize(mu:Tensor,     # Mean of the normal distribution, shape (batch_size, latent_dim)
                   logvar:Tensor  # Diagonal log variance of the normal distribution, shape (batch_size, latent_dim)
                  ) -> Tensor:    # Sampled latent z, shape (batch_size, latent_dim)
    """Sample latent tensor using the reparameterization trick:
    $z=\epsilon\sigma+\mu$, where $\epsilon\sim\mathcal{N}(0,1)$."""
    std = logvar.mul(0.5).exp_()
    eps = std.data.new(std.size()).normal_()
    return eps.mul(std).add_(mu)

# %% ../nbs/source/01_models.ipynb #d6a3eca3
def output_size_conv(input_size:int, kernel_size:int,
                     stride:int=1, padding:int=0, dilation:int=1)-> int:
    return int(np.floor((input_size+2*padding-dilation*(kernel_size-1)-1)/stride+1))

def output_size_after_n_conv(n:int, input_size:int, kernel_size:int,
                             stride:int=1, padding:int=0, dilation:int=1)-> int:
    out=[]
    for i in range(n):
        input_size = output_size_conv(input_size, kernel_size, stride, padding, dilation)
        out.append(input_size)
    return out

def output_size_convt(input_size:int, kernel_size:int,
                      stride:int=1, padding:int=0, output_padding:int=0, dilation:int=1)-> int:
    return int((input_size-1)*stride-2*padding+dilation*(kernel_size-1)+output_padding+1)

def output_size_after_n_convt(n:int, input_size:int, kernel_size:int,
                              stride:int=1, padding:int=0, output_padding:int=0, dilation:int=1)-> int:
    out=[]
    for i in range(n):
        input_size = output_size_convt(input_size, kernel_size, stride, padding, output_padding, dilation)
        out.append(input_size)
    return out

# %% ../nbs/source/01_models.ipynb #03a262d2
class View(nn.Module):
    """Use as (un)flattening layer"""
    def __init__(self, size:tuple[int]):
        super(View, self).__init__()
        self.size = size

    def forward(self, tensor):  return tensor.reshape(self.size)

# %% ../nbs/source/01_models.ipynb #7ecf2714
class VAEConv1d(nn.Module):
    """1-dimensional convolutional VAE architecture"""
    def __init__(self,
                 nf:list[int],   # number of filters
                 encoder:list,   # Encoder's dense layers sizes
                 decoder:list,   # Decoder's dense layers sizes
                 o_dim: int,     # input size (T)
                 nc_in: int=1,   # number of input channels
                 nc_out:int=6,   # number of output channels
                 z_dim: int=6,   # number of latent neurons
                 beta:float=0.,  # weight of the KLD loss
                 avg_size:int=24,# output size of the pooling layers
                 **kwargs
                ):
        super().__init__()
        store_attr()
        self.nc = nc_in
        self.input_size = input_size = o_dim;
        activation = nn.ReLU(inplace=True)

        # with Kernel_size=3, stride=1, padding=0
        output_size = output_size_after_n_conv(len(nf),input_size,3,1,0)
        self.output_size = output_size
        if output_size[-1]<5:
          raise ValueError(f"Convolutional output_size is {output_size}, considered too low.")

        # ndn: number of dense neurons after n=4 conv2d
        ndn = nf[-1]*output_size[-1]

        conv_sizes = [nc_in] + nf
        conv_activations = [nn.ReLU(inplace=True) for i in range(len(conv_sizes) - 1)]
        conv_layers = [nn.Conv1d(dim_in, dim_out, 3, 1, 0) for dim_in, dim_out in zip(conv_sizes[:-1], conv_sizes[1:])]
        conv = [item for pair in zip(conv_layers, conv_activations) for item in pair]
        avg_pooling = avg_size
        conv.append(AdaptiveConcatPool1d(avg_pooling))
        avg_pooling = avg_size*2  # twice due to the concat of both poolings
        self.scale_factor = output_size[-1]/avg_pooling
        output_size.append(avg_pooling)
        ndn = nf[-1]*output_size[-1]

        # add flatten layer:
        conv.append(View((-1, ndn)))

        encoder_sizes = [ndn] + encoder + [z_dim*2]
        encoder_activations = [nn.ReLU(inplace=True) for i in range(len(encoder_sizes) - 2)] +[None]
        encoder_layers = [nn.Linear(dim_in, dim_out) for dim_in, dim_out in zip(encoder_sizes[:-1], encoder_sizes[1:])]
        encoder = [item for pair in zip(encoder_layers, encoder_activations) for item in pair]

        if encoder[-1]==None: encoder = encoder[:-1] # clean None

        encoder = conv + encoder

        self.encoder = nn.Sequential(*encoder)

        decoder_sizes = [z_dim] + decoder + [ndn]

        decoder_activations = [nn.ReLU(inplace=True) for i in range(len(decoder_sizes) - 1)]
        decoder_layers = [nn.Linear(dim_in, dim_out) for dim_in, dim_out in
                          zip(decoder_sizes[:-1], decoder_sizes[1:])]
        decoder = [item for pair in zip(decoder_layers, decoder_activations) for item in pair]

        # Unflatten
        decoder.append(View((-1, nf[-1], output_size[-1])))  # BS, nf, size
        self.decoder = nn.Sequential(*decoder)
        # Interpolation to reverse poolings

        convt_sizes = nf[::-1] + [nc_out]
        convt_activations = [nn.ReLU(inplace=True) for i in range(len(convt_sizes) - 2)] + [nn.ReLU(inplace=True)]
        convt_layers = [nn.ConvTranspose1d(dim_in, dim_out, 3, 1, 0) for
                        dim_in, dim_out in zip(convt_sizes[:-1], convt_sizes[1:])]
        convt = [item for pair in zip(convt_layers, convt_activations) for item in pair]

        if convt[-1]==None: convt = convt[:-1]   # clean None

        self.convt = nn.Sequential(*convt)

    def forward(self, x, question=None):
        distributions = self._encode(x)
        mu     = distributions[:, :self.z_dim]
        logvar = distributions[:, self.z_dim:]
        z      = reparameterize(mu, logvar)  # B, |z|
        x_logit= self._decode(z)             # B, channels, T
        return x_logit, mu, logvar, z

    def _encode(self, x):  return self.encoder(x)  # Aproximate posterior q(z|x)

    def _decode(self, z, question=None):
        if question is not None:  qz = torch.cat((question, z), dim=-1)
        else:  qz = z
        x = self.decoder(qz)
        if not hasattr(self,'Tgen') or self.Tgen is None:
            x = F.interpolate(x, scale_factor=self.scale_factor)
        else:
            x = F.interpolate(x, size=self.Tgen-(self.input_size-self.output_size[-2]))
        x = self.convt(x)
        return x

# %% ../nbs/source/01_models.ipynb #1a326b01
from torch.distributions import Normal

# %% ../nbs/source/01_models.ipynb #d471b6ce
# modified from https://github.com/r9y9/wavenet_vocoder/blob/master/wavenet_vocoder/mixture.py#L221
def sample_from_mix_gaussian(y:Tensor, # Mixture of Gaussians parameters. Shape (B x C x T)
                             log_scale_min:float=-12.0, # Log scale minimum value.
                            )-> Tensor:
    """Sample from (discretized) mixture of gaussian distributions."""
    C = y.size(1)
    if C == 2:
        nr_mix = 1
    else:
        assert y.size(1) % 3 == 0
        nr_mix = y.size(1) // 3

    # B x T x C
    y = y.transpose(1, 2)

    if C == 2:
        logit_probs = None
    else:
        logit_probs = y[:, :, :nr_mix]

    if nr_mix > 1:
        # sample mixture indicator from softmax
        temp = logit_probs.data.new(logit_probs.size()).uniform_(
            1e-5, 1.0 - 1e-5)
        temp = logit_probs.data - torch.log(-torch.log(temp))
        _, argmax = temp.max(dim=-1)

        # (B, T) -> (B, T, nr_mix)
        one_hot = to_one_hot(argmax, nr_mix)

        # Select means and log scales
        means = torch.sum(y[:, :, nr_mix:2 * nr_mix] * one_hot, dim=-1)
        log_scales = torch.sum(y[:, :, 2 * nr_mix:3 * nr_mix] * one_hot,
                               dim=-1)
    else:
        if C == 2:
            means, log_scales = y[:, :, 0], y[:, :, 1]
        elif C == 3:
            means, log_scales = y[:, :, 1], y[:, :, 2]
        else:
            assert False, f"shouldn't happen C={C}"

    # bound log_scales to avoid -inf going to scale=0 and Normal raising an error
    log_scales = torch.clamp(log_scales, min=log_scale_min)
    scales = torch.exp(log_scales)
    dist = Normal(loc=means, scale=scales)
    x = dist.sample()
    # x = torch.clamp(x, min=-1.0, max=1.0) # [-1,1] removed to have no truncation in sampling
    return x

# %% ../nbs/source/01_models.ipynb #b0abaab0
# Implementation of WaveNet modified from
# https://github.com/ryujaehun/wavenet
# and Pieter Abbeel course
# https://github.com/rll/deepul/blob/master/demos/lecture2_autoregressive_models_demos.ipynb

class DilatedCausalConv1d(nn.Module):
    """Dilated causal convolution for WaveNet"""
    def __init__(self, mask_type:str, in_channels:int, out_channels:int,
                 kernel_size:int=2, dilation:int=1, bias:bool=True, use_pad:bool=True):
        super().__init__()
        self.use_pad = use_pad
        self.conv = nn.Conv1d(in_channels,out_channels,
                              kernel_size=kernel_size,dilation=dilation,
                              padding=0,bias=bias)
        self.mask_type,self.dilation = mask_type,dilation
        assert mask_type in ['A', 'B']
        self.left_padding = dilation * (kernel_size - 1) if use_pad else 0

    def forward(self, x):
        if self.mask_type == 'A':
            # Pad only on the left side   and discard the last data-point that implements the self shift
            # Hence, we are able to make loss(output[...:-1], input) or input[...:-1] with loss(output, input)
            if self.use_pad:
                return self.conv( F.pad(x[..., :-1],[self.left_padding + 1, 0]) )
            else:
                return self.conv(x[...,:-1])
            # No pad and discard the last data-point that implements the self shift
            # Hence, we are able to make loss(output, input)
        else:
            if self.use_pad:
                return self.conv(F.pad(x, [self.left_padding, 0]))
            else:
                return self.conv(x)

# %% ../nbs/source/01_models.ipynb #8aaed028
class ResidualBlock(nn.Module):
    """Residual block with conditions and gate mechanism"""
    def __init__(self, res_channels:int, skip_channels:int, kernel_size:int, dilation:int,
                 c_channels:int=0, g_channels:int=0, bias:bool=True, use_pad:bool=True):
        super().__init__()
        self.dilated = DilatedCausalConv1d('B',res_channels,2 * res_channels,
                                           kernel_size=kernel_size,dilation=dilation,
                                           bias=True, use_pad=use_pad)
        # conditions:
        # local conditioning
        if c_channels > 0:
            self.conv_c = nn.Conv1d(
                c_channels, 2 * res_channels, kernel_size=1,
                bias=False)  # they are a bias in themselves
        else:
            self.conv_c = None
        # global conditioning
        if g_channels > 0:
            self.conv_g = nn.Conv1d(
                g_channels, 2 * res_channels, kernel_size=1,
                bias=False)  # they are a bias in themselves
        else:
            self.conv_g = None

        self.conv_res = nn.Conv1d(res_channels,res_channels,
                                  kernel_size=1,bias=bias)
        self.conv_skip = nn.Conv1d(res_channels,skip_channels,
                                   kernel_size=1,bias=bias)

    def forward(self, x, c=None, g=None):
        output = self.dilated(x)
        # PixelCNN gate:
        # split channels
        channel_dim = 1
        o1, o2 = output.chunk(2, dim=channel_dim) # also named filter and gate

        # local conditioning
        if c is not None:
            if self.conv_c is not None:
                c = self.conv_c(c)
                c = c[...,-output.shape[-1]:]
                c1, c2 = c.chunk(2, dim=channel_dim)
                o1, o2 = o1 + c1, o2 + c2

        # global conditioning
        if g is not None:
            assert self.conv_g is not None, 'You should enable global conditioning with g_channels>0.'
            g = self.conv_g(g)
            g = g[...,-output.shape[-1]:]
            g1, g2 = g.chunk(2, dim=channel_dim)
            o1, o2 = o1 + g1, o2 + g2
        else: assert self.conv_g is None

        # GLU from https://arxiv.org/abs/1612.08083
        # gate # id * sigmoid # tanh * sigmoid from
        gated = o1 * torch.sigmoid(o2)  # torch.tanh(o1) * torch.sigmoid(o2)
        # skip conv 1x1
        skip = self.conv_skip(gated)
        # residual conv 1x1
        residual = self.conv_res(gated)
        # residual sum to block input
        residual = x[...,-residual.shape[-1]:] + residual

        return residual, skip

# %% ../nbs/source/01_models.ipynb #b6c72853
class VAEWaveNet(nn.Module):
    """VAE with autoregressive decoder"""
    name = "VAEWaveNet"
    def __init__(self,
                 in_channels:  int=1,   # input channels
                 res_channels: int=16,  # residual block channels
                 skip_channels:int=16,  # skip connection channels
                 c_channels:   int=6,   # local conditioning (time-wise)
                 g_channels:   int=0,   # global conditioning (the same for the whole sequence)
                 out_channels: int=1,   # output channels
                 res_kernel_size:int=3, # kernel_size of dilated layers in residual blocks
                 layer_size:   int=4,   # largest dilation is 2^layer_size
                 stack_size:   int=1,   # number of layers stacks
                 out_distribution:str="normal",
                 discrete_channels:int=256,
                 num_mixtures: int=1,   # 1=no mixture
                 use_pad: bool=False,
                 weight_norm:bool=False,
                 **kwargs):
        store_attr()
        super().__init__()
        self.init_channels = in_channels
        self.out_distribution = out_distribution.lower()
        if self.out_distribution == "normal":
            out_channels = out_channels * num_mixtures * 3  # (p_i, mu, logstd)
            # p_i is the probability of each distribution in the mixture
        elif self.out_distribution == "logistic":
            out_channels = out_channels * num_mixtures * 3
        elif self.out_distribution == "categorical":
            out_channels = out_channels * discrete_channels
        elif self.out_distribution == "binary":
            out_channels = out_channels
        self.out_channels = out_channels

        # x->vae->c
        if c_channels>0:
            self.vae = VAEConv1d(**kwargs)
            self.beta = self.vae.beta; self.z_dim = self.vae.z_dim
            self.input_size = self.vae.input_size
            self.output_size = self.vae.output_size

        # x-> init -> causal
        self.init_conv = nn.Conv1d(in_channels, res_channels,
                                   kernel_size=1, bias=True
                                   )
        self.causal = DilatedCausalConv1d('A',res_channels,res_channels,
                                          kernel_size=2,dilation=1, use_pad=use_pad)

        # x_init, c, g -> ResStack...->skips sumed-> out_conv -> (softmax or equivalent in the loss)-> prob
        self.res_stack = nn.ModuleList(
            [ResidualBlock(res_channels, skip_channels, res_kernel_size,
                           2**i, c_channels,g_channels, use_pad=use_pad)
             for i in range(layer_size)] * stack_size)
        self.out_conv = nn.Sequential(*[nn.ReLU(inplace=True),
                                        nn.Conv1d(skip_channels, skip_channels, kernel_size=1),
                                        nn.ReLU(inplace=True),
                                        nn.Conv1d(skip_channels, out_channels, kernel_size=1,bias=True)])
        if weight_norm:
            for m in flatten_model(self):
                if has_params(m):    m = nn.utils.weight_norm(m)
        self.receptive_field = self.receptive_field_size()

        self.weight_init()  # custom weight init for VAE (to prevent logvar explosion)

    def weight_init(self):
        for m in flatten_model(self):
            if has_params(m):    init_cnn(m,)  #  method='kaiming'
        # to avoid explosion in VAE latent logvar we initialize it to zero
        if hasattr(self,'vae'):
            self.vae.encoder[-1].weight.data[self.z_dim:,:].fill_(0)  # logvar_w=0
            self.vae.encoder[-1].bias.data.fill_(0)  # mu and logvar biases    = 0
        # change to fan_in (16) for mu_out because fan_out (3) is very little
        fan = init._calculate_correct_fan(self.out_conv[-1].weight.data, mode='fan_in')
        gain = init.calculate_gain('relu', 0.)
        std = gain / math.sqrt(fan)
        init.trunc_normal_(self.out_conv[-1].weight.data, std=std)

    def forward(self, x:Tensor, c=None, g=None):
        batch_size, x_size = x.shape[0], x.shape[2:]
        # input x
        if x.dim() >= 4:
            # Flatten
            x = x.reshape(batch_size, self.init_channels,-1)  # B, in_ch, -1=28*28
        x_ = x.detach().clone()
        # x -> VAE -> c
        c, mu, logvar, z = self.vae(x) if c is None and self.c_channels>0 else (c, None, None, None)
        if self.c_channels==0: c=None # supress c input from fastai dataloader when no VAE

        if g is not None:
            # expand to all time steps
            g = g.unsqueeze(-1).expand(
                batch_size, -1, np.prod(x_size))  # B, C, 28*28

        # x-> init_conv-> causal_conv
        output = self.init_conv(x_)   # B, res, 28*28
        output = self.causal(output)  # B, res, 28*28
        # output = self.res_stack(output)         # B, res, 28*28

        skip = torch.zeros((batch_size, self.skip_channels, *output.shape[2:]), device=x.device)
        sqrt_half = 0.5**0.5
        # -> dilated residual stack...-> skips sum
        for res_block in self.res_stack:
            res_out, skip_out = res_block(output, c, g)
            output = (output[...,-res_out.shape[-1]:] + res_out) * sqrt_half
            skip = skip[...,-skip_out.shape[-1]:] + skip_out
        skip = skip * (1. / len(self.res_stack))**0.5
        output = skip
        # out_conv -> (softmax or equivalent in the loss)-> prob
        output = self.out_conv(output)  # B, out, 28*28

        if self.c_channels>0: return output.reshape(batch_size, self.out_channels, -1), mu, logvar, c, z
        else: return output.reshape(batch_size, self.out_channels, -1)

    def set_use_pad(self, v: bool):
        """
        Recursively sets the `use_pad` attribute
        for all instances of `DilatedCausalConv1d` in a model.
        """
        for m in self.modules():
            if isinstance(m, DilatedCausalConv1d):  m.use_pad = v

    def sample_batch(self, n:int,bs:int=100,T:int=200, c=None,g=None, device=None):
        """Generates new samples from zeros and conditions in batches."""
        max_value, min_value = 1e3,-1e3
        T_in = 1  # additional initial points for predictions
        # assure sampling correctly
        if c is None:
            assert self.c_channels==0
        else:
            assert self.c_channels>0
        if g is None:
            assert self.g_channels==0
        else:
            assert self.g_channels>0
        device = device if device is not None else next(self.parameters()).device
        if self.use_pad: self.set_use_pad(False) # avoid padding in forward pass

        with torch.no_grad():
            samples_all = torch.tensor([], device=device,dtype=torch.float) # all batches
            for i in trange(n//bs):
                samples = torch.zeros(bs, # one batch of size bs
                                      self.in_channels,
                                      T+ self.receptive_field + T_in, # desired T sample, RF of zeros, Initial T
                                      device=device,
                                      dtype=torch.float)
                c_b = c[i*bs:(i+1)*bs] if c is not None else c
                g_b = g[i*bs:(i+1)*bs] if g is not None else g

                for r in range(T+ T_in): # sequential loop
                    c_b_ = c_b[...,r:r+self.receptive_field+1] if c is not None else c_b # handle None in c_b
                    pred = self(samples[...,r:r+self.receptive_field+1],
                                c=c_b_,g=g_b)
                    if self.c_channels>0:  # VAE + WaveNet
                        pred, mu, logvar, c_, z = pred
                    # remember model already processes input[...:-1] hence +1 to include last index (r + RF)
                    # else: p_logit, mu, logvar = pred
                    p_logit = pred
                    logits  = p_logit.squeeze(-1) # the next timestep

                    if self.out_distribution == "normal":
                        logits = logits.unsqueeze(-1)
                        sample_ = sample_from_mix_gaussian(logits)
                    elif self.out_distribution == "logistic":
                        logits = logits.unsqueeze(-1)
                        sample_ = sample_from_discretized_mix_logistic(logits)

                    elif self.out_distribution == "binary":
                        probs = torch.sigmoid(logits)
                        sample_ = torch.bernoulli(probs)
                    elif self.out_distribution == "categorical":
                        logits = logits.unsqueeze(-1)
                        sample_ = sample_from_multinomial(logits, dim=1)
                        sample_ = class_q_value(sample_, self.discrete_channels, )
                    else:
                        raise NotImplementedError(
                            "Current output distributions are: [logistic, normal, categorical, binary]"
                        )

                    # clamp to avoid divergence
                    sample_ = torch.clamp(sample_, min=min_value,max=max_value)

                    samples[..., r+self.receptive_field] = sample_
                samples=samples[..., self.receptive_field+ T_in:] # cut zeros and T_in
                # takes only sampled times from t=RF+T_in
                samples_all = torch.cat((samples_all,samples),dim=0)

        if self.use_pad: self.set_use_pad(True)  #  Leave the model as it was

        return samples_all.cpu(), logits

    def receptive_field_size(self, show:bool=False):
        """Compute theoretical receptive field size of the model."""
        # (L_in−1)×stride−2×more_than_causal_padding+dilation×(kernel_size−1)+1
        receptive_field = 1 # this accounts for oneself shift.
        # RF past points + oneself, the RF definition is past seen points and cannot be oneself because of causality
        if show: print(f"{receptive_field=}")
        for r in self.res_stack[::-1]:
            receptive_field += r.dilated.conv.dilation[0]*(r.dilated.conv.kernel_size[0]-1)
            if show: print(f"{receptive_field=}, dilation={r.dilated.conv.dilation[0]}, kernel={r.dilated.conv.kernel_size[0]}")

        receptive_field += self.causal.conv.dilation[0]*(self.causal.conv.kernel_size[0]-1)
        if show: print(f"{receptive_field=}")
        return receptive_field
