"""
    TODO:
    1. Model probability in the future.
"""
import logging
import os
from pathlib import Path
from typing import Tuple

import click
import numpy as np
import pandas as pd
import torch
from tqdm import tqdm
from transformers import AutoModelWithHeads, AutoTokenizer

from food_ke.entailment.custom_typing import PathLike
from food_ke.entailment.predict import get_predictions_for_candidates_df

logging.basicConfig(
    format="[%(filename)s:%(lineno)s - %(funcName)20s() ] %(message)s",
    level=logging.DEBUG,
)

if torch.cuda.is_available():
    device = torch.device("cuda")
else:
    device = torch.device("cpu")


def _compute_hyp_novelty_score(
    foodb_food_id: str, foodb_chemical_id: str, contents_kg: pd.DataFrame
):
    """Helper function to compute novelty score for a hypothesis.

    Args:
        hypothesis (pd.Series): The hypothesis to be scored.
        contents_kg (pd.DataFrame): The content in knowledge graph.

    Returns:
        (int): The novelty score of the hypothesis.

    """
    if foodb_food_id is None or foodb_chemical_id is None:
        # Novel if the food or the chemical is not in the FooDB.
        return 1

    content = contents_kg[
        (contents_kg["food_id"] == foodb_food_id)
        & (contents_kg["chemical_id"] == foodb_chemical_id)
    ]

    if content.empty:
        # Novel if the relation is not in the FooDB.
        return 1

    return 0


def rank_by_novelty(data, contents_kg, aggregation="sum"):
    """Rank the premises based on their hypothesis novelty.

    Args:
        data (pd.DataFrame): The data being ranked.
        contents_kg (pd.DataFrame): The content in knowledge graph.
        scoring (str): The scoring method.

    Returns:
        (list): Novelty scores.

    """
    data = data.copy()
    logging.info(
        f"Computing novelty scores and aggregating by {aggregation}...\n"
    )

    scores = []

    novelty_score = [
        _compute_hyp_novelty_score(
            foodb_food_id, foodb_chemical_id, contents_kg
        )
        for foodb_food_id, foodb_chemical_id in zip(
            data["hypothesis_food_id"],
            data["hypothesis_chemical_id"],
        )
    ]

    data["hypothesis_novelty"] = novelty_score

    premise_novelty = data.groupby("premise")["hypothesis_novelty"].agg(
        aggregation
    )

    data["premise_novelty"] = data.index.map(premise_novelty)

    return data


def rank_by_completeness(data: pd.DataFrame, scoring: str):
    """Rank the data into complete and incomplete dataframes in terms of
    the metadata, (e.g., food parts).

    Args:
        data (pd.DataFrame): The dataframe to be ranked.
        scoring (str): [mean|sum], Completeness scoring method.

    Returns:
        (list): Scores for premises.

    """
    data = data.copy()
    data["hypothesis_completeness"] = np.nan
    data["premise_completeness"] = np.nan
    logging.info(f"Computing completeness scores with {scoring} metric...")

    scores = []

    for premise in tqdm(data["premise"].unique().tolist()):
        try:
            counts_food_part = data["hypothesis_food_part"][
                premise
            ].value_counts()
        except AttributeError:
            # Only 1 hypothesis associated with the premise, so score = 0.
            scores.append(0)
            continue

        logging.info("counts food part:" + str(counts_food_part))

        if scoring == "mean":
            score = 1 - (
                counts_food_part["(whole plant)"] / counts_food_part.sum()
            )
        elif scoring == "sum":
            score = counts_food_part.sum() - counts_food_part["(whole plant)"]
        else:
            raise ValueError(f"Unknown scoring method: {scoring}")
        scores.append(score)

    data["hypothesis_completeness"] = scores

    return data


def rank_by_confidence_and_uncertainty(
    data, model, tokenizer, aggregation="sum", confidence_mode: str = "entails"
) -> Tuple[list, list]:
    data = data.copy()

    preds = get_predictions_for_candidates_df(data, model, tokenizer)
    preds.index = data.index
    data["proba_entails"] = preds["proba_entails"]
    data["proba_not_entails"] = preds["proba_not_entails"]

    if confidence_mode == "entails":
        logging.info("Using proba_entails as confidence value")
        data["hypothesis_confidence"] = preds["proba_entails"]

    elif confidence_mode == "both":
        logging.info(
            "Using max of proba_entails or proba_not_entails as confidence value"
        )
        data["hypotheses_confidence"] = preds[
            ["proba_entails", "proba_not_entails"]
        ].max(axis=1)
    else:
        raise ValueError(
            f"confidence_mode should be either 'entails' or 'both'"
        )

    logging.info(
        "Using min of proba_entails and proba_not_entails as uncertainty"
    )
    data["hypothesis_uncertainty"] = preds[
        ["proba_entails", "proba_not_entails"]
    ].min(axis=1)

    logging.info(
        "Computing premise confidence and uncertainty scores with "
        f"'{aggregation}' metric...\n"
    )

    data["premise_confidence"] = data["premise"].map(
        data.groupby("premise")["hypothesis_confidence"]
        .agg(aggregation)
        .to_dict()
    )
    data["premise_uncertainty"] = data["premise"].map(
        data.groupby("premise")["hypothesis_uncertainty"]
        .agg(aggregation)
        .to_dict()
    )

    return data


def rank_data_unlabelled(
    data_path: PathLike,
    model_path: PathLike,
    adapter_path: PathLike,
    tokenizer_name: PathLike,
    kg_content_path: PathLike,
    rank_novelty: bool = True,
    agg_novelty: str = "sum",
    rank_confidence: bool = True,
    agg_confidence: str = "sum",
    rank_completeness: bool = True,
    agg_completeness: str = "sum",
):
    data = pd.read_csv(data_path)
    data.index = data["premise"]
    kg_content = pd.read_csv(kg_content_path)

    model = AutoModelWithHeads.from_pretrained(model_path)
    model.active_adapters = model.load_adapter(
        adapter_path, model_name="roberta-base"
    )
    model = model.to(device)
    model.eval()
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)

    if rank_novelty:
        data = rank_by_novelty(data, kg_content, agg_novelty)
    if rank_confidence:
        data = rank_by_confidence_and_uncertainty(
            data,
            model,
            tokenizer,
            agg_confidence,
            confidence_mode="entails",
        )
    if rank_completeness:
        data = rank_by_completeness(data, agg_completeness)

    return data


@click.command()
@click.option("--data-path", type=str)
@click.option("--output-path", type=str)
@click.option("--model-path", type=str)
@click.option("--adapter-path", type=str)
@click.option("--tokenizer-name", type=str, default="roberta-base")
@click.option(
    "--kg-content-path",
    type=str,
    default=os.path.join(
        f"{os.path.abspath(os.path.dirname(__file__))}",
        "..",
        "..",
        "data",
        "knowledge_bases",
        "foodb",
        "foodb_2020_04_07_csv",
        "FooDB.csv",
    ),
)
@click.option("--score-weights", default=None)
@click.option("--scoring", default=None)
@click.option("--rank-novelty", type=bool, default=True)
@click.option("--agg-novelty", type=bool, default=True)
@click.option("--rank-confidence", type=bool, default=True)
@click.option("--agg-confidence", type=bool, default=True)
@click.option("--rank-completeness", type=bool, default=False)
@click.option("--agg-completeness", type=bool, default=True)
def main(
    data_path: PathLike,
    output_path: PathLike,
    model_path: PathLike,
    adapter_path: PathLike,
    tokenizer_name: str,
    kg_content_path: PathLike,
    score_weights=None,
    scoring=None,
    rank_novelty=True,
    agg_novelty="sum",
    rank_confidence=True,
    agg_confidence="sum",
    rank_completeness=False,
    agg_completeness="sum",
):
    if score_weights is not None or scoring is not None:
        raise NotImplementedError(
            "Score_weights and scoring are used to calculate premise scores which are "
            "not implemented yet."
        )

    data = rank_data_unlabelled(
        data_path,
        model_path,
        adapter_path,
        tokenizer_name,
        kg_content_path,
        rank_novelty,
        agg_novelty,
        rank_confidence,
        agg_confidence,
        rank_completeness,
        agg_completeness,
    )

    if output_path:
        data.to_csv(output_path, index=False)

    return data


if __name__ == "__main__":
    main()
