import logging
import optparse
import os

import dedupe
import polars as pl


def clean_strings(df: pl.DataFrame, columns=["title", "abstract"]) -> pl.DataFrame:
    """
    Some simple string cleaning to remove meaningless differences.
    Separates abstract into chunks to make edit distance calculation less time
    consuming when using dedupe.
    """
    for col in columns:
        # lowercase the strings
        df = df.with_columns(pl.col(col).str.to_lowercase().alias(f"{col}_clean"))

        # remove extra whitespace
        df = df.with_columns(pl.col(f"{col}_clean").str.replace(r"\s+", ""))

    df = df.with_columns(
        pl.col("abstract")
        .str.slice(0, length=100)
        .str.zfill(100)
        .alias("abstract_first100")
    )
    df = df.with_columns(
        pl.col("abstract")
        .str.slice(100, length=100)
        .str.zfill(100)
        .alias("abstract_second100")
    )
    df = df.with_columns(
        pl.col("abstract")
        .str.slice(200, length=100)
        .str.zfill(100)
        .alias("abstract_third100")
    )
    df = df.with_columns(
        pl.col("abstract")
        .str.slice(300, length=100)
        .str.zfill(100)
        .alias("abstract_fourth100")
    )
    df = df.with_columns(
        pl.col("abstract")
        .str.slice(400, length=100)
        .str.zfill(100)
        .alias("abstract_fifth100")
    )

    return df


def read_data(filename: str) -> pl.DataFrame:
    """
    Read in our data from parquet and do some simple cleaning.
    """
    df = pl.read_parquet(filename)

    # remove blank fields
    df = df.filter((pl.col("title") != "") & (pl.col("abstract") != ""))

    # remove exact duplicates
    df = df.unique(subset=["title"])
    df = df.unique(subset=["abstract"])

    # add a record ID column, which we'll use later to merge dedupe results
    df = df.with_row_count(name="record_id")

    return df


def data_to_dict(df: pl.DataFrame) -> dict:
    """
    Format the data in the dict structure that dedupe expects.
    """
    df = df.to_dicts()
    indexes = range(len(df))
    return {i: row for i, row in zip(indexes, df)}


if __name__ == "__main__":
    # ## Logging

    # Dedupe uses Python logging to show or suppress verbose output. This
    # code block lets you change the level of loggin on the command
    # line. You don't need it if you don't want that. To enable verbose
    # logging, run `python examples/csv_example/csv_example.py -v`
    optp = optparse.OptionParser()
    optp.add_option(
        "-v",
        "--verbose",
        dest="verbose",
        action="count",
        help="Increase verbosity (specify multiple times for more)",
    )
    (opts, args) = optp.parse_args()
    log_level = logging.WARNING
    if opts.verbose:
        if opts.verbose == 1:
            log_level = logging.INFO
        elif opts.verbose >= 2:
            log_level = logging.DEBUG
    logging.getLogger().setLevel(log_level)

    # ## Setup

    input_file = "data/NSF/nsf_2022.parquet"
    output_file = "data/NSF/nsf_2022_deduped.parquet"
    settings_file = "data/nsf_dedupe_learned_settings"
    training_file = "data/nsf_dedupe_training.json"

    print("importing data ...")
    df = read_data(input_file)
    df_cleaned = clean_strings(df)
    data_d = data_to_dict(df_cleaned)

    # If a settings file already exists, we'll just load that and skip training
    if os.path.exists(settings_file):
        print("reading from", settings_file)
        with open(settings_file, "rb") as f:
            deduper = dedupe.StaticDedupe(f)
    else:
        # ## Training

        # Define the fields dedupe will pay attention to
        fields = [
            {"field": "title_clean", "type": "String"},
            {"field": "abstract_first100", "type": "String"},
            {"field": "abstract_second100", "type": "String"},
            {"field": "abstract_third100", "type": "String"},
            {"field": "abstract_fourth100", "type": "String"},
            {"field": "abstract_fifth100", "type": "String"},
        ]

        # Create a new deduper object and pass our data model to it.
        deduper = dedupe.Dedupe(fields)

        # If we have training data saved from a previous run of dedupe,
        # look for it and load it in.
        # __Note:__ if you want to train from scratch, delete the training_file
        if os.path.exists(training_file):
            print("reading labeled examples from ", training_file)
            with open(training_file, "rb") as f:
                deduper.prepare_training(data_d, f)
        else:
            print("preparing for training...")
            deduper.prepare_training(data_d)

        # ## Active learning
        # Dedupe will find the next pair of records
        # it is least certain about and ask you to label them as duplicates
        # or not.
        # use 'y', 'n' and 'u' keys to flag duplicates
        # press 'f' when you are finished
        print("starting active labeling...")
        dedupe.console_label(deduper)

        # Using the examples we just labeled, train the deduper and learn
        # blocking predicates
        print("training...")
        deduper.train()

        # When finished, save our training to disk
        with open(training_file, "w") as tf:
            deduper.write_training(tf)

        # Save our weights and predicates to disk.  If the settings file
        # exists, we will skip all the training and learning next time we run
        # this file.
        with open(settings_file, "wb") as sf:
            deduper.write_settings(sf)

    # ## Clustering

    # `partition` will return sets of records that dedupe
    # believes are all referring to the same entity.

    print("clustering...")
    clustered_dupes = deduper.partition(data_d, 0.5)

    print("# duplicate sets", len(clustered_dupes))

    # ## Writing Results

    # put the partition output into a more usable format
    cluster_membership = {}
    for cluster_id, (records, scores) in enumerate(clustered_dupes):
        for record_id, score in zip(records, scores):
            cluster_membership[record_id] = {
                "cluster_id": cluster_id,
                "confidence_score": score,
            }

    # make a dataframe of the cluster assignments and merge to main data
    membership_data = [
        {
            "record_id": int(record_id),
            "cluster_id": info["cluster_id"],
            "confidence_score": info["confidence_score"],
        }
        for record_id, info in cluster_membership.items()
    ]
    membership_df = pl.DataFrame(membership_data)
    membership_df = membership_df.with_columns(pl.col("record_id").cast(pl.UInt32))

    df = df.join(membership_df, on="record_id", how="left")

    df.write_parquet(output_file)
