import torch
import pandas as pd
import numpy as np
from tqdm import tqdm
from config import CFG
import itertools
from clip import CLIPModel
from utils import AvgMeter, get_lr
from data import *
from transformers import DistilBertTokenizer
import torch.nn.functional as F
from utils import *
import pickle


# Reference for loading dataset and model architecture.
# https://towardsdatascience.com/simple-implementation-of-openai-clip-model-a-tutorial-ace6ff01d9f2
#


def train_epoch(model, train_data_loader, optimizer):
    """
    Method to start a train epoch.
    """
    loss_sum = 0
    loss_count = 0
    for batch in tqdm(train_data_loader, total=len(train_data_loader)):
        batch_data = {k: v.to(CFG.device) for k, v in batch.items() if k != "caption"}
        loss = model(batch_data)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        batch_size = batch_data["image"].size(0)
        loss_count += max(1, batch_size)
        loss_sum += loss.item() * max(1,batch_size)
    return loss_sum/loss_count


def valid_epoch(model, valid_data_loader):
    """
    Method to start a validation epoch.
    """
    loss_sum = 0
    loss_count = 0
    #tqdm_object = tqdm(valid_loader, total=len(valid_loader))
    for batch in tqdm(valid_data_loader, total=len(valid_data_loader)):
        batch_data = {k: v.to(CFG.device) for k, v in batch.items() if k != "caption"}
        loss = model(batch)
        batch_size = batch_data["image"].size(0)
        loss_count += max(1, batch_size)
        loss_sum += loss.item() * max(1,batch_size)
        #tqdm_object.set_postfix(valid_loss=loss_sum/loss_count)
    return loss_sum/loss_count

def store_image_embeddings(valid_df, model):
    """
    Generate and store embeddings of validation
    """
    tokenizer = DistilBertTokenizer.from_pretrained(CFG.text_tokenizer)
    valid_loader = get_dataset_loader(valid_df, tokenizer, mode="valid")
    
    #model = CLIPModel().to(CFG.device)
    #model.load_state_dict(torch.load(model_path, map_location=CFG.device))
    model.eval()
    
    valid_image_embeddings = []
    with torch.no_grad():
        for batch in tqdm(valid_loader):
            image_features = model.image_encoder(batch["image"].to(CFG.device))
            image_embeddings = model.image_projection(image_features)
            valid_image_embeddings.append(image_embeddings)
    
    torch.save(torch.cat(image_embeddings), CFG.img_emb_path)
    return

def main():

    # Generate train and valid dataframe
    train_data, valid_data = gen_train_valid_dfs()
    
    # Use pretrained tokeniser to tokensize the captions
    tokenizer = DistilBertTokenizer.from_pretrained(CFG.text_tokenizer)

    # Create data loaders for both train and validation data
    train_loader = get_dataset_loader(train_data, tokenizer, mode="train")
    valid_loader = get_dataset_loader(valid_data, tokenizer, mode="valid")

    model = CLIPModel().to(CFG.device)
    params = [
        {"params": model.image_encoder.parameters(), "lr": CFG.image_encoder_lr},
        {"params": model.text_encoder.parameters(), "lr": CFG.text_encoder_lr},
        {"params": itertools.chain( model.image_projection.parameters(), model.text_projection.parameters()
        ), "lr": CFG.head_lr, "weight_decay": CFG.weight_decay}
    ]
    optimizer = torch.optim.AdamW(params, weight_decay=0.)
    lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode="min", patience=CFG.patience, factor=CFG.factor
    )
    min_loss = float('inf')
    for epoch in range(CFG.epochs):
        print(f"Epoch: {epoch + 1}")
        model.train()
        train_loss = train_epoch(model, train_loader, optimizer)
        model.eval()
        with torch.no_grad():
            valid_loss = valid_epoch(model, valid_loader)

        if valid_loss < min_loss:
            min_loss = valid_loss
            torch.save(model.state_dict(), CFG.model_path)
            print("Saved Best Model!")
    
        lr_scheduler.step(valid_loss)

main()
