# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/03a_model_core.ipynb (unless otherwise specified).

__all__ = ['conv2d', 'conv_block', 'multi_conv_block', 'ConvResBlock', 'multi_ConvResBlock', 'up_block',
           'multi_up_block', 'to_rgb_block', 'from_rgb_block', 'multi_from_rgb_block', 'AttnResBlock', 'RevAttnBlock',
           'AddNoise', 'AddChannels', 'NoOpLayer']

# Cell
import torch
import torch.nn as nn
from fastcore.all import *

from .torch_utils import *

# Cell
def conv2d(ni, nf, ks, stride, padding, bias=True):
    conv = nn.Conv2d(ni, nf, ks, stride, padding, bias=bias)
    nn.init.kaiming_normal_(conv.weight)
    return conv
def conv_block(ni, nf, ks=3, stride=1, padding=1):
    ''' (bs, ni, _, _) -> (bs, nf, _, _) '''
    return nn.Sequential(
        nn.utils.spectral_norm(conv2d(ni, nf, ks, stride, padding, bias=False)),
        nn.LeakyReLU(0.2),
    )
def multi_conv_block(ni, nf, ks=3, stride=1, padding=1):
    ''' (bs, ni, _, _), word_emb, src_mask -> (bs, nf, _, _), word_emb, src_mask '''
    return MultiWrapper(conv_block(ni, nf, ks, stride, padding), 3)

# Cell
class ConvResBlock(nn.Module):
    def __init__(self, nf):
        super().__init__()
        self.conv = nn.Sequential(
            conv_block(nf, nf*2),
            conv_block(nf*2, nf),
        )
    def forward(self, x):
        " x: (bs, nf, _, _), returns: (bs, nf, _, _) "
        out = self.conv(x)
        return x + out
def multi_ConvResBlock(nf):
    ''' (bs, nf, _, _), word_emb, src_mask -> (bs, nf, _, _), word_emb, src_mask '''
    return MultiWrapper(ConvResBlock(nf), 3)

# Cell
def up_block():
    ''' (bs, _, _, _) -> (bs, _, _*2, _*2) '''
    return nn.Upsample(scale_factor=2)
def multi_up_block():
    ''' (bs, _, _, _), word_emb, src_mask -> (bs, _, _*2, _*2), word_emb, src_mask'''
    return MultiWrapper(up_block(), 3)
# def down_block():
#     ''' (bs, _, _, _), word_emb, src_mask -> (bs, _, _//2, _//2) '''
#     return nn.AvgPool2d(2)
# def multi_down_block():
#     ''' (bs, _, _, _), word_emb, src_mask -> (bs, _, _//2, _//2), word_emb, src_mask'''
#     return MultiWrapper(down_block(), 3)

# Cell
def to_rgb_block(ni):
    ''' (bs, ni, _, _) -> (bs, 3, _, _) '''
    return nn.Sequential(
        conv2d(ni, 3, 1, 1, 0, bias=False),
        nn.Tanh(),
    )
# def multi_to_rgb_block(ni):
#     ''' (bs, ni, _, _), word_emb, src_mask -> (bs, 3, _, _), word_emb, src_mask '''
#     return MultiWrapper(to_rgb_block(ni), 3)

# Cell
def from_rgb_block(nf):
    ''' (bs, 3, _, _) -> (bs, nf, _, _) '''
    return nn.Sequential(
        conv2d(3, nf, 1, 1, 0, bias=False),
        nn.LeakyReLU(0.2),
    )
def multi_from_rgb_block(nf):
    ''' (bs, 3, _, _), word_emb, src_mask -> (bs, nf, _, _), word_emb, src_mask '''
    return MultiWrapper(from_rgb_block(nf), 3)

# Internal Cell
def simple_attn(tgt, src, src_mask=None):
    ''' tgt: (bs, tgt_seq_len, emb_sz), src: (bs, src_seq_len, emb_sz), src_mask: (bs, src_seq_len)
        returns: (bs, tgt_seq_len, emb_sz), (bs, tgt_seq_len, src_seq_len) '''
    attn_w = torch.bmm(tgt, src.permute(0, 2, 1)) # (bs, tgt_seq_len, src_seq_len)
    attn_w = attn_w.masked_fill_(src_mask.unsqueeze(1), -float('inf')) if src_mask is not None else attn_w # (bs, tgt_seq_len, src_seq_len)
    attn_w = torch.softmax(attn_w, 2) # (bs, tgt_seq_len, src_seq_len)

    attn_out = torch.bmm(attn_w, src) # (bs, tgt_seq_len, emb_sz)

    return attn_out, attn_w

# Internal Cell
class AttnBlock(nn.Module):
    def __init__(self, emb_sz, nc):
        super().__init__()
        self.emb_sz = emb_sz
        self.nc = nc
        self.conv1 = conv_block(emb_sz, nc, 1, 1, 0)
    def forward(self, x, word_emb, src_mask):
        """ Attention from x to word
            x: (bs, nc, w, h), word_emb: (bs, seq_len, emb_sz), src_mask: (bs, seq_len)
            attn_out: (bs, nc, w, h), self.attn_w: (bs, w, h, seq_len) """
        bs, nc, w, h = x.shape
        _, _, emb_sz = word_emb.shape
        assert (nc, emb_sz) == (self.nc, self.emb_sz)
        tgt = x.view(bs, nc, -1).permute(0, 2, 1).contiguous() # (bs, w*h, nc)

        # Use Conv1x1 to reduce word_emb dim, Do Not Use Linear!!!
        word_emb_t = word_emb.permute(0, 2, 1).unsqueeze(3).contiguous() # (bs, emb_sz, seq_len, 1)
        src = self.conv1(word_emb_t).squeeze(3) # (bs, nc, seq_len)
        src = src.permute(0, 2, 1).contiguous() # (bs, seq_len, nc)

        attn_out, attn_w = simple_attn(tgt, src, src_mask=src_mask) # (bs, w*h, nc), (bs, w*h, seq_len)
        attn_out = attn_out.permute(0, 2, 1).view(bs, nc, w, h) # (bs, nc, w, h)
        self.attn_w = attn_w.view(bs, w, h, -1) # (bs, w, h, seq_len)
        return attn_out

# Cell
class AttnResBlock(nn.Module):
    def __init__(self, emb_sz, nc, return_multi=False):
        super().__init__()
        self.return_multi = return_multi
        self.attn_block = AttnBlock(emb_sz, nc)
    def forward(self, x, word_emb, src_mask):
        """ x: (bs, nc, w, h), word_emb: (bs, seq_len, emb_sz), src_mask: (bs, seq_len)
            out: (bs, nc, w, h), word_emb, src_mask """
        attn_out = self.attn_block(x, word_emb, src_mask) # (bs, nc, w, h)
#         out = torch.cat([x, attn_out], dim=1) # (bs, 2*nc, w, h)
        out = x + attn_out # (bs, nc, w, h)
        return out if not self.return_multi else (out, word_emb, src_mask)

# Cell
class RevAttnBlock(nn.Module):
    def __init__(self, emb_sz, nc):
        super().__init__()
        self.emb_sz = emb_sz
        self.nc = nc
        self.conv1 = conv_block(nc, emb_sz, 1, 1, 0)
    def forward(self, word_emb, x):
        """ Attention from word to x
            word_emb: (bs, seq_len, emb_sz), x: (bs, nc, w, h)
            attn_out: (bs, seq_len, emb_sz), attn_w: (bs, )
        """
        bs, nc, w, h = x.shape
        _, seq_len, emb_sz = word_emb.shape
        assert (nc, emb_sz) == (self.nc, self.emb_sz)

        tgt = word_emb # (bs, seq_len, emb_sz)
        src = self.conv1(x).view(bs, emb_sz, -1).permute(0, 2, 1).contiguous() # (bs, w*h, emb_sz)

        attn_out, attn_w = simple_attn(tgt, src) # (bs, seq_len, emb_sz), (bs, seq_len, w*h)
        self.attn_w = attn_w.view(bs, seq_len, w, h) # (bs, seq_len, w, h)
        return attn_out

# Cell
class AddNoise(nn.Module):
    def __init__(self, noise_sz):
        super().__init__()
        self.noise_sz = noise_sz
        self.normal = torch.distributions.normal.Normal(0, torch.exp(torch.tensor(-1/np.pi)))
    def forward(self, x):
        ''' x: (bs, _), returns (bs, _+noise_sz) '''
        bs = x.shape[0]
        device = x.device
        noise = self.normal.sample((bs, self.noise_sz)).to(device)
#         noise = torch.randn(bs, self.noise_sz, device=device)
        out = torch.cat([x, noise], 1) # (bs, _+noise_sz)
        return out

# Cell
class AddChannels(nn.Module):
    "Add `n_dim` channels at the end of the input."
    def __init__(self, n_dim):
        super().__init__()
        self.n_dim=n_dim
    def forward(self, x):
        return x.view(*(list(x.shape)+[1]*self.n_dim))

# Cell
class NoOpLayer(nn.Module):
    def forward(self, *inp): return inp