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

__all__ = []

# Cell
import torch
import matplotlib.pyplot as plt
from fastcore.all import *

from .torch_utils import Normalizer
from .data_birds import BertTokenizer
from .model_birds import ExportedModel
from .constants import Birds

# Internal Cell
normalizer = Normalizer()
tokenizer = BertTokenizer()

# Internal Cell
def decode_img(img):
    "img: (bs, 3, _, _), returns: (bs, _, _, 3)"
    imgs = normalizer.decode(img).permute(0, 2, 3, 1).cpu()
    return imgs

# Internal Cell
def get_attn_w(model):
    "return: (bs, seq_len, w, h)"
    attn_w = model.g_net.gs[-1].conv[0].attn_block.attn_w # (bs, w, h, seq_len)
    return attn_w.permute(0, 3, 1, 2) # (bs, seq_len, w, h)

# Cell
@patch
@torch.no_grad()
def predict(self: ExportedModel, cap):
    ''' cap: 'this is a white bird with a grey cheek patch and a black eyebrow.'
        returns: img: (256, 256, 3), attn_w: (60, 128, 128) '''
    cap = tokenizer.encode(re.sub(r'[^\w\s]','',cap).lower())
    cap = torch.tensor([cap])
    self.eval()
    img = self(cap)
    img = decode_img(img)
    attn_w = get_attn_w(self)
    return img[0], attn_w[0]

# Internal Cell
def show_pred(img, attn_w, cap, not_attn_ids):
    ''' img: (256, 256, 3), attn_w: (60, 128, 128), cap: 'this is a white bird with a grey cheek patch and a black eyebrow.' '''
    idxs = []
    cap = tokenizer.encode(re.sub(r'[^\w\s]','',cap).lower())
    vals = tokenizer.encode(cap)
    for idx, val in enumerate(vals):
        if len(idxs) >= 19: break
        if val not in not_attn_ids: idxs.append(idx)

    ncols=4
    nrows=math.ceil((len(idxs)+1)/ncols)
    figsize = (ncols*4, nrows*4)
    fig, ax = plt.subplots(nrows=nrows, ncols=ncols, figsize=figsize)
    axs = ax.flatten()
    axs[0].imshow(img)
    for i in range(len(idxs)):
        idx = idxs[i]
        val = vals[idx]
        axs[i+1].set_title(tokenizer.tok.convert_ids_to_tokens([val]))
        axs[i+1].imshow(attn_w[idx])

# Cell
@patch
def pred_and_show(self: ExportedModel, cap):
    "cap: 'white hair yellow eyes' "
    img, attn_w = self.predict(cap)
    show_pred(img, attn_w, cap, self.c_encoder.not_attn_ids)