import pickle

import os

import numpy as np

from tqdm import tqdm

from sklearn.cluster import spectral_clustering, KMeans

import torch

from transformers import AutoTokenizer, AutoModel, AutoConfig

import argparse

device = torch.device("cuda:0")

batch_size = 64

MAX_CLUSTER_COUNT = 50

mode = 'ratio'

model = AutoModel.from_pretrained('GanjinZero/coder_eng_pp')

tokenizer = AutoTokenizer.from_pretrained('GanjinZero/coder_eng_pp')

def load_pickle(file_path):

    with open(file_path, "rb") as f:

        df = pickle.load(f)

    return df

def get_bert_embed(phrase_list, normalize=True, summary_method="MEAN", tqdm_bar=False):

    global model

    input_ids = []

    for phrase in phrase_list:

        input_ids.append(tokenizer.encode_plus(

            phrase, max_length=32, add_special_tokens=True,

            truncation=True, padding='max_length')['input_ids'])

        # print(len(input_ids))

    model.eval()

    model = model.to(device)

    count = len(input_ids)

    now_count = 0

    output_list = []

    with torch.no_grad():

        if tqdm_bar:

            pbar = tqdm(total=count)

        while now_count < count:

            input_gpu_0 = torch.LongTensor(input_ids[now_count:min(

                now_count + batch_size, count)]).to(device)

            if summary_method == "CLS":

                embed = model(input_gpu_0)[1]

            if summary_method == "MEAN":

                embed = torch.mean(model(input_gpu_0)[0], dim=1)

            if normalize:

                embed_norm = torch.norm(

                    embed, p=2, dim=1, keepdim=True).clamp(min=1e-12)

                embed = embed / embed_norm

            if now_count % 1000000 == 0:

                if now_count != 0:

                    output_list.append(output.cpu().numpy())

                    del output

                    torch.cuda.empty_cache()

                output = embed

            else:

                output = torch.cat((output, embed), dim=0)

            if tqdm_bar:

                pbar.update(min(now_count + batch_size, count) - now_count)

            now_count = min(now_count + batch_size, count)

            del input_gpu_0

            torch.cuda.empty_cache()

        if tqdm_bar:

            pbar.close()

    output_list.append(output.cpu().numpy())

    del output

    torch.cuda.empty_cache()

    return np.mean(np.concatenate(output_list, axis=0), axis=0)

def re_cluster(terms_list, mode, similarity, threshold):

    ready = [terms_list]

    res = []

    while ready:

        now = ready.pop()

        clu0, clu1 = cut(now, mode, similarity)

        membed_0 = get_bert_embed(clu0)

        membed_1 = get_bert_embed(clu1)

        if np.dot(membed_0, membed_1) > threshold or len(clu0) <= 1 or len(clu1) <= 1:

            res.append(clu0)

            res.append(clu1)

        else:

            # ready.append(clu0)

            # ready.append(clu1)

            if len(clu0) <= MAX_CLUSTER_COUNT:

                res.append(clu0)

            else:

                ready.append(clu0)

            if len(clu1) <= MAX_CLUSTER_COUNT:

                res.append(clu1)

            else:

                ready.append(clu1)

        # for clu in [clu0, clu1]:

        #     if len(clu) <= MAX_CLUSTER_COUNT:

        #         res.append(clu)

        #     else:

        #         ready.append(clu)

    return res

def cut(terms_list, mode, similarity):

    if mode == 'ratio':

        clu0, clu1 = ratio_cut(terms_list, similarity)

    else:

        clu0, clu1 = normalize_cut(terms_list, similarity)

    return clu0, clu1

def get_sim(terms_list, similarity):

    idx = [phrase2id[x] for x in terms_list]

    sim = np.zeros(shape=(len(idx), len(idx)))

    cnt = len(idx)

    for i in range(cnt):

        for j in range(cnt):

            if idx[j] in indices[idx[i]]:

                sim[i][j] = similarity[idx[i]][np.argwhere(indices[idx[i]]==idx[j])]

            elif idx[i] in indices[idx[j]]:

                sim[i][j] = similarity[idx[j]][np.argwhere(indices[idx[j]]==idx[i])]

    return sim

def laplacian(matrix, normalize=False):

    d_val = matrix.sum(axis=0)

    d = np.diag(d_val)

    l = d - matrix

    if normalize:

        d_inverse_root_val = d_val ** (-1/2)

        d_inverse_root = np.diag(d_inverse_root_val)

        l = np.matmul(np.matmul(d_inverse_root, l), d_inverse_root)

    return l

def ratio_cut(terms_list, similarity):

    sim = get_sim(terms_list, similarity)

    l = laplacian(sim)

    u, v = np.linalg.eig(l)

    index = np.argsort(u.real)

    feat = v[:,index[0:2]].real

    feat_norm = np.linalg.norm(feat, ord=2, axis=1, keepdims=True)

    feat = feat / feat_norm

    cluster = KMeans(n_clusters=2).fit_predict(feat)

    clu0 = np.array(terms_list)[cluster==0].tolist()

    clu1 = np.array(terms_list)[cluster==1].tolist()

    return clu0, clu1

def normalize_cut(terms_list, similarity):

    sim = get_sim(terms_list, similarity)

    l = laplacian(sim, True)

    u, v = np.linalg.eig(l)

    index = np.argsort(u.real)

    feat = v[:,index[0:2]].real

    feat_norm = np.linalg.norm(feat, ord=2, axis=1, keepdims=True)

    feat = feat / feat_norm

    cluster = KMeans(n_clusters=2).fit_predict(feat)

    clu0 = np.array(terms_list)[cluster==0].tolist()

    clu1 = np.array(terms_list)[cluster==1].tolist()

    return clu0, clu1

def print_cluster_to_file(f, one_cluster_result):

    for idx, term in enumerate(one_cluster_result):

        f.write(term)

        if idx != len(one_cluster_result) - 1:

            f.write('|')

        else:

            f.write('\n')

if __name__ == '__main__':

    parser = argparse.ArgumentParser()

    parser.add_argument(

        "--use_data_dir",

        default="../use_data/",

        type=str,

        help="Directory to indices and similarity and idx2phrase"

    )

    parser.add_argument(

        "--result_dir",

        default="../result/",

        type=str,

        help="Directory to save clustering result"

    )

    args = parser.parse_args()

    args.indices_path = args.use_data_dir + 'indices.npy'

    args.similarity_path = args.use_data_dir + 'similarity.npy'

    args.idx2phrase_path = args.use_data_dir + 'idx2phrase.pkl'

    args.result_path = args.result_dir + 'clustering_result.pkl'

    args.phrase2idx_path = args.use_data_dir + 'phrase2idx.pkl'

    cluster_res = load_pickle(args.result_path)

    id2phrase = load_pickle(args.idx2phrase_path)

    phrase2id = load_pickle(args.phrase2idx_path)

    similarity = np.load(args.similarity_path)

    indices = np.load(args.indices_path)

    need_cluster_list = []

    need_cluster_length_list = []

    for key in tqdm(cluster_res):

        if len(cluster_res[key]) > MAX_CLUSTER_COUNT:

            need_cluster_list.append(key)

            need_cluster_length_list.append(len(cluster_res[key]))

            # break

    print(len(need_cluster_list))

    print(np.mean(need_cluster_length_list))

    threshold_list = [0.60]

    for threshold in threshold_list:

        print('threshold=', threshold)

        final_res = []

        for key in tqdm(cluster_res):

            if key not in need_cluster_list:

                final_res.append(list(cluster_res[key]))

                # print_cluster_to_file(f, list(cluster_res[key]))

            else:

                re_cluster_list = re_cluster(list(cluster_res[key]), mode, similarity, threshold)

                for cluster in re_cluster_list:

                    final_res.append(cluster)

                    # print_cluster_to_file(f, cluster)

        with open('../result/final_cluster_res.txt', 'w') as f:

            for cluster in tqdm(final_res):

                print_cluster_to_file(f, cluster)

                # f.write('-------')

                # f.write('RATIO\n')

                # f.write(str(re_cluster(list(cluster_res[key]), mode, similarity, threshold))+'\n')

                # f.write('NORMALIZE\n')

                # f.write(str(re_cluster(list(cluster_res[key]), 'normalize', similarity, threshold))+'\n')