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

__all__ = ['AnimeHeadsTrainer', 'BirdsTrainer']

# Cell
import time
from tqdm.auto import tqdm
import torch
import torch.nn as nn
import torch.optim as optim
from torch import autograd
from kornia.augmentation.functional import random_hflip
from kornia.geometry.transform import resize
from fastcore.all import *

from .data_anime_heads import Datasets as AnimeHeadsDatasets, DataLoaders as AnimeHeadsDataLoaders
from .model_anime_heads import OnehotEncoder, G_Net as AnimeHeads_G_Net, D_Net as AnimeHeads_D_Net, Attn_D_Net as AnimeHeads_Attn_D_Net, is_d_attn as anime_heads_is_d_attn
from .data_birds import Datasets as BirdsDatasets, DataLoaders as BirdsDataLoaders
from .model_birds import BertEncoder, G_Net as Birds_G_Net, D_Net as Birds_D_Net, Attn_D_Net as Birds_Attn_D_Net, is_d_attn as birds_is_d_attn
from .torch_utils import *

# Cell
import os
if os.getenv('COLAB_TPU_ADDR'):
    import torch_xla
    import torch_xla.core.xla_model as xm

# Internal Cell
bce_loss = nn.BCEWithLogitsLoss()

# Internal Cell
def is_d_attn(m):
    return anime_heads_is_d_attn(m) or birds_is_d_attn(m)

# Internal Cell
def compute_gradient_penalty(d, true_img, fake_img, *other_inps):
    " true_img: (bs, 3, w, h), fake_img: (bs, 3, w, h), returns: uncond_gp, cond_gp "
    bs = true_img.shape[0]
    device = true_img.device

    alpha = torch.rand(bs, 1, 1, 1, device=device)
    interpolate = (alpha * true_img.detach() + ((1 - alpha) * fake_img.detach())).requires_grad_(True)
    uncond_logit, cond_logit = d(interpolate, *other_inps) # (bs, 1)

    uncond_grad = autograd.grad(
        outputs=uncond_logit,
        inputs=interpolate,
        grad_outputs=torch.ones_like(uncond_logit),
        create_graph=True,
        retain_graph=True,
        only_inputs=True,
    )[0].contiguous()
    uncond_grad = uncond_grad.view(uncond_grad.size(0), -1)
    uncond_gp = ((uncond_grad.norm(2, dim=1) - 1) ** 2).mean()

    cond_grad = autograd.grad(
        outputs=cond_logit,
        inputs=interpolate,
        grad_outputs=torch.ones_like(cond_logit),
        create_graph=True,
        retain_graph=True,
        only_inputs=True,
    )[0].contiguous()
    cond_grad = cond_grad.view(cond_grad.size(0), -1)
    cond_gp = ((cond_grad.norm(2, dim=1) - 1) ** 2).mean()

    return uncond_gp, cond_gp

# Internal Cell
def compute_d_loss(d, fake_img, true_img, sent_emb, gp_lambda=None):
    fake_img = fake_img.detach()
    bs = fake_img.shape[0]
    device = fake_img.device
    ones = torch.ones((bs, 1), device=device)
    zeros = torch.zeros((bs, 1), device=device)
    false_sent_emb = torch.cat([sent_emb[bs//2:], sent_emb[:bs//2]])

    true_sent_code = d.get_sent_code(true_img)
    fake_sent_code = d.get_sent_code(fake_img)
    uncond_true_logit = d.uncond_cls(true_sent_code) # (bs, 1)
    uncond_fake_logit = d.uncond_cls(fake_sent_code)
    cond_tt_logit = d.cond_cls(true_sent_code, sent_emb)
    cond_tf_logit = d.cond_cls(true_sent_code, false_sent_emb)
    cond_ft_logit = d.cond_cls(fake_sent_code, sent_emb)
    cond_ff_logit = d.cond_cls(fake_sent_code, false_sent_emb)

    uncond_true_loss = bce_loss(uncond_true_logit, ones)
    uncond_fake_loss = bce_loss(uncond_fake_logit, zeros)
    cond_tt_loss = bce_loss(cond_tt_logit, ones)
    cond_tf_loss = bce_loss(cond_tf_logit, zeros)
    cond_ft_loss = bce_loss(cond_ft_logit, zeros)
    cond_ff_loss = bce_loss(cond_ff_logit, zeros)
#     loss = (uncond_true_loss+cond_tt_loss)/2 + (uncond_fake_loss+cond_tf_loss+cond_ft_loss)/3
    loss = (uncond_true_loss+uncond_fake_loss) + (cond_tt_loss+(cond_tf_loss+cond_ft_loss+cond_ff_loss)/3)

    if gp_lambda:
        uncond_gp, cond_gp = compute_gradient_penalty(d, true_img, fake_img, sent_emb)
        uncond_gp_loss = uncond_gp*gp_lambda
        cond_gp_loss = cond_gp*gp_lambda
        loss += uncond_gp_loss+cond_gp_loss
    return loss

# Internal Cell
def compute_attn_d_loss(d, fake_img, true_img, sent_emb, word_emb, src_mask, gp_lambda=None):
    fake_img = fake_img.detach()
    bs = fake_img.shape[0]
    device = fake_img.device
    ones = torch.ones((bs, 1), device=device)
    zeros = torch.zeros((bs, 1), device=device)
    false_sent_emb = torch.cat([sent_emb[bs//2:], sent_emb[:bs//2]])
    false_word_emb = torch.cat([word_emb[bs//2:], word_emb[:bs//2]])
    false_src_mask = torch.cat([src_mask[bs//2:], src_mask[:bs//2]])

    true_sent_code, true_word_code = d.get_code(true_img)
    fake_sent_code, fake_word_code = d.get_code(fake_img)
    uncond_true_logit = d.uncond_cls(true_sent_code) # (bs, 1)
    uncond_fake_logit = d.uncond_cls(fake_sent_code)
    cond_tt_logit = d.cond_cls(true_sent_code, sent_emb, true_word_code, word_emb, src_mask)
    cond_tf_logit = d.cond_cls(true_sent_code, false_sent_emb, true_word_code, false_word_emb, false_src_mask)
    cond_ft_logit = d.cond_cls(fake_sent_code, sent_emb, fake_word_code, word_emb, src_mask)
    cond_ff_logit = d.cond_cls(fake_sent_code, false_sent_emb, fake_word_code, false_word_emb, false_src_mask)

    uncond_true_loss = bce_loss(uncond_true_logit, ones)
    uncond_fake_loss = bce_loss(uncond_fake_logit, zeros)
    cond_tt_loss = bce_loss(cond_tt_logit, ones)
    cond_tf_loss = bce_loss(cond_tf_logit, zeros)
    cond_ft_loss = bce_loss(cond_ft_logit, zeros)
    cond_ff_loss = bce_loss(cond_ff_logit, zeros)
#     loss = (uncond_true_loss+cond_tt_loss)/2 + (uncond_fake_loss+cond_tf_loss+cond_ft_loss)/3
    loss = (uncond_true_loss+uncond_fake_loss) + (cond_tt_loss+(cond_tf_loss+cond_ft_loss+cond_ff_loss)/3)

    if gp_lambda:
        uncond_gp, cond_gp = compute_gradient_penalty(d, true_img, fake_img, sent_emb, word_emb, src_mask)
        uncond_gp_loss = uncond_gp*gp_lambda
        cond_gp_loss = cond_gp*gp_lambda
        loss += uncond_gp_loss+cond_gp_loss
    return loss

# Internal Cell
def compute_g_loss(d_net, fake_imgs, *other_inps):
    bs = fake_imgs[0].shape[0]
    device = fake_imgs[0].device
    ones = torch.ones((bs, 1), device=device)
    zeros = torch.zeros((bs, 1), device=device)

    logits = d_net(fake_imgs, *other_inps) # list of (uncond_logit, cond_logit)
    l = 0
    for uncond_logit, cond_logit in logits:
        l += bce_loss(uncond_logit, ones)
        l += bce_loss(cond_logit, ones)
    return l

# Cell
class AnimeHeadsTrainer():
    def __init__(self, data_dir, bs, data_pct=1, g_lr=1e-3, d_lr=4e-3, gp_lambda=None, g_clip=None, device='cpu', d_attn=False):
        self.gp_lambda = gp_lambda
        self.g_clip = g_clip
        self.device = torch.device(device)
        self.normalizer = Normalizer(device=self.device)

        dsets = AnimeHeadsDatasets(data_dir, pct=data_pct, valid_pct=0)
        self.dls = AnimeHeadsDataLoaders(dsets, bs)

        vocab_sz = len(dsets.train.tokenizer.vocab)
        emb_sz = vocab_sz
        self.c_encoder = OnehotEncoder(vocab_sz).requires_grad_(False).eval().to(self.device)
        self.g_net = AnimeHeads_G_Net(emb_sz).to(self.device)
        self.d_net = AnimeHeads_D_Net(emb_sz).to(self.device) if d_attn==False else AnimeHeads_Attn_D_Net(emb_sz).to(self.device)

        self.optim_g = optim.Adam(self.g_net.parameters(), lr=g_lr, betas=(0.5, 0.999))
        self.optim_ds = [optim.Adam(d.parameters(), lr=d_lr, betas=(0.5, 0.999)) for d in self.d_net.ds]
    def after_batch_tfm(self, b):
        tag, img = b # (bs, 2), (bs, 64, 64, 3)
        img = img.permute(0, 3, 1, 2).float() # (bs, 3, 64, 64)
        img64 = random_hflip(img)
        img16 = resize(img64, size=(16, 16))
        img32 = resize(img64, size=(32, 32))
        img16 = self.normalizer.encode(img16) # (bs, 3, 16, 16)
        img32 = self.normalizer.encode(img32) # (bs, 3, 32, 32)
        img64 = self.normalizer.encode(img64) # (bs, 3, 64, 64)
        return tag, (img16, img32, img64)
    def export(self, path):
        state = {
            'emb_sz': self.c_encoder.emb_sz,
            'g_net': self.g_net.state_dict(),
        }
        torch.save(state, path)

# Cell
class BirdsTrainer():
    def __init__(self, data_dir, bs, data_pct=1, g_lr=1e-3, d_lr=4e-3, gp_lambda=None, g_clip=None, device='cpu', d_attn=False):
        self.gp_lambda = gp_lambda
        self.g_clip = g_clip
        self.device = torch.device(device)
        self.normalizer = Normalizer(device=self.device)

        dsets = BirdsDatasets(data_dir, pct=data_pct)
        self.dls = BirdsDataLoaders(dsets, bs)

        tok = dsets.train.tokenizer
        pad_id = tok.pad_id
        not_attn_ids = [tok.bos_id, tok.eos_id, tok.pad_id]
        self.c_encoder = BertEncoder(not_attn_ids).requires_grad_(False).eval().to(self.device)
        emb_sz = self.c_encoder.emb_sz
        self.g_net = Birds_G_Net(emb_sz).to(self.device)
        self.d_net = Birds_D_Net(emb_sz).to(self.device) if d_attn==False else Birds_Attn_D_Net(emb_sz).to(self.device)

        self.optim_g = optim.Adam(self.g_net.parameters(), lr=g_lr, betas=(0.5, 0.999))
        self.optim_ds = [optim.Adam(d.parameters(), lr=d_lr, betas=(0.5, 0.999)) for d in self.d_net.ds]
    def after_batch_tfm(self, b):
        cap, img = b # (bs, 60), (bs, 256, 256, 3)
        img = img.permute(0, 3, 1, 2).float() # (bs, 3, 256, 256)
        img256 = random_hflip(img)
        img64 = resize(img256, size=(64, 64))
        img128 = resize(img256, size=(128, 128))
        img64 = self.normalizer.encode(img64) # (bs, 3, 64, 64)
        img128 = self.normalizer.encode(img128) # (bs, 3, 128, 128)
        img256 = self.normalizer.encode(img256) # (bs, 3, 256, 256)
        return cap, (img64, img128, img256)
    def export(self, path):
        state = {
            'not_attn_ids': self.c_encoder.not_attn_ids,
            'g_net': self.g_net.state_dict(),
        }
        torch.save(state, path)

# Internal Cell
@patch
def set_device(self: [AnimeHeadsTrainer, BirdsTrainer], device='cpu'):
    self.device = torch.device(device)
    self.normalizer.set_device(self.device)
    self.c_encoder.to(self.device)
    self.g_net.to(self.device)
    self.d_net.to(self.device)
@patch
def set_model_mode(self: [AnimeHeadsTrainer, BirdsTrainer], mode):
    if mode=='D':
        self.g_net.eval()
        self.g_net.requires_grad_(False)
        self.d_net.train()
        self.d_net.requires_grad_(True)
    elif mode=='G':
        self.g_net.train()
        self.g_net.requires_grad_(True)
        self.d_net.eval()
        self.d_net.requires_grad_(False)
    else:
        raise Exception('Oops!!')
@patch
def save_checkpoint(self: [AnimeHeadsTrainer, BirdsTrainer], path):
    optim_ds = [optim_d.state_dict() for optim_d in self.optim_ds]
    state = {
        'g_net': self.g_net.state_dict(),
        'd_net': self.d_net.state_dict(),
        'optim_g': self.optim_g.state_dict(),
        'optim_ds': optim_ds,
    }
    torch.save(state, path)
@patch
def load_checkpoint(self: [AnimeHeadsTrainer, BirdsTrainer], path):
    state = torch.load(path, map_location=self.device)
    self.g_net.load_state_dict(state['g_net'])
    self.d_net.load_state_dict(state['d_net'])
    self.optim_g.load_state_dict(state['optim_g'])
    for optim_d, state in zip(self.optim_ds, state['optim_ds']):
        optim_d.load_state_dict(state)
@patch
@torch.no_grad()
def check_d(self: [AnimeHeadsTrainer, BirdsTrainer]):
    b = iter(self.dls.train).next()
    b = to_device(b, device='cuda')

    cap, true_imgs = self.after_batch_tfm(b)
    sent_emb, word_emb, src_mask = detach(self.c_encoder(cap))
    self.g_net.eval()
    fake_imgs = self.g_net(sent_emb, word_emb, src_mask)
    self.d_net.eval()

    t_logits = self.d_net(true_imgs, sent_emb) if not is_d_attn(self.d_net) else self.d_net(true_imgs, sent_emb, word_emb, src_mask)
    f_logits = self.d_net(fake_imgs, sent_emb) if not is_d_attn(self.d_net) else self.d_net(fake_imgs, sent_emb, word_emb, src_mask)

    ts = [torch.sigmoid(l) for logit in t_logits for l in logit] # t_uncond_s64, t_cond_s64, t_uncond_s128, t_cond_s128, t_uncond_s256, t_cond_s256
    fs = [torch.sigmoid(l) for logit in f_logits for l in logit]

    return list(zip(ts, fs))


# Internal Cell
@patch
def decode(self: [AnimeHeadsTrainer, BirdsTrainer], tag, true_img, *fake_imgs):
    caps = [self.dls.dsets.train.tokenizer.decode(t) for t in tag]
    true_imgs = self.normalizer.decode(true_img).permute(0, 2, 3, 1).cpu()
    fake_imgs = [self.normalizer.decode(fake_img).permute(0, 2, 3, 1).cpu() for fake_img in fake_imgs]
    return caps, true_imgs, fake_imgs
@patch
@torch.no_grad()
def fig_for_show(self: [AnimeHeadsTrainer, BirdsTrainer], is_valid=False):
    dl = self.dls.valid if is_valid else self.dls.train
    b = iter(dl).next()
    b = to_device(b, device=self.device)
    cap, true_imgs = self.after_batch_tfm(b)
    sent_emb, word_emb, src_mask = detach(self.c_encoder(cap))
    self.g_net.eval()
    fake_imgs = self.g_net(sent_emb, word_emb, src_mask)
    caps, true_img, fake_imgs = self.decode(cap, true_imgs[-1], *fake_imgs) # List of cap, (bs,), List of (bs,)

    ncols = len(fake_imgs)+1
    nrows = 4
    figsize = (ncols*4, nrows*4)
    fig, ax = plt.subplots(nrows=nrows, ncols=ncols, figsize=figsize)
    axs = ax.flatten()
    for i in range(nrows):
        ax[i][0].imshow(true_img[i])
        for j in range(len(fake_imgs)):
            ax[i][j+1].imshow(fake_imgs[j][i])
        ax[i][0].text(0, -1.5, caps[i])
    return fig
@patch
def show(self: [AnimeHeadsTrainer, BirdsTrainer], is_valid=False):
    fig = self.fig_for_show(is_valid)
    display(fig)
    plt.close()
@patch
def save_jpg(self: [AnimeHeadsTrainer, BirdsTrainer], path, is_valid=False):
    fig = self.fig_for_show(is_valid)
    fig.savefig(path)
    plt.close()

# Internal Cell
@patch
def train(self: [AnimeHeadsTrainer, BirdsTrainer], n_epoch, log_every=500, savejpg_every=None, jpg_path:str=None, saveck_every=None, ck_path:str=None):
    num_batch = 0
    g_losses = []
    d_lossess = [[] for _ in range(len(self.optim_ds))]
    total_start_t = time.time()
    log_start_t = time.time()

    mb = tqdm(range(n_epoch))
    for epoch in mb:
        pb = tqdm(self.dls.train, leave=False)
        for b in pb:
            num_batch += 1
            b = to_device(b, device=self.device)
            tag, true_imgs = self.after_batch_tfm(b)

            sent_emb, word_emb, src_mask = detach(self.c_encoder(tag))
            fake_imgs = self.g_net(sent_emb, word_emb, src_mask)

            # Train Discriminator
            self.set_model_mode('D')
            for i in range(len(true_imgs)):
                optim_d = self.optim_ds[i]
                d = self.d_net.ds[i]
                fake_img = fake_imgs[i]
                true_img = true_imgs[i]
                d_losses = d_lossess[i]

                optim_d.zero_grad()
                d_loss = compute_d_loss(d, fake_img, true_img, sent_emb, gp_lambda=self.gp_lambda) if not is_d_attn(d) \
                        else compute_attn_d_loss(d, fake_img, true_img, sent_emb, word_emb, src_mask, gp_lambda=self.gp_lambda)
                d_losses.append(d_loss.detach().cpu())
                d_loss.backward()
                if self.g_clip: torch.nn.utils.clip_grad_value_(d.parameters(), self.g_clip)
                optim_d.step() if self.device.type!='xla' else xm.optimizer_step(optim_d, barrier=True)

            # Train Generator
            self.set_model_mode('G')
            self.optim_g.zero_grad()
            g_loss = compute_g_loss(self.d_net, fake_imgs, sent_emb) if not is_d_attn(self.d_net) \
                    else compute_g_loss(self.d_net, fake_imgs, sent_emb, word_emb, src_mask)
            g_losses.append(g_loss.detach().cpu())
            pb.set_postfix(g_loss=g_loss.detach().cpu().numpy())
            g_loss.backward()
            if self.g_clip: torch.nn.utils.clip_grad_value_(self.g_net.parameters(), self.g_clip)
            self.optim_g.step() if self.device.type!='xla' else xm.optimizer_step(self.optim_g, barrier=True)

            # show or save something
            if savejpg_every and num_batch%savejpg_every == 0:
                self.save_jpg(path=jpg_path+f'-{num_batch//savejpg_every}.jpg')
            if saveck_every and num_batch%saveck_every == 0:
                self.save_checkpoint(path=ck_path+f'-{num_batch//saveck_every}.pt')
            if num_batch%log_every == 0:
                duration = time.time() - log_start_t
                msg = f'{num_batch//log_every}, time: {duration:.1f}s, g_loss: {torch.tensor(g_losses).mean():.4f}, '
                for i in range(len(d_lossess)):
                    msg += f'd_{i}_loss: {torch.tensor(d_lossess[i]).mean():.4f}, '
                tqdm.write(msg)
                g_losses = []
                d_lossess = [[] for _ in range(len(self.optim_ds))]
                log_start_t = time.time()
        pb.close()
    mb.close()
    tqdm.write(f'total_time: {(time.time()-total_start_t)/60:.1f}min')