import logging
import math
from enum import Enum

from astabench.evals.paper_finder.datamodel import (
    BroadScorerCriteria,
    CorpusId,
    ExpectedAgentOutput,
    SpecificPaperScorerCriteria,
)
from astabench.evals.paper_finder.paper_finder_utils import (
    init_references,
    update_references,
)
from astabench.evals.paper_finder.relevance import (
    Document,
    Relevance,
    bench_rj_2_pf_rj,
    load_relevance_judgement,
)

LOGGER = logging.getLogger(__name__)

MAX_NORMALIZATION_MULTIPLIER_FACTOR = 10
MIN_NORMALIZATION_MULTIPLIER_FACTOR = 2

_normalizer_reference, _detailed_reference = None, None


class KTypes(Enum):
    KNOWN = "known"
    ESTIMATED = "estimated"


class KValues(Enum):
    AT_ESTIMATE = "estimate"
    AT_30 = "30"
    AT_FULL = "full"


def get_normalizer_references() -> tuple[dict[str, int], dict[str, dict[str, str]]]:
    global _normalizer_reference, _detailed_reference
    if _normalizer_reference is None or _detailed_reference is None:
        # we only want this to happen once during initialization phase
        _normalizer_reference, _detailed_reference = init_references(
            MAX_NORMALIZATION_MULTIPLIER_FACTOR, MIN_NORMALIZATION_MULTIPLIER_FACTOR
        )
    return _normalizer_reference, _detailed_reference


async def get_llm_relevance(
    output: ExpectedAgentOutput, metric: BroadScorerCriteria
) -> dict[CorpusId, str]:
    _, detailed_reference = get_normalizer_references()

    judgements = dict()
    to_judge = []
    qid = output.query_id
    for result in output.results:
        if result.paper_id in metric.known_to_be_good:
            judgements[result.paper_id] = Relevance.PERFECT.value
        elif (
            detailed_reference
            and detailed_reference.get(qid)
            and detailed_reference[qid].get(result.paper_id)
        ):
            judgements[result.paper_id] = detailed_reference[qid][result.paper_id]
        else:
            to_judge.append(
                Document(corpus_id=result.paper_id, markdown=result.markdown_evidence)
            )
    new_judgements = dict()
    if to_judge:
        new_judgements = await load_relevance_judgement(
            to_judge, metric.relevance_criteria
        )
        await update_references(qid, new_judgements)
        judgements.update(new_judgements)

    new_perfects = len(
        [j for j in new_judgements.values() if j == Relevance.PERFECT.value]
    )
    LOGGER.info(f"New judgements: {new_perfects}")

    return judgements


def count_relevant(
    judgements: dict[CorpusId, str], top_k_results: list[str] | None = None
) -> int:

    relevant_predictions = len(
        [
            paper_id
            for paper_id, judgement in judgements.items()
            if Relevance.PERFECT.value == judgement
            and (not top_k_results or paper_id in top_k_results)
        ]
    )

    return relevant_predictions


def get_predefined_judgements(
    output: ExpectedAgentOutput, metric: SpecificPaperScorerCriteria
) -> dict[str, str]:
    return {
        result.paper_id: (
            Relevance.PERFECT
            if result.paper_id in metric.corpus_ids
            else Relevance.NOT_RELEVANT
        )
        for result in output.results
    }


def calc_recall_at_k(
    output: ExpectedAgentOutput,
    query_id: str,
    judgements: dict[CorpusId, str],
    k_type: KTypes,
) -> dict[str, float]:

    # NOTE: the judge might have failures on its own fault - so no reason to punish the caller. that's why we:
    #   (1) judge all and then take top k.
    #   (2) treat llm_judgements and not output.results as the full set

    to_return = dict()

    normalizer_reference, _ = get_normalizer_references()

    total_relevant = normalizer_reference.get(query_id)
    relevant_predictions = count_relevant(judgements)
    to_return[f"relevant_predictions_at_{KValues.AT_FULL.value}"] = relevant_predictions
    to_return[f"{k_type.value}_recall_at_{KValues.AT_FULL.value}"] = (
        relevant_predictions / total_relevant
    )

    for name, k in {
        KValues.AT_ESTIMATE.value: total_relevant,
        KValues.AT_30.value: 30,
    }.items():
        top_k_results = [
            result.paper_id
            for result in output.results
            if result.paper_id in judgements.keys()
        ][:k]
        relevant_predictions_at_k = count_relevant(judgements, top_k_results)
        to_return[f"{k_type.value}_recall_at_{name}"] = (
            relevant_predictions_at_k / total_relevant
        )

    return to_return


# computes the precision of perfect results out of all predictions
def calc_precision(
    judgements: dict[CorpusId, str],
) -> dict[str, float]:
    relevant_predictions = count_relevant(judgements)
    return {
        "precision": relevant_predictions / len(judgements) if len(judgements) else 0.0
    }


def find_dcg(element_list: list[int]) -> float:
    score = 0.0
    for order, rank in enumerate(element_list, 1):
        score += float(rank) / math.log(order + 1)
    return score


def lower_bound_corrected_ndcg(hypothesis: list[int]) -> dict[str, float]:
    max_dcg = find_dcg(sorted(hypothesis, reverse=True))
    min_dcg = find_dcg(sorted(hypothesis))
    return {
        "rank": (
            ((find_dcg(hypothesis) - min_dcg) / (max_dcg - min_dcg))
            if (max_dcg - min_dcg)
            else 0.0
        )
    }


def _calc_any_f(f_elements: list[float]) -> float:
    return (
        (len(f_elements) / sum([(1 / element) for element in f_elements]))
        if all(f_elements)
        else 0.0
    )


def calc_adjusted_f1(
    output: ExpectedAgentOutput,
    query_id: str,
    judgements: dict[CorpusId, str],
    k_type: KTypes,
    k_value: KValues,
) -> dict[str, float]:
    recall = calc_recall_at_k(output, query_id, judgements, k_type)
    rank = lower_bound_corrected_ndcg(
        [bench_rj_2_pf_rj[j] for j in judgements.values()]
    )

    f1 = _calc_any_f(
        [rank["rank"], recall[f"{k_type.value}_recall_at_{k_value.value}"]]
    )
    return {"adjusted_f1": f1, **recall, **rank}


def calc_standard_f1(
    output: ExpectedAgentOutput,
    query_id: str,
    judgements: dict[CorpusId, str],
    k_type: KTypes,
    k_value: KValues,
) -> dict[str, float]:
    recall = calc_recall_at_k(output, query_id, judgements, k_type)
    precision = calc_precision(judgements)

    f1 = _calc_any_f(
        [precision["precision"], recall[f"{k_type.value}_recall_at_{k_value.value}"]]
    )
    return {"standard_f1": f1, **recall, **precision}
