import numpy as np

from tqdm import tqdm

import pickle

import argparse

# reference: https://stackoverflow.com/questions/3067529/a-set-union-find-algorithm

class DisjointSet(object):

    def __init__(self):

        self.leader = dict()     # maps a member to the group's leader

        self.group = dict()  # maps a group leader to the group (which is a set)

    def add(self, a, b):

        leadera = self.leader.get(a)   

        leaderb = self.leader.get(b)   

        if leadera is not None:

            if leaderb is not None:

                if leadera == leaderb: return # nothing to do

                groupa = self.group[leadera]

                groupb = self.group[leaderb]

                if len(groupa) < len(groupb):

                    a, leadera, groupa, b, leaderb, groupb = b, leaderb, groupb, a, leadera, groupa

                groupa |= groupb

                del self.group[leaderb]

                for k in groupb:

                    self.leader[k] = leadera

            else:

                self.group[leadera].add(b)

                self.leader[b] = leadera

        else:

            if leaderb is not None:

                self.group[leaderb].add(a)

                self.leader[a] = leaderb

            else:

                self.leader[a] = self.leader[b] = a

                self.group[a] = set([a, b])

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'

    indices = np.load(args.indices_path)

    similarity = np.load(args.similarity_path)

    with open(args.idx2phrase_path, 'rb') as f:

        idx2phrase = pickle.load(f)

    ds = DisjointSet()

    for idxi in tqdm(range(indices.shape[0])):

        a = idx2phrase[idxi]

        if len(a) <= 3:

            continue

        for idxj in range(indices.shape[1]):

            if similarity[idxi, idxj] > 0.8:

                b = idx2phrase[indices[idxi][idxj]]

                if len(b) <= 3:

                    continue

                ds.add(a, b)

    with open(args.result_path, 'wb') as f:

        pickle.dump(ds.group, f)

    print('done')