from sentence_transformers import SentenceTransformer, util
import numpy as np
import csv 
import seaborn
import matplotlib
import matplotlib.pyplot as plt

embedder = SentenceTransformer('roberta-large-nli-stsb-mean-tokens') 


with open("result.csv", 'w') as csvfile:
      csvwriter = csv.writer(csvfile)

ifcn_s = []
bot_s = []

with open("test.csv", 'r') as csvfile: 
    csvreader = csv.reader(csvfile) 
    
    for row in csvreader:
      corpus = [row[0]]
      ifcn = [row[1]]
      bot = [row[2]]

      corpus_embeddings = embedder.encode(corpus, convert_to_tensor=True)
      ifcn_embedding = embedder.encode(ifcn, convert_to_tensor=True)
      bot_embedding = embedder.encode(bot, convert_to_tensor=True)

      cos_scores = util.pytorch_cos_sim(ifcn_embedding, corpus_embeddings)[0]
      cos_scores = cos_scores.cpu()
      top_results = np.argpartition(-cos_scores, range(1))[0:1]
      ifcn_score = cos_scores[0]
      ifcn_s.append(float(ifcn_score))

      cos_scores = util.pytorch_cos_sim(bot_embedding, corpus_embeddings)[0]
      cos_scores = cos_scores.cpu()
      top_results = np.argpartition(-cos_scores, range(1))[0:1]
      bot_score = cos_scores[0]
      bot_s.append(float(bot_score))

      temp = [row[0],row[1],row[2],float(ifcn_score),float(bot_score)]



      print(temp)
      with open("result.csv", 'a+') as csvfile:
        csvwriter = csv.writer(csvfile)
        csvwriter.writerow(temp)