import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
import streamlit as st

import torchvision
from torchvision import transforms
import cv2
import math

from collections import Counter
from PIL import Image
import PIL

import zipfile
import io

import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.utils.data import Dataset, DataLoader
import torch.nn.functional as fun

from tqdm.notebook import tqdm
import matplotlib.pyplot as plt

# import nltk
# import ssl

# try:
#     _create_unverified_https_context = ssl._create_unverified_context
# except AttributeError:
#     pass
# else:
#     ssl._create_default_https_context = _create_unverified_https_context

# nltk.download()
from nltk.tokenize import word_tokenize

import match
import pickle
import gc
import random

data = pd.read_csv("captions.txt", sep=',')

#Removes Single Char
def remove_single_char(caption_list):
    list = []
    for word in caption_list:
        if len(word)>1:
            list.append(word)
    return list

#Make an array of words out of caption and then remove useless single char words

data['caption'] = data['caption'].apply(lambda caption :word_tokenize(caption))

data['caption'] = data['caption'].apply(lambda word : remove_single_char(word))

#We need to make sure size of all the captions arrays is same so we add <cell> to cover up
lengths = []
lengths = data['caption'].apply(lambda caption : len(caption))

max_length = lengths.max()

data['caption'] = data['caption'].apply(lambda caption : ['<start>'] + caption + ['<cell>']*(max_length-len(caption)) + ['<end>'])

#For non truncated dataframe to appear
pd.set_option('display.max_colwidth', None)

#Extracting words 
words = data['caption'].apply(lambda word : " ".join(word)).str.cat(sep = ' ').split(' ')

#Arranging the words in order of their frequency
word_dict = sorted(Counter(words), key=Counter(words).get, reverse=True)

dict_size = len(word_dict)
vocab_threshold = 5

#Encoding the words with index in dictionary made above
data['sequence'] = data['caption'].apply(lambda caption : [word_dict.index(word) for word in caption])
data = data.sort_values(by = 'image')

class PositionalEncoding(nn.Module):

    def __init__(self, d_model, dropout=0.1, max_len=(max_length+2)):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)
        

    def forward(self, x):
        if self.pe.size(0) < x.size(0):
            self.pe = self.pe.repeat(x.size(0), 1, 1)
        self.pe = self.pe[:x.size(0), : , : ]
        
        x = x + self.pe
        return self.dropout(x)
    
class ImageCaptionModel(nn.Module):
    def __init__(self, n_head, n_decoder_layer, vocab_size, embedding_size):
        super(ImageCaptionModel, self).__init__()
        self.pos_encoder = PositionalEncoding(embedding_size, 0.1)
        self.TransformerDecoderLayer = nn.TransformerDecoderLayer(d_model =  embedding_size, nhead = n_head)
        self.TransformerDecoder = nn.TransformerDecoder(decoder_layer = self.TransformerDecoderLayer, num_layers = n_decoder_layer)
        self.embedding_size = embedding_size
        self.embedding = nn.Embedding(vocab_size , embedding_size)
        self.last_linear_layer = nn.Linear(embedding_size, vocab_size)
        self.init_weights()
        self.n_head = n_head

    def init_weights(self):
        initrange = 0.1
        self.embedding.weight.data.uniform_(-initrange, initrange)
        self.last_linear_layer.bias.data.zero_()
        self.last_linear_layer.weight.data.uniform_(-initrange, initrange)

    def generate_Mask(self, size, decoder_inp):
        decoder_input_mask = (torch.triu(torch.ones(size, size)) == 1).transpose(0, 1)
        decoder_input_mask = decoder_input_mask.float().masked_fill(decoder_input_mask == 0, float('-inf')).masked_fill(decoder_input_mask == 1, float(0.0))

        decoder_input_pad_mask = decoder_inp.float().masked_fill(decoder_inp == 0, float(0.0)).masked_fill(decoder_inp > 0, float(1.0))
        decoder_input_pad_mask_bool = decoder_inp == 0

        return decoder_input_mask, decoder_input_pad_mask, decoder_input_pad_mask_bool

    def forward(self, encoded_image, decoder_inp):
        # display(decoder_inp)
        encoded_image = encoded_image.permute(1,0,2)

        decoder_inp = torch.clamp(decoder_inp, 0, self.embedding.num_embeddings - 1)


        decoder_inp_embed = self.embedding(decoder_inp)* math.sqrt(self.embedding_size)
        decoder_inp_embed = self.embedding(decoder_inp)
        
        decoder_inp_embed = self.pos_encoder(decoder_inp_embed)
        decoder_inp_embed = decoder_inp_embed.permute(1,0,2)
        
        
        decoder_input_mask, decoder_input_pad_mask, decoder_input_pad_mask_bool = self.generate_Mask(decoder_inp.size(1), decoder_inp)
        decoder_input_mask = decoder_input_mask
        decoder_input_pad_mask = decoder_input_pad_mask
        decoder_input_pad_mask_bool = decoder_input_pad_mask_bool
        
        decoder_output = self.TransformerDecoder(tgt = decoder_inp_embed, memory = encoded_image, tgt_mask = decoder_input_mask, tgt_key_padding_mask = decoder_input_pad_mask_bool)
        
        final_output = self.last_linear_layer(decoder_output)

        return final_output,  decoder_input_pad_mask
    

model = pd.read_pickle('ImageCaptioning_Model.pkl')
model.eval()
start_token = 2
end_token = 3
cell_token = 1
max_seq_len = 34

validation = pd.read_pickle('Image_Features_Embed_ResNet_Valid.pkl')

def process_image_from_zip(zip_path, image_name):
    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        with zip_ref.open(image_name) as file:
            # Use BytesIO to keep the file-like object open
            image_data = io.BytesIO(file.read())
            image = Image.open(image_data)
            image = image.convert("RGB")  # Convert to RGB if needed
            return image

def generate_caption(K, image_path): 

    model.eval() 
    image__path = 'Images/' + image_path
    # image = Image.open('Images/' + image_path).convert("RGB")
    image = process_image_from_zip('Images.zip', image__path)
    plt.imshow(image)

    valid_img_df = validation[validation['image']==image_path]
    print("Actual Caption : ")
    actual_caption_list = valid_img_df['caption'].tolist()
    filtered_caption_list = [word for word in actual_caption_list[0] if word not in ['<start>', '<end>', '<cell>']]
    actual_caption = " ".join(filtered_caption_list)
    print(actual_caption)

    valid_img_embed = validation[validation['image'] == image_path]
        
    img_embed = valid_img_embed['embedded'].tolist()

    img_embed = torch.tensor(img_embed)


    input_seq = [cell_token]*max_seq_len
    input_seq[0] = start_token

    input_seq = torch.tensor(input_seq).unsqueeze(0)
    predicted_sentence = []
    
    with torch.no_grad():
        for eval_iter in range(0, max_seq_len):
            img_embed_dense = img_embed.to_dense()

            output, padding_mask = model.forward(img_embed, input_seq)

            output = output[eval_iter, 0, :]

            values = torch.topk(output, K).values.tolist()
            indices = torch.topk(output, K).indices.tolist()

            next_word_index = random.choices(indices, values, k = 1)[0]

            index_to_word = {index: word for index, word in enumerate(word_dict)}
            next_word = index_to_word[next_word_index]

            if eval_iter + 1 < max_seq_len:
                input_seq[:, eval_iter + 1] = next_word_index

            if next_word == '<end>' :
                break

            predicted_sentence.append(next_word)
    
    print("\n")
    print("Predicted caption : ")

    filtered_caption_list = [word for word in predicted_sentence if word not in ['<start>', '<end>', '<cell>']]
    print(" ".join(filtered_caption_list))


st.title('Image Captioning')
st.write('Generate Caption for Random Image')
generate_caption_button = st.button('Generate Caption')

if generate_caption_button: 
    try:
        random_row = validation.sample()
        random_image = random_row.iloc[0]['image']
        
        generate_caption(1, random_image)
            
    except RuntimeError as e:
        print("TRY AGAIN")