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

__all__ = ['BertEncoder', 'G_Net', 'D_Net', 'Attn_D_Net', 'is_d_attn', 'ExportedModel']

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

from .torch_utils import *
from .model_core import *

# Cell
class BertEncoder(nn.Module):
    def __init__(self, not_attn_ids=[]):
        super().__init__()
        self.bert = AutoModel.from_pretrained('albert-base-v1')
        self.not_attn_ids = not_attn_ids
        self.emb_sz = self.bert.config.hidden_size
    def forward(self, inp_ids):
        ''' inp_ids: (bs, max_seq_len)
            returns: sent_emb(bs, emb_sz), word_emb(bs, seq_len, emb_sz), src_mask(bs, seq_len) '''
        src_mask = isin(inp_ids, ids=self.not_attn_ids) # (bs, seq_len)

        attn_mask = (~src_mask).bool() # (bs, max_seq_len)
        word_emb = self.bert(inp_ids, attn_mask)[0] # (bs, max_seq_len, emb_sz)

        not_attn_mask = src_mask[..., None] # (bs, max_seq_len, 1)
        sent_emb = word_emb.masked_fill(not_attn_mask, 0.) # (bs, max_seq_len, emb_sz)
        sent_emb = sent_emb.mean(1) # (bs, emb_sz)

        return sent_emb, word_emb, src_mask
#     def gen_attention_mask(self, inp_ids):
#         device = inp_ids.device
#         one = torch.tensor(1, device=device)
#         zero = torch.tensor(0, device=device)
#         return torch.where(inp_ids==self.pad_id, zero, one)

# Internal Cell
CH_TABLE = [768, 512, 256, 128, 64]
NOISE_SZ = 100

# Internal Cell
class G_Init(nn.Module):
    def __init__(self, emb_sz):
        super().__init__()
        self.add_noise = nn.Sequential(AddNoise(NOISE_SZ), AddChannels(2)) # (bs, emb_sz+NOISE_SZ, 1, 1)
        self.conv = MultiSequential(
            conv_block(emb_sz+100, CH_TABLE[0], 4, 1, 3), # (bs, CH_TABLE[0], 4, 4)
            up_block(), # (bs, CH_TABLE[0], 8, 8)
            conv_block(CH_TABLE[0], CH_TABLE[1]), # (bs, CH_TABLE[1], 8, 8)
            up_block(), # (bs, CH_TABLE[1], 16, 16)
            conv_block(CH_TABLE[1], CH_TABLE[2]), # (bs, CH_TABLE[2], 16, 16)
            up_block(), # (bs, CH_TABLE[2], 32, 32)
            conv_block(CH_TABLE[2], CH_TABLE[3]), # (bs, CH_TABLE[3], 32, 32)
            up_block(), # (bs, CH_TABLE[3], 64, 64)
            conv_block(CH_TABLE[3], CH_TABLE[4]), # (bs, CH_TABLE[4], 64, 64)
        )
        self.to_rgb = to_rgb_block(CH_TABLE[-1]) # (bs, 3, 16, 16)
    def forward(self, sent_emb):
        ''' sent_emb: (bs, emb_sz)
            returns: img64(bs, 3, 64, 64), code64(bs, CH_TABLE[-1], 64, 64) '''
        code = self.add_noise(sent_emb) # (bs, emb_sz+100, 1, 1)
        code64 = self.conv(code) # (bs, CH_TABLE[-1], 64, 64)
        img64 = self.to_rgb(code64) # (bs, 3, 64, 64)
        return img64, code64

# Internal Cell
class G_Next(nn.Module):
    def __init__(self, emb_sz):
        super().__init__()
        self.conv = MultiSequential(
            AttnResBlock(emb_sz, CH_TABLE[-1]), # (bs, CH_TABLE[-1], inp_sz, inp_sz)
            ConvResBlock(CH_TABLE[-1]), # (bs, CH_TABLE[-1], inp_sz, inp_sz)
            up_block(), # (bs, CH_TABLE[-1], inp_sz*2, inp_sz*2)
            ConvResBlock(CH_TABLE[-1]), # (bs, CH_TABLE[-1], inp_sz*2, inp_sz*2)
        )
        self.to_rgb = to_rgb_block(CH_TABLE[-1]) # (bs, 3, inp_sz*2, inp_sz*2)
    def forward(self, code, word_emb, src_mask):
        ''' code: (bs, n_ch, inp_sz, inp_sz), word_emb: (bs, seq_len, emb_sz), src_mask: (bs, seq_len)
            returns: img(bs, 3, inp_sz*2, inp_sz*2), code(bs, CH_TABLE[-1], inp_sz*2, inp_sz*2) '''
        code = self.conv(code, word_emb, src_mask) # (bs, CH_TABLE[-1], inp_sz*2, inp_sz*2)
        img = self.to_rgb(code) # (bs, 3, inp_sz*2, inp_sz*2)
        return img, code

# Cell
class G_Net(nn.Module):
    def __init__(self, emb_sz):
        super().__init__()
        self.gs = nn.ModuleList([
            G_Init(emb_sz),
            G_Next(emb_sz),
            G_Next(emb_sz),
        ])
    def forward(self, sent_emb, word_emb, src_mask):
        """ sent_emb: (bs, emb_sz), word_emb: (bs, seq_len, emb_sz), src_mask: (bs, seq_len)
            returns: img64(bs, 3, 64, 64), img128(bs, 3, 128, 128), img256(bs, 3, 256, 256) """
        img, code = self.gs[0](sent_emb) # (bs, 3, 64, 64), (bs, CH_TABLE[-1], 64, 64)
        imgs = [img]
        for g in self.gs[1:]:
            img, code = g(code, word_emb, src_mask)
            imgs.append(img)
        return imgs

# Internal Cell
class UncondCls(nn.Module):
    def __init__(self):
        super().__init__()
        self.cls = nn.Sequential(
            conv2d(CH_TABLE[0], 1, 4, 1, 0, bias=False), # (bs, 1, 1, 1)
            nn.Flatten(), # (bs, 1)
        )
    def forward(self, sent_code):
        ''' sent_code: (bs, CH_TABLE[0], 4, 4), returns: (bs, 1) '''
        return self.cls(sent_code)


# Internal Cell
class CondCls(nn.Module):
    def __init__(self, emb_sz):
        super().__init__()
        self.up = nn.Upsample(4)
        self.cls = nn.Sequential(
            conv_block(CH_TABLE[0]+emb_sz, CH_TABLE[0]), # (bs, CH_TABLE[0], 4, 4)
            conv2d(CH_TABLE[0], 1, 4, 1, 0, bias=False), # (bs, 1, 1, 1)
            nn.Flatten(), # (bs, 1)
        )
    def forward(self, sent_code, sent_emb):
        ''' sent_code: (bs, CH_TABLE[0], 4, 4), sent_emb: (bs, emb_sz)
            returns: (bs, 1) '''
        sent_emb = self.up(sent_emb[..., None, None]) # (bs, emb_sz, 4, 4)
        sent_code = torch.cat([sent_code, sent_emb], 1) # (bs, emb_sz+CH_TABLE[0], 4, 4)
        return self.cls(sent_code)
class AttnCondCls(nn.Module):
    def __init__(self, emb_sz, attn_ch):
        super().__init__()
        self.cond_cls = CondCls(emb_sz)
        self.rev_attn_block = RevAttnBlock(emb_sz, attn_ch)
        self.word_cls = conv_block(emb_sz, 1, 1, 1, 0)
    def forward(self, sent_code, sent_emb, word_code, word_emb, src_mask):
        ''' sent_code: (bs, CH_TABLE[0], 4, 4), sent_emb: (bs, emb_sz),
            word_code: (bs, attn_ch, w, h), word_emb: (bs, seq_len, emb_sz), src_mask: (bs, seq_len)
            returns: (bs, 1) '''
        sent_logit = self.cond_cls(sent_code, sent_emb) # (bs, 1)

        word_code = self.rev_attn_block(word_emb, word_code) # (bs, seq_len, emb_sz)
        word_code = word_code.masked_fill(src_mask[..., None], 0.).mean(1)[..., None, None] # (bs, emb_sz, 1, 1)
        word_logit = self.word_cls(word_code).squeeze()[..., None] # (bs, 1)

        return torch.cat([sent_logit, word_logit], 1).mean(1, keepdim=True)

# Internal Cell
class D_64(nn.Module):
    def __init__(self, emb_sz):
        super().__init__()
        self.sent_code_m = nn.Sequential(
            from_rgb_block(CH_TABLE[-1]),  # (bs, CH_TABLE[-1], 64, 64)
            *self.create_conv(), # (bs, CH_TABLE[0], 4, 4)
        )
        self.uncond_cls = UncondCls()
        self.cond_cls = CondCls(emb_sz)
    @classmethod
    def create_conv(cls):
        return [
            conv_block(CH_TABLE[-1], CH_TABLE[-2], 4, 2, 1), # (bs, CH_TABLE[-2], 32, 32)
            conv_block(CH_TABLE[-2], CH_TABLE[-3], 4, 2, 1), # (bs, CH_TABLE[-3], 16, 16)
            conv_block(CH_TABLE[-3], CH_TABLE[-4], 4, 2, 1), # (bs, CH_TABLE[-4], 8, 8)
            conv_block(CH_TABLE[-4], CH_TABLE[-5], 4, 2, 1), # (bs, CH_TABLE[-5], 4, 4)
        ]
    def get_sent_code(self, img):
        """ img: (bs, 3, 64, 64), returns (bs, CH_TABLE[0], 4, 4) """
        return self.sent_code_m(img)
    def forward(self, img, sent_emb):
        """ img: (bs, 3, 64, 64), sent_emb: (bs, emb_sz)
            returns: uncond_logit(bs, 1), cond_logit(bs, 1) """
        sent_code = self.get_sent_code(img) # (bs, CH_TABLE[0], 4, 4)
        uncond_logit = self.uncond_cls(sent_code)
        cond_logit = self.cond_cls(sent_code, sent_emb)
        return uncond_logit, cond_logit # (bs, 1), (bs, 1)

# Internal Cell
class D_128(nn.Module):
    def __init__(self, emb_sz):
        super().__init__()
        self.sent_code_m = nn.Sequential(
            from_rgb_block(CH_TABLE[-1]),  # (bs, CH_TABLE[-1], 128, 128)
            *self.create_conv(), # (bs, CH_TABLE[0], 4, 4)
        )
        self.uncond_cls = UncondCls()
        self.cond_cls = CondCls(emb_sz)
    @classmethod
    def create_conv(cls):
        return [
            conv_block(CH_TABLE[-1], CH_TABLE[-1], 4, 2, 1), # (bs, CH_TABLE[-1], 64, 64)
            *D_64.create_conv(), # (bs, CH_TABLE[-3], 4, 4)
        ]
    def get_sent_code(self, img):
        """ img: (bs, 3, 128, 128), returns (bs, CH_TABLE[0], 4, 4) """
        return self.sent_code_m(img)
    def forward(self, img, sent_emb):
        """ img: (bs, 3, 128, 128), sent_emb: (bs, emb_sz)
            returns: uncond_logit(bs, 1), cond_logit(bs, 1) """
        sent_code = self.get_sent_code(img) # (bs, CH_TABLE[0], 4, 4)
        uncond_logit = self.uncond_cls(sent_code)
        cond_logit = self.cond_cls(sent_code, sent_emb)
        return uncond_logit, cond_logit # (bs, 1), (bs, 1)

# Internal Cell
class D_256(nn.Module):
    def __init__(self, emb_sz):
        super().__init__()
        self.sent_code_m = nn.Sequential(
            from_rgb_block(CH_TABLE[-1]),  # (bs, CH_TABLE[-1], 256, 256)
            *self.create_conv(), # (bs, CH_TABLE[0], 4, 4)
        )
        self.uncond_cls = UncondCls()
        self.cond_cls = CondCls(emb_sz)
    @classmethod
    def create_conv(cls):
        return [
            conv_block(CH_TABLE[-1], CH_TABLE[-1], 4, 2, 1), # (bs, CH_TABLE[-1], 128, 128)
            *D_128.create_conv(), # (bs, CH_TABLE[-3], 4, 4)
        ]
    def get_sent_code(self, img):
        """ img: (bs, 3, 256, 256), returns (bs, CH_TABLE[0], 4, 4) """
        return self.sent_code_m(img)
    def forward(self, img, sent_emb):
        """ img: (bs, 3, 256, 256), sent_emb: (bs, emb_sz)
            returns: uncond_logit(bs, 1), cond_logit(bs, 1) """
        sent_code = self.get_sent_code(img) # (bs, CH_TABLE[0], 4, 4)
        uncond_logit = self.uncond_cls(sent_code)
        cond_logit = self.cond_cls(sent_code, sent_emb)
        return uncond_logit, cond_logit # (bs, 1), (bs, 1)

# Internal Cell
class Attn_D_256(nn.Module):
    def __init__(self, emb_sz):
        super().__init__()
        self.word_code_m = nn.Sequential(
            from_rgb_block(CH_TABLE[-1]),  # (bs, CH_TABLE[-1], 256, 256)
            conv_block(CH_TABLE[-1], CH_TABLE[-1], 4, 2, 1), # (bs, CH_TABLE[-1], 128, 128)
            conv_block(CH_TABLE[-1], CH_TABLE[-1], 4, 2, 1), # (bs, CH_TABLE[-1], 64, 64)
            conv_block(CH_TABLE[-1], CH_TABLE[-2], 4, 2, 1), # (bs, CH_TABLE[-2], 32, 32)
            conv_block(CH_TABLE[-2], CH_TABLE[-3], 4, 2, 1), # (bs, CH_TABLE[-3], 16, 16)
        )
        self.sent_code_m = nn.Sequential(
            conv_block(CH_TABLE[-3], CH_TABLE[-4], 4, 2, 1), # (bs, CH_TABLE[-4], 8, 8)
            conv_block(CH_TABLE[-4], CH_TABLE[-5], 4, 2, 1), # (bs, CH_TABLE[-5], 4, 4)
        )

        self.uncond_cls = UncondCls()
        self.cond_cls = AttnCondCls(emb_sz, CH_TABLE[-3])
    def get_code(self, img):
        """ img: (bs, 3, 256, 256), returns (bs, CH_TABLE[0], 4, 4), (bs, CH_TABLE[-3], 16, 16) """
        word_code = self.word_code_m(img)
        sent_code = self.sent_code_m(word_code)
        return sent_code, word_code
    def forward(self, img, sent_emb, word_emb, src_mask):
        """ img: (bs, 3, 256, 256), sent_emb: (bs, emb_sz), word_emb: (bs, seq_len, emb_sz), src_mask: (bs, seq_len)
            returns: uncond_logit(bs, 1), cond_logit(bs, 1) """
        sent_code, word_code = self.get_code(img) # (bs, CH_TABLE[0], 4, 4), (bs, CH_TABLE[-3], 16, 16)
        uncond_logit = self.uncond_cls(sent_code)
        cond_logit = self.cond_cls(sent_code, sent_emb, word_code, word_emb, src_mask)
        return uncond_logit, cond_logit # (bs, 1), (bs, 1)

# Cell
class D_Net(nn.Module):
    def __init__(self, emb_sz):
        super().__init__()
        self.ds = nn.ModuleList([D_64(emb_sz), D_128(emb_sz), D_256(emb_sz)])
    def forward(self, imgs, sent_emb):
        """ imgs: [img64, img128, img256],
            returns: [(uncond_logit64, cond_logit64), (uncond_logit128, cond_logit128), (uncond_logit256, cond_logit256)] """
        assert len(imgs) == len(self.ds)
        logits = [d(img, sent_emb) for d, img in zip(self.ds, imgs)]
        return logits
class Attn_D_Net(nn.Module):
    def __init__(self, emb_sz):
        super().__init__()
        self.ds = nn.ModuleList([D_64(emb_sz), D_128(emb_sz), Attn_D_256(emb_sz)])
    def forward(self, imgs, sent_emb, word_emb, src_mask):
        """ imgs: [img64, img128, img256],
            returns: [(uncond_logit64, cond_logit64), (uncond_logit128, cond_logit128), (uncond_logit256, cond_logit256)] """
        assert len(imgs) == len(self.ds)
        logits = [d(img, sent_emb) for d, img in zip(self.ds[:-1], imgs[:-1])]
        logits.append(self.ds[-1](imgs[-1], sent_emb, word_emb, src_mask))
        return logits

# Cell
def is_d_attn(m):
    return True if isinstance(m, Attn_D_Net) or isinstance(m, Attn_D_256) else False

# Cell
class ExportedModel(nn.Module):
    def __init__(self, not_attn_ids):
        super().__init__()
        self.c_encoder = BertEncoder(not_attn_ids)
        self.g_net = G_Net(self.c_encoder.emb_sz)
    def forward(self, inp_ids):
        """ inp_ids: (bs, seq_len)
            (bs, 3, 256, 256) """
        sent_emb, word_emb, src_mask = self.c_encoder(inp_ids)
        imgs = self.g_net(sent_emb, word_emb, src_mask)
        return imgs[-1]
    @classmethod
    def from_pretrained(cls, path, device='cpu'):
        state = torch.load(path, map_location=device)
        not_attn_ids = state['not_attn_ids']
        m = cls(not_attn_ids)
        m.g_net.load_state_dict(state['g_net'])
        return m