from transformers import BertModel, BertTokenizer, Trainer, TrainingArguments
import torch
import torch.nn as nn
from datasets import load_dataset
from torch.utils.data import Dataset
from torch.nn.functional import cosine_similarity
import json
import numpy as np
import sys
sys.path.append("./")
from tqdm import tqdm
from embedder.train_contrastive_retriever import TripletNetwork

embedder_dir = 'embedder/embedder_margin_2/checkpoint-300'


model = TripletNetwork().to("cuda")
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

# Load the weights
model.load_state_dict(torch.load(embedder_dir + "/pytorch_model.bin"))
model.eval()  # for inference


dataset_path = "data/finetune/contrastive_preprocess_test.json"

ds = load_dataset("json", data_files=dataset_path)
dataset = ds['train']

success = 0
for idx, item in tqdm(enumerate(dataset), total=len(dataset)):
    # if idx == 100:
    #     break
    tokenized_input = tokenizer(item["query_prompt"], padding='max_length', truncation=True, max_length=512, return_tensors="pt")
    input_ids = tokenized_input["input_ids"].to("cuda")
    attention_mask = tokenized_input["attention_mask"].to("cuda")

    positive_input = tokenizer(item["positive_sample"], padding='max_length', truncation=True, max_length=512, return_tensors="pt")
    positive_input_ids = positive_input["input_ids"].to("cuda")
    positive_attention_mask = positive_input["attention_mask"].to("cuda")

    negative_input = tokenizer(item["negative_sample"], padding='max_length', truncation=True, max_length=512, return_tensors="pt")
    negative_input_ids = negative_input["input_ids"].to("cuda")
    negative_attention_mask = negative_input["attention_mask"].to("cuda")

    with torch.no_grad():
        # print("input_ids", input_ids)
        # print("attention_mask", attention_mask)
        query_embedding = model(input_ids, attention_mask)
        positive_embedding = model(positive_input_ids, positive_attention_mask)
        negative_embedding = model(negative_input_ids, negative_attention_mask)

        # calculate similarity
        positive_similarity = cosine_similarity(query_embedding, positive_embedding)
        negative_similarity = cosine_similarity(query_embedding, negative_embedding)

        # print("positive_similarity: ", positive_similarity)
        # print("negative_similarity: ", negative_similarity)

        if positive_similarity > negative_similarity:
            success += 1

print("Success rate: ", success / len(dataset))
        # input()
