import itertools
import random
from typing import List

import pandas as pd
from datasketch import MinHash, MinHashLSH
from Levenshtein import ratio
from ttdata.common import display_side_by_side


class DisjointSet:
    def __init__(self, n: int):
        self.parent = [i for i in range(n)]
        self.rank = [0 for _ in range(n)]

    def find(self, i: int) -> int:
        if self.parent[i] != i:
            self.parent[i] = self.find(self.parent[i])
        return self.parent[i]

    def union(self, i: int, j: int) -> None:
        i_root = self.find(i)
        j_root = self.find(j)

        if i_root == j_root:
            return

        if self.rank[i_root] < self.rank[j_root]:
            self.parent[i_root] = j_root
        elif self.rank[i_root] > self.rank[j_root]:
            self.parent[j_root] = i_root
        else:
            self.parent[j_root] = i_root
            self.rank[i_root] += 1

    def get_sets(self) -> list[list[int]]:
        sets = {}
        for i in range(len(self.parent)):
            root = self.find(i)
            if root in sets:
                sets[root].append(i)
            else:
                sets[root] = [i]
        return list(sets.values())


def preprocess(document: str):
    return document.lower().strip().replace("\r", "").replace("\n", "")


# Shingle size of 4-9 is normally used.
# Since we expect documents to be small, choose 4 as default.
# UPDATE: moving to size of 5 since it's less likely to divide nicely
def shingling(document: str, size: int = 5) -> set[bytes]:
    shingles = set()
    for i in range(len(document) - size + 1):
        shingles.add(document[i : i + size].encode("utf8"))
    return shingles


def get_clusters(
    documents: list[str],
    num_perm=256,
    lsh_threshold=0.9,
    similarity_threshold=0.9,
    filter_result=True,
) -> list[list[str]]:
    """
    Returns clusters of similar documents.

    Args:
        documents: A list of documents.
        num_perm: Number of permutations to use in MinHash.
            Greater num_perm -> greater accuracy, but slower and higher memory usage.
        lsh_threshold: Jaccard similarity threshold for LSH.
        similarity_threshold: Ratio for two strings to be considered similar, if filtering is used.
        filter_result: Whether to filter the candidates in the LSH buckets. This significantly increases the time required.
    """
    lsh = MinHashLSH(threshold=lsh_threshold, num_perm=num_perm)

    # For each document, create a MinHash object and add it to the LSH index
    # In the LSH index, documents that are similar will likely be hashed to the same bucket for some band.
    for i, document in enumerate(documents):
        minhash = MinHash(num_perm=num_perm)
        shingles = filter(None, shingling(preprocess(document)))
        minhash.update_batch(shingles)
        lsh.insert(i, minhash)
        if i > 0 and i % 5000 == 0:
            print(f"Inserted {i} documents")

    print("Processing LSH...")

    disjoint_set = DisjointSet(len(documents))
    for hashtable in lsh.hashtables:
        assert hashtable is not None, f"{lsh=}"  # keeps pyright happy
        for key, count in hashtable.itemcounts().items():
            # We consider documents that are hashed to the same bucket in at least one band to be candidates.
            if count > 1:
                similar_docs = list(hashtable.get(key))
                # Filter out candidate pairs below a specified similarity ratio since LSH is probabilistic.
                # This produces exact results at the expense of time.
                if filter_result:
                    for doc1, doc2 in itertools.combinations(similar_docs, 2):
                        if (
                            ratio(
                                documents[doc1],
                                documents[doc2],
                                score_cutoff=similarity_threshold,
                            )
                            > 0
                        ):
                            disjoint_set.union(doc1, doc2)
                else:
                    for i in range(1, len(similar_docs)):
                        disjoint_set.union(similar_docs[0], similar_docs[i])

    sets_of_indices = disjoint_set.get_sets()
    print(f"{len(sets_of_indices)} clusters found")

    sets_of_documents = [[documents[index] for index in s] for s in sets_of_indices]

    if True:  # debug branch, should disable or make optional later
        print("Results of cluster detection with LSH/minhash")
        print("=============================================")
        for _ in range(5):
            # pick a random cluster with >1 document in it
            cluster = random.choice([s for s in sets_of_indices if len(s) > 1])

            # pick two distinct random documents from the cluster
            doc1_idx, doc2_idx = random.sample(cluster, 2)

            # order so last comes second
            doc1_idx, doc2_idx = sorted([doc1_idx, doc2_idx])

            doc1 = documents[doc1_idx]
            doc2 = documents[doc2_idx]

            print("\n\nExample of two duplicates found:\n")
            print(display_side_by_side(doc1, doc2, colwidth=80))

    return sets_of_documents


def approximate_deduplicate_keep_last(
    frame: pd.DataFrame,
    dedup_cols: List[str],
    threshold: float,
) -> pd.DataFrame:
    """Deduplicate a dataframe by approximate Jaccard similarity of joined
    dedup_cols. Keeps just the last duplicate in each case."""

    # join the columns together into single strings
    joined_cols = frame[dedup_cols].apply(lambda row: " | ".join(row), axis=1)
    joined_cols_list = joined_cols.to_list()

    # assign cluster IDs to each string
    print(f"Clustering {len(joined_cols_list)} samples (may take a while)")
    clusters = get_clusters(
        documents=joined_cols_list,
        similarity_threshold=threshold,
        lsh_threshold=threshold,
    )
    doc_to_cluster_id = {
        doc: cluster_id for cluster_id, docs in enumerate(clusters) for doc in docs
    }

    # add a column with the cluster ID
    cluster_id_column = joined_cols.map(doc_to_cluster_id)
    cluster_id_column.name = "_approx_dedup_cluster_id"
    rv_frame = pd.concat([frame, cluster_id_column], axis=1)

    # deduplicate based on cluster ID, keeping last
    rv_frame = rv_frame.drop_duplicates(subset=[cluster_id_column.name], keep="last")
    del rv_frame[cluster_id_column.name]

    return rv_frame
