# coding: utf-8

# Base Dependencies
# ------------------
import numpy as np
from collections import Counter
from typing import Dict
from pathlib import Path
from os.path import join as pjoin
from tqdm import tqdm

# Local Dependencies
# ------------------
from models.relation_collection import RelationCollection

# 3rd-Party Dependencies
# ----------------------
import pandas as pd
from tabulate import tabulate

# Constants
# ---------
from constants import N2C2_REL_TYPES, DDI_ALL_TYPES, N2C2_PATH, DDI_PATH

TABLE_FORMAT = "latex"


# Main Functions
# ---------------
def generate_statistics(dataset: str, collections: Dict[str, RelationCollection]):
    if dataset == "n2c2":
        return generate_statistics_n2c2(collections)
    elif dataset == "ddi":
        return generate_statistics_ddi(collections)
    else:
        raise ValueError("unsupported dataset '{}'".format(dataset))


def generate_statistics_n2c2(collections: Dict[str, RelationCollection]):
    """Generates the statistics for the n2c2 dataset"""

    df_counts = {
        "relation": [],
        "train_positive": [],
        "train_negative": [],
        "test_positive": [],
        "test_negative": [],
    }
    df_seq_lengths = {
        "relation": [],
        "train_min": [],
        "train_avg": [],
        "train_max": [],
        "test_min": [],
        "test_avg": [],
        "test_max": [],
    }

    # number of relations per type of relation
    for rel_type in tqdm(N2C2_REL_TYPES):
        df_counts["relation"].append(rel_type)
        df_seq_lengths["relation"].append(rel_type)

        for split, collection in collections.items():
            subcollection = collection.type_subcollection(rel_type)

            # add counts to data
            count_labels = Counter(subcollection.labels)

            df_counts[split + "_negative"].append(count_labels[0])
            df_counts[split + "_positive"].append(count_labels[1])

            # add sequence length to dataframe
            seq_lengths = list(
                map(lambda rel: len(rel.text.split()), subcollection.relations)
            )

            df_seq_lengths[split + "_min"].append(min(seq_lengths))
            df_seq_lengths[split + "_avg"].append(sum(seq_lengths) / len(subcollection))
            df_seq_lengths[split + "_max"].append(max(seq_lengths))

    df_counts = pd.DataFrame(df_counts)
    df_seq_lengths = pd.DataFrame(df_seq_lengths)

    # add totals to counts
    df_counts["train_total"] = df_counts["train_positive"] + df_counts["train_negative"]
    df_counts["test_total"] = df_counts["test_positive"] + df_counts["test_negative"]
    df_counts["total_positive"] = (
        df_counts["train_positive"] + df_counts["test_positive"]
    )
    df_counts["total_negative"] = (
        df_counts["train_negative"] + df_counts["test_negative"]
    )
    df_counts["total"] = df_counts["total_positive"] + df_counts["total_negative"]
    df_counts.loc[len(df_counts)] = ["Total"] + [
        df_counts[col].sum() for col in df_counts.columns[1:]
    ]

    all_train_seq_lengths = list(map(lambda rel: len(rel.text.split()), collections["train"].relations))
    all_test_seq_lengths = list(map(lambda rel: len(rel.text.split()), collections["test"].relations))
    df_seq_lengths = df_seq_lengths.append(
        {
            "relation": "Overall",
            "train_min": min(all_train_seq_lengths),
            "train_avg": sum(all_train_seq_lengths) / len(all_train_seq_lengths),
            "train_max": max(all_train_seq_lengths),
            "test_min": min(all_test_seq_lengths),
            "test_avg": sum(all_test_seq_lengths) / len(all_test_seq_lengths),
            "test_max": max(all_test_seq_lengths),
        },
        ignore_index=True,
    )

    # select and reorder columns
    df_counts = df_counts.loc[:, ["relation", "train_positive", "train_negative", "train_total", "test_positive", "test_negative", "test_total", "total"]]

    # save data to csv
    df_counts.to_csv(Path(pjoin(N2C2_PATH, "counts.csv")), index=False)
    df_seq_lengths.to_csv(Path(pjoin(N2C2_PATH, "seq_length.csv")), index=False)

    # print statistics
    print("\n **** Statistics of the N2C2 Dataset ****")
    print("Counts:")
    print(tabulate(df_counts, headers="keys", tablefmt=TABLE_FORMAT))
    print("Seq Length:")
    print(tabulate(df_seq_lengths, headers="keys", tablefmt=TABLE_FORMAT))


def generate_statistics_ddi(collections: Dict[str, RelationCollection]) -> None:
    """Generates the statistics of the DDI dataset"""

    df_counts = {"relation": [], "train": [], "test": []}
    df_seq_lengths = {
        "relation": [],
        "train_min": [],
        "train_avg": [],
        "train_max": [],
        "test_min": [],
        "test_avg": [],
        "test_max": [],
    }

    for rel_type in DDI_ALL_TYPES:
        df_counts["relation"].append(rel_type)
        df_seq_lengths["relation"].append(rel_type)

        for split, collection in collections.items():
            subcollection = collection.type_subcollection(rel_type)

            df_counts[split].append(len(subcollection))

            seq_lengths = list(
                map(lambda rel: len(rel.text.split()), subcollection.relations)
            )
            df_seq_lengths[split + "_min"].append(min(seq_lengths))
            df_seq_lengths[split + "_avg"].append(sum(seq_lengths) / len(subcollection))
            df_seq_lengths[split + "_max"].append(max(seq_lengths))

    # convert to dataframes
    df_counts = pd.DataFrame(df_counts)
    df_seq_lengths = pd.DataFrame(df_seq_lengths)

    # add totals
    train_negative = df_counts.loc[(df_counts["relation"] == "NO-REL"), "train"].values[
        0
    ]
    train_positive = df_counts.loc[(df_counts["relation"] != "NO-REL"), "train"].sum()
    test_negative = df_counts.loc[(df_counts["relation"] == "NO-REL"), "test"].values[0]
    test_positive = df_counts.loc[(df_counts["relation"] != "NO-REL"), "test"].sum()
    train_total = train_positive + train_negative
    test_total = test_positive + test_negative
    total = train_total + test_total

    # add positive row
    df_counts.loc[len(df_counts)] = ["Total Positive", train_positive, test_positive]

    # add totals
    df_counts["total"] = df_counts["train"] + df_counts["test"]
    df_counts.loc[len(df_counts)] = [" Total", train_total, test_total, total]

    all_train_seq_lengths = list(map(lambda rel: len(rel.text.split()), collections["train"].relations))
    all_test_seq_lengths = list(map(lambda rel: len(rel.text.split()), collections["test"].relations))
    df_seq_lengths = df_seq_lengths.append(
        {
            "relation": "Overall",
            "train_min": min(all_train_seq_lengths),
            "train_avg": sum(all_train_seq_lengths) / len(all_train_seq_lengths),
            "train_max": max(all_train_seq_lengths),
            "test_min": min(all_test_seq_lengths),
            "test_avg": sum(all_test_seq_lengths) / len(all_test_seq_lengths),
            "test_max": max(all_test_seq_lengths),
        },
        ignore_index=True,
    )

    # save data to csv
    df_counts.to_csv(Path(pjoin(DDI_PATH, "counts.csv")), index=False)
    df_seq_lengths.to_csv(Path(pjoin(DDI_PATH, "seq_length.csv")), index=False)

    # print statistics
    print("\n **** Statistics of the DDI Dataset ****")
    print("Counts:")
    print(tabulate(df_counts, headers="keys", tablefmt=TABLE_FORMAT))
    print("Seq Length:")
    print(tabulate(df_seq_lengths, headers="keys", tablefmt=TABLE_FORMAT))