import torch
import torch.nn as nn

import numpy as np

import json

import captioning.utils.opts as opts
import captioning.models as models
import captioning.utils.misc as utils

import pytorch_lightning as pl

import gradio as gr


# Checkpoint class
class ModelCheckpoint(pl.callbacks.ModelCheckpoint):

    def on_keyboard_interrupt(self, trainer, pl_module):
        # Save model when keyboard interrupt
        filepath = os.path.join(self.dirpath, self.prefix + 'interrupt.ckpt')
        self._save_model(filepath)
        
device = 'cpu' #@param ["cuda", "cpu"] {allow-input: true}

reward = 'clips_grammar' #@param ["mle", "cider", "clips", "cider_clips", "clips_grammar"] {allow-input: true}

if reward == 'mle':
    cfg = f'./configs/phase1/clipRN50_{reward}.yml'
else:
    cfg = f'./configs/phase2/clipRN50_{reward}.yml'

print("Loading cfg from", cfg)

opt = opts.parse_opt(parse=False, cfg=cfg)

import gdown

if reward == "mle":
  url = "https://drive.google.com/drive/folders/1hfHWDn5iXsdjB63E5zdZBAoRLWHQC3LD"
elif reward == "cider":
  url = "https://drive.google.com/drive/folders/1MnSmCd8HFnBvQq_4K-q4vsVkzEw0OIOs"
elif reward == "clips":
  url = "https://drive.google.com/drive/folders/1toceycN-qilHsbYjKalBLtHJck1acQVe"
elif reward == "cider_clips":
  url = "https://drive.google.com/drive/folders/1toceycN-qilHsbYjKalBLtHJck1acQVe"
elif reward == "clips_grammar":
  url = "https://drive.google.com/drive/folders/1nSX9aS7pPK4-OTHYtsUD_uEkwIQVIV7W"
gdown.download_folder(url, quiet=True, use_cookies=False, output="save/")

url = "https://drive.google.com/uc?id=1HNRE1MYO9wxmtMHLC8zURraoNFu157Dp"
gdown.download(url, quiet=True, use_cookies=False, output="data/")

dict_json = json.load(open('./data/cocotalk.json'))
print(dict_json.keys())

ix_to_word = dict_json['ix_to_word']
vocab_size = len(ix_to_word)
print('vocab size:', vocab_size)

seq_length = 1

opt.vocab_size = vocab_size
opt.seq_length = seq_length

opt.batch_size = 1
opt.vocab = ix_to_word
# opt.use_grammar = False

model = models.setup(opt)
del opt.vocab

ckpt_path = opt.checkpoint_path + '-last.ckpt'

print("Loading checkpoint from", ckpt_path)
raw_state_dict = torch.load(
    ckpt_path,
    map_location=device)

strict = True

state_dict = raw_state_dict['state_dict']

if '_vocab' in state_dict:
    model.vocab = utils.deserialize(state_dict['_vocab'])
    del state_dict['_vocab']
elif strict:
    raise KeyError
if '_opt' in state_dict:
    saved_model_opt = utils.deserialize(state_dict['_opt'])
    del state_dict['_opt']
    # Make sure the saved opt is compatible with the curren topt
    need_be_same = ["caption_model",
                    "rnn_type", "rnn_size", "num_layers"]
    for checkme in need_be_same:
        if getattr(saved_model_opt, checkme) in ['updown', 'topdown'] and \
                getattr(opt, checkme) in ['updown', 'topdown']:
            continue
        assert getattr(saved_model_opt, checkme) == getattr(
            opt, checkme), "Command line argument and saved model disagree on '%s' " % checkme
elif strict:
    raise KeyError
res = model.load_state_dict(state_dict, strict)
print(res)

model = model.to(device)
model.eval();

import clip
from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
from PIL import Image
from timm.models.vision_transformer import resize_pos_embed

clip_model, clip_transform = clip.load("RN50", jit=False, device=device)

preprocess = Compose([
    Resize((448, 448), interpolation=Image.BICUBIC),
    CenterCrop((448, 448)),
    ToTensor()
])

image_mean = torch.Tensor([0.48145466, 0.4578275, 0.40821073]).to(device).reshape(3, 1, 1)
image_std = torch.Tensor([0.26862954, 0.26130258, 0.27577711]).to(device).reshape(3, 1, 1)

num_patches = 196 #600 * 1000 // 32 // 32
pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, clip_model.visual.attnpool.positional_embedding.shape[-1],  device=device),)
pos_embed.weight = resize_pos_embed(clip_model.visual.attnpool.positional_embedding.unsqueeze(0), pos_embed)
clip_model.visual.attnpool.positional_embedding = pos_embed

def inference(img):

  with torch.no_grad():
      image = preprocess(img)
      image = torch.tensor(np.stack([image])).to(device)
      image -= image_mean
      image /= image_std
      
      tmp_att, tmp_fc = clip_model.encode_image(image)
      tmp_att = tmp_att[0].permute(1, 2, 0)
      tmp_fc = tmp_fc[0]
      
      att_feat = tmp_att
      fc_feat = tmp_fc
      
      
  # Inference configurations
  eval_kwargs = {}
  eval_kwargs.update(vars(opt))
  
  verbose = eval_kwargs.get('verbose', True)
  verbose_beam = eval_kwargs.get('verbose_beam', 0)
  verbose_loss = eval_kwargs.get('verbose_loss', 1)
  
  # dataset = eval_kwargs.get('dataset', 'coco')
  beam_size = eval_kwargs.get('beam_size', 1)
  sample_n = eval_kwargs.get('sample_n', 1)
  remove_bad_endings = eval_kwargs.get('remove_bad_endings', 0)
  
  with torch.no_grad():
      fc_feats = torch.zeros((1,0)).to(device)
      att_feats = att_feat.view(1, 196, 2048).float().to(device)
      att_masks = None
  
      # forward the model to also get generated samples for each image
      # Only leave one feature for each image, in case duplicate sample
      tmp_eval_kwargs = eval_kwargs.copy()
      tmp_eval_kwargs.update({'sample_n': 1})
      seq, seq_logprobs = model(
          fc_feats, att_feats, att_masks, opt=tmp_eval_kwargs, mode='sample')
      seq = seq.data
  
      sents = utils.decode_sequence(model.vocab, seq)

      return sents

demo = gr.Blocks()

with demo:
    gr.Markdown(
    """
    # Demo for CLIP-Caption-Reward
    """)
    inp = gr.Image(type="pil")
    out = gr.Textbox()

    image_button = gr.Button("Run")
    image_button.click(fn=inference,
               inputs=inp, 
               outputs=out)
    

demo.launch()