from sentence_transformers import SentenceTransformer, util
from scipy.stats import spearmanr, pearsonr
from sklearn import metrics
from sklearn.decomposition import PCA
from sklearn.cluster import AgglomerativeClustering
import umap
import matplotlib.pyplot as plt
import numpy as np

import time
from datetime import datetime
import os
import spacy


class SBERT_Model:
    def __init__(self, name, filepath, dataset, probing=False):
        """
        param: dataset must be pandas dataframe
        """
        self.name = name
        self.filepath = filepath
        self.dataset = dataset
        self.model = SentenceTransformer(filepath)
        self.sentences1 = dataset['sentence1'].tolist()
        self.sentences2 = dataset['sentence2'].tolist()
        self.labels = dataset['label'].tolist()
        self.embeddings1 = self.get_embeddings(self.sentences1)
        self.embeddings2 = self.get_embeddings(self.sentences2)
        self.cosine_scores = self.get_cosine_scores(self.embeddings1, self.embeddings2)
        self.preds = self.get_preds()

    def visualize_embeddings(self, sentences, highlights=[], figsize=(80, 200)):
        """
        Parameters
        ----------
        sentences: list of strings
            list of sentences that will be visualized and colored in blue
        highlights: list of strings or a string
            list of sentences that will be visualized and colored in red
        figsize: tuple
            size of the final image. The default is set to (80,200),
            which is 8000x20000 pixels and is appropriate for
            the stsb test dataset (2757 data points).
        """
        plt.figure(figsize=figsize)
        t1 = time.time()
        #Convert into list if it is a string
        if isinstance(highlights, str):
            highlights = [highlights]
        if len(highlights) == 1:
            highlights_embs = self.get_embeddings(highlights).reshape(1, -1).cpu()
        elif len(highlights) > 1:
            highlights_embs = self.get_embeddings(highlights).cpu()
        embs = self.get_embeddings(sentences).cpu()
        umap_model = umap.UMAP(n_neighbors=15, n_components=2, min_dist=2, spread=20, metric='cosine', random_state=42).fit(embs)
        umap_data_transformed = umap_model.transform(embs)
        folder = "outputs"
        plot_name = f"plot-{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.png"
        plot_path = os.path.join(folder,plot_name)
        legend_name = f'legend-{plot_name[:-4]}.txt'
        legend_path = os.path.join(folder,legend_name)
        with open(legend_path, 'w') as reader:
            for i, (x, y) in enumerate(umap_data_transformed):
                name = sentences[i]
                reader.write(f"{i}: {name} ({x:.2f},{y:.2f})\n")
                plt.plot(x, y, 'bo', alpha=0.3)
            if highlights:
                highlights_transformed = umap_model.transform(highlights_embs)
                for i, (x, y) in enumerate(highlights_transformed):
                    name = highlights[i]
                    label = f"{i}: {name}"
                    reader.write(f"Highlights: {i}: {name} ({x:.2f},{y:.2f})\n")
                    plt.plot(x, y, 'ro', markersize=20)
                    plt.text(x, y, label , backgroundcolor="white", alpha=0.4)
        plt.savefig(plot_path)
        t2 = time.time()
        print(f"Plot saved to {plot_path}. Time it took to generate: {t2-t1:.2f} seconds")
        print(f'{legend_path} stores the information about the data points on the plot.')

    def visualize_preds(self):
        print(f"Average cosine similarity: {np.average(self.preds)}; Standard deviation: {np.std(self.preds)}")
        fig, ax = plt.subplots(1, 1, figsize=(8, 5), tight_layout=True)
        ax.hist(self.preds, bins=20)
        plt.xlim([0, 1])
        ax.set_title(
            f"{self.name} - histogram " + r"$\mu$" + f"={round(np.average(self.preds), 3)}, " + r"$\sigma$" f"={round(np.std(self.preds), 3)}")

        plt.ylabel('Frequency')
        plt.xlabel('Cosine similarity between sentence pair')
        plt.savefig(f'{self.name}_histogram.png', bbox_inches='tight')

    def get_word_importance(self,sentence_pair, mask="[MASK]"):
        print(f"Getting word importance of {sentence_pair}")
        nlp = spacy.load("en_core_web_sm")
        sentence1 = nlp(sentence_pair[0])
        sentence2 = nlp(sentence_pair[1])
        unmasked_embs = self.get_embeddings([sentence_pair[0], sentence_pair[1]])
        unmasked_similarity = self.get_cosine_scores(unmasked_embs[0], unmasked_embs[1])
        masked_pairs = []
        for i in range(len(sentence1)):
            masked_sentence = []
            for y in range(len(sentence1)):
                if y == i:
                    masked_sentence.append(mask)
                    masked_word = sentence1[y]
                else:
                    masked_sentence.append(sentence1[y].text)
            masked_sentence = " ".join(masked_sentence)
            masked_pair = {"masked word": masked_word, "sentence1": masked_sentence, "sentence2": sentence_pair[1], "masked sentence": 0}
            masked_pairs.append(masked_pair)
        for i in range(len(sentence2)):
            masked_sentence = []
            for y in range(len(sentence2)):
                if y == i:
                    masked_sentence.append(mask)
                    masked_word = sentence2[y]
                else:
                    masked_sentence.append(sentence2[y].text)
            masked_sentence = " ".join(masked_sentence)
            masked_pair = {"masked word": masked_word, "sentence1": sentence_pair[0], "sentence2": masked_sentence, "masked sentence": 1}
            masked_pairs.append(masked_pair)
        scores = []
        for masked_pair in masked_pairs:
            embs = self.get_embeddings([masked_pair["sentence1"], masked_pair["sentence2"]])
            scores.append((masked_pair, float(abs(unmasked_similarity - self.get_cosine_scores(embs[0], embs[1])))))
        token_scores = []
        for i, score in enumerate(scores):
            masked_pair = score[0]
            if masked_pair["masked sentence"]:
                doc2 = 0
            else:
                doc2 = 1
            token_score = {"doc":[sentence1, sentence2][masked_pair["masked sentence"]], "doc2":[sentence1, sentence2][doc2], "word": masked_pair["masked word"].text, "i": masked_pair["masked word"].i,  "score": score[1], "sentence_pair": sentence_pair}
            token_scores.append(token_score)
        
        html_output = []
        html_output.append('<div class="sentence">\n')
        for i, score in enumerate(token_scores):
            s1_words = [s.text for s in score["doc"]]
            s2_words = [s.text for s in score["doc2"]]
            token = score["doc"][score["i"]]
            cos_score = score["score"]
            if s1_words.count(token.text) and s2_words.count(token.text):
              both_sentences = True
            else:
              both_sentences = False
            html_output.append(f'<div id="{randrange(0,10000)}_{token.text}_{token.lemma_}_{both_sentences}_{token.pos_}_{token.tag_}_{token.dep_}_{cos_score:.2f}" style="background-color: rgba(0, 160, 252, {cos_score:.2f});">{token.text}<br>{cos_score:.2f}</div>\n')
            if i == len(sentence1)-1:
                html_output.append('</div>\n<div class="sentence">\n')
        html_output.append('</div>\n')
        return html_output

    def get_html_word_importance(self, sentence_pairs, mask="[MASK]"):
        top = """<!doctype html>
<html lang="en">
<head>
  <meta charset="utf-8">
  <title>SBERT Word Importance</title>
  <meta name="description" content="Visualizing SBERT sentence pairs importance of words">
<style>
.sentence{
    display:flex;
}
.sentence > div{
    border: 1px solid gray;
    
    min-width: 100px;
    min-height: 20px;
    padding: 20px;
    align-items: center;
    text-align: center;
}

</style>
</head>

<body>
    <h1>SBERT Word Importance Visualization</h1>
"""

        bottom = "</body>\n</html>"

        folder = "outputs"
        html_name = f"heatmap-{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.html"
        html_path = os.path.join(folder,html_name)

        with open(html_path, 'w') as reader:
            reader.write(top)
            reader.write(f'<h2>{self.name}</h2>\n')
            reader.write(f'<h3>Token used for masking:{mask}</h3>\n')
            for i, sentence_pair in enumerate(sentence_pairs):
                print(i)
                reader.write(f'<div id="{i}" class="sentence_pair">{i+1}.\n')
                sentence_divs = self.get_word_importance(sentence_pair, mask=mask)
                for sentence_div in sentence_divs:
                    reader.write(sentence_div)
                reader.write('</div>\n')
            reader.write(bottom)
            
    def get_embeddings(self, sentences):
        return self.model.encode(sentences, convert_to_tensor=True)

    def get_preds(self):
        """This method extracts the scores for similarity between sentence pairs
        Cosine_scores have scores for similarity between all of the sentences, 
        but we only need similarity between each sentence pair.
        """
        preds = []
        for i in range(len(self.cosine_scores[0])):
            preds.append(float(self.cosine_scores[i][i]))
        return preds

    def get_cosine_scores(self, emb1, emb2):
        return util.pytorch_cos_sim(emb1, emb2)

    def get_pearson(self):
        return pearsonr(self.labels, self.preds)[0]

    def get_spearman(self):
        return spearmanr(self.labels, self.preds)[0]

    def get_MSE(self):
        return metrics.mean_squared_error(self.labels, self.preds)
    
    def print_statistics(self):
        print(f"{self.name}: MSE:{self.get_MSE()}; Pearson:{self.get_pearson()}; Spearman:{self.get_spearman()}")
