import torch
from transformers import BertTokenizerFast, BertModel, BertConfig, BertTokenizer
import numpy as np
import pandas as pd
from tqdm.notebook import tqdm as tqdm
import glob
import os
from sklearn.decomposition import PCA
import time
import plotly.express as px
from sklearn.manifold import TSNE
from scipy.spatial.distance import cdist, cosine
from gpytorch.kernels.rq_kernel import RQKernel
import torch.nn as nn
from torch.nn import functional as F
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader

class BertClassifier(torch.nn.Module):
    
    def __init__(self, config, model, dim=256, num_classes=2):
        super(BertClassifier, self).__init__()
        
        # create the model config and BERT initialize the pretrained BERT, also layers wise outputs
        self.config = config
        self.base = model
        
        # classifier head [not useful]
        self.head = torch.nn.Sequential(*[
            torch.nn.Dropout(p=self.config.hidden_dropout_prob),
            torch.nn.Linear(in_features=self.config.hidden_size, out_features=dim),
            torch.nn.ReLU(),
            torch.nn.Dropout(p=self.config.hidden_dropout_prob),
            torch.nn.Linear(in_features=dim, out_features=num_classes)
        ])
    
    def forward(self, input_ids, attention_mask=None):
        
        # first output is top layer output, second output is context of input seq and third output will be layerwise tokens 
        top_layer, pooled, layers = self.base(input_ids, attention_mask)
        outputs = self.head(pooled)
        return top_layer, outputs, layers

class SentimentDataset(Dataset):
    def __init__(self, df, tokenizer, max_len=512):
        self.tokenizer = tokenizer
        self.text = df.review_text.values
        self.max_len =  max_len
    
    def __len__(self):
        return len(self.text)
    
    def __getitem__(self, idx):
        text  = self.text[idx]
        
        # encode the text and target into tensors return the attention masks as well
        encoding = self.tokenizer.encode_plus(
            text=text,
            add_special_tokens=True,
            max_length=self.max_len,
            return_token_type_ids=False,
            pad_to_max_length=True,
            return_attention_mask=True,
            return_tensors='pt',
            truncation=True
        )
        
        return {
          'text': text,
          'input_ids': encoding['input_ids'].flatten(),
          'attention_mask': encoding['attention_mask'].flatten(),
        }
        

def model_predict(trained):
        dictionary_list = []

        df = pd.read_csv("./amazon-review/dvd-UL.csv")
        df = df.sample(n=1000, random_state=42) #number of samples
        df = df.reset_index(drop=True)

        dataset = SentimentDataset(df=df, tokenizer=tokenizer)

        data_loader = torch.utils.data.DataLoader(
            dataset=dataset,
            batch_size= 4,
            shuffle=False,
            num_workers=8
        )

        for bi, d in enumerate(tqdm(data_loader)):
            input_ids = d["input_ids"]
            attention_mask = d["attention_mask"]

            _, _, output = classifier_trained(input_ids, attention_mask)

            output = output[1:]

            for zeta in range(len(output[0])):
                for i in range(0,12):
                    new_row = {'embeddings':output[i][zeta][0].cpu().detach().numpy(), 'layers': i+1}
                    dictionary_list.append(new_row)

            dictionary_list = np.save(f"./data/batch_{bi}", dictionary_list, allow_pickle=True)
            dictionary_list = []

if(os.path.exists("./data")):
    files = glob.glob('./data/*')
    for f in files:
        os.remove(f)
else:
    os.makedirs("./data")

dictionary_list = []

files = glob.glob("./data/*.npy")

for j in range(len(files)):
    alpha = np.load(f"./data/batch_{j}.npy", allow_pickle = True)
    for i in range(len(alpha)):
        new_row = {'embeddings':alpha[i]["embeddings"], 'layers': alpha[i]["layers"]}
        dictionary_list.append(new_row)

df.to_pickle("./dvd-ft-b.pkl") #naming convention dvd-ft-b (dvd dataset on books finetuned)

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', use_lower_case=True)
config = BertConfig.from_pretrained("bert-base-uncased", output_hidden_states=True)
model = BertModel.from_pretrained('bert-base-uncased', config=config)

if(os.path.exists("./data")):
    files = glob.glob('./data/*')
    for f in files:
        os.remove(f)
else:
    os.makedirs("./data")

files = "./amazon-review/dvd-UL.csv"
dataset = pd.read_csv(files)
dataset = dataset.sample(n=1000, random_state=42)
dataset = dataset.reset_index(drop=True)

dictionary_list = []

files = glob.glob("./data/*.npy")

for j in range(len(files)):
    alpha = np.load(f"./data/batch_{j}.npy", allow_pickle = True)
    for i in range(len(alpha)):
        new_row = {'embeddings':alpha[i]["embeddings"], 'layers': alpha[i]["layers"]}
        dictionary_list.append(new_row)

df = pd.DataFrame.from_dict(dictionary_list)
df.to_pickle("./dvd-pt.pkl") ##naming convention dvd-pt (dvd dataset on pretrain model)

df = pd.read_pickle("./books-ft-b.pkl")

df = pd.read_pickle("./books-pt.pkl")

fig = px.scatter_3d(
    components, x=0, y=1, z=2, color=df['layers'],
    title=f'Domain representation plotting PCA',
    labels={'color': 'layers'}
)
fig.update_layout(
    autosize=False,
    width=1080,
    height=720,)

fig.show()

fig = px.scatter_3d(
    components, x=0, y=1, z=2, color=df['layers'],
    title=f'Domain representation plotting PCA',
    labels={'color': 'layers'}
)
fig.update_layout(
    autosize=False,
    width=1080,
    height=720,)

fig.show()

df = pd.read_pickle("./dvd-pt.pkl")
df1 = pd.read_pickle("./dvd-ft-b.pkl")

def mmd(first_tensor: torch.Tensor, second_tensor: torch.tensor, kernel) -> float:
    """ MMD will be calculated between two sets of tensors
    Parameters
    ----------
    first_tensor : torch.Tensor
        Tensor of shape batch_size, m
        m is the dimension of the first tensor
    second_tensor: torch.Tensor
        Tensor of shape batch_size, n
        n is the dimension of the target tensor 
    Returns
    -------
    float
        MMD between two samples
    """
    first_tensor_correlation = kernel(first_tensor, first_tensor).evaluate()
    second_tensor_correlation = kernel(second_tensor, second_tensor).evaluate()
    first_second_tensor_correlation = kernel(first_tensor, second_tensor).evaluate()
    m = first_tensor.size(0)
    n = second_tensor.size(0)
    sum_first_corr = first_tensor_correlation.sum().item()
    sum_second_corr = second_tensor_correlation.sum().item()
    sum_first_second_corr = first_second_tensor_correlation.sum().item()
    first_term = (1 / (m ** 2)) * sum_first_corr
    second_term = (1 / (n ** 2)) * sum_second_corr
    third_term = (2 / (m * n)) * sum_first_second_corr
    divergence = first_term + second_term - third_term
    return divergence

def closest_rows(a):
    # Get euclidean distances as 2D array
    dists = cdist(a, a, 'cosine')

    # Fill diagonals with something greater than all elements as we intend
    # to get argmin indices later on and then index into input array with those
    # indices to get the closest rows
#     dists.ravel()[::dists.shape[1]+1] = dists.max()+1
    return dists

pretrained_books_dvd_RSA = {}
pretrained_books_dvd_MMD = {} #for storing MMD between 2 models for each layer

MMD between bert representation and fine tuned bert (books) - DVD dataset

{1: 0.025523517691246278,
 2: 0.017031062632806515,
 3: 0.034405121243115344,
 4: 0.0743630493421803,
 5: 0.0658974319876604,
 6: 0.21873141347150904,
 7: 0.23004184150353701,
 8: 0.16440297895429068,
 9: 0.11399606384908612,
 10: 0.2680937669944564,
 11: 0.31493303470159223,
 12: 0.397320439848271}

MMD between bert representation and fine tuned bert (books) - books dataset

{1: 0.008592661067214902,
 2: 0.008466682016583071,
 3: 0.01607070255527243,
 4: 0.03613209326039435,
 5: 0.059290986788250954,
 6: 0.0553308317517649,
 7: 0.0686747228871778,
 8: 0.060072450953288725,
 9: 0.09658271662444762,
 10: 0.2985128726786279,
 11: 0.3239971706890749,
 12: 0.4579824415026016}