import json
from typing import Dict, List, Tuple
from datetime import datetime

import nltk

from inspect_ai.model import Model, GenerateConfig
from numpy.f2py.crackfortran import quiet

from astabench.evals.sqa.rubric import extract_json_from_response
from astabench.evals.sqa.pydantic_models import (
    assess_citation_group_schema,
    citation_eval_json_schema,
    precision_json_schema,
    AssessCitationGroup,
)
from astabench.evals.sqa.retry_utils import generate_with_retry


JUST_HAS_A_TITLE = "Paper content unavailable. The paper's title is: "


CITATION_EVAL_PROMPT = """As an Attribution Validator, your task is to verify whether a given reference can support the given claim. 
A claim can be either a plain sentence or a question followed by its answer. 
Specifically, your response should clearly indicate the relationship: Attributable, Contradictory or Extrapolatory. 
A contradictory error occurs when you can infer that the answer contradicts the fact presented in the context, while an extrapolatory error means that you cannot infer the correctness of the answer based on the information provided in the context. 
Output your response as a json with only a single key "output" and a value of one among - ("Attributable", "Contradictory", "Extrapolatory").
Claim: {claim}
Reference: {ref_excerpt}
"""


def append_prompt_logs(prompt_logs, new_prompt_logs):
    if "prompt_id" not in new_prompt_logs:
        return
    prompt_id = new_prompt_logs["prompt_id"]
    for k in new_prompt_logs.keys():
        if k == "prompt_id":
            continue
        prompt_logs[k].append({"prompt_id": prompt_id, "data": {k: new_prompt_logs[k]}})


def just_has_a_title(citation):
    return citation["snippets"].startswith(JUST_HAS_A_TITLE)


def clean_citation(c: str) -> str:
    chars_to_remove = "[]()"
    for char in chars_to_remove:
        c = c.replace(char, "")
    return c


def citation_intersection(supporting_cits, bad_cits):
    supporting_cits = {clean_citation(str(c)) for c in supporting_cits}
    bad_cits = {clean_citation(str(c)) for c in bad_cits}
    return list(supporting_cits.intersection(bad_cits))


class CitationEval:
    def __init__(self, model: Model, temperature: float = 0.5, top_p: float = 0.95):
        self.grader_model = model
        self.temperature = temperature
        self.top_p = top_p

    @staticmethod
    def map_output_to_score(output) -> int:
        if output:
            output = output["output"]
        return 1 if output == "Attributable" else 0

    async def recall(
        self, claim: str, ref_excerpts: List[str], cit_ids: List[str]
    ) -> int:
        """
        Given a sentence claim with citations and the corresponding citation excerpts,
        the recall score is 1 if the concatenated excerpts entail the claim and 0 otherwise.
        """
        excerpt_concat = "\n" + "\n".join(
            "{}: {}".format(cit_id, t) for cit_id, t in zip(cit_ids, ref_excerpts)
        )
        prompt = CITATION_EVAL_PROMPT.format(claim=claim, ref_excerpt=excerpt_concat)
        config = GenerateConfig(
            response_schema=citation_eval_json_schema,
            temperature=self.temperature,
            top_p=self.top_p,
        )
        result, parsed, num_retries = await generate_with_retry(
            self.grader_model, prompt, config
        )
        return self.map_output_to_score(parsed)

    async def precision_all_together(self, claim: str, ref_excerpts: List[str]) -> int:
        prompt = (
            "Which references support the following claim?\nClaim: "
            + claim
            + "\n\n"
            + "\n".join(
                "Reference {}: {}".format(i, c) for i, c in enumerate(ref_excerpts, 1)
            )
            + "\nReturn your response as a JSON object with two keys `supporting` and `non_supporting`. `supporting` is a list of the citations that support the claim. `non_supporting` is a list of the citations that do not support the claim. Each citation should appear in either `supporting` or `non_supporting` but not both."
        )
        config = GenerateConfig(
            response_schema=precision_json_schema,
            temperature=self.temperature,
            top_p=self.top_p,
        )
        result, parsed, num_retries = await generate_with_retry(
            self.grader_model, prompt, config
        )
        o = parsed.get("supporting", [])
        return len(o)

    async def score_claim(
        self, claim: str, ref_excerpts: List[str], cit_ids: List[str]
    ) -> Dict[str, int]:
        if not ref_excerpts:
            return {"recall_ct": 0, "precision_ct": 0}
        recall_score = await self.recall(claim, ref_excerpts, cit_ids)
        precision_score = (
            0
            if not recall_score
            else (
                1
                if len(ref_excerpts) == 1
                else await self.precision_all_together(claim, ref_excerpts)
            )
        )
        return {"recall_ct": recall_score, "precision_ct": precision_score}

    async def evaluate(
        self, claim_evidence_pairs: List[Tuple[str, List[str], List[str]]]
    ) -> Dict[str, float]:
        recall_ct, precision_rt, total_attributions = 0, 0, 1e-10
        for claim, excerpts, cit_ids in claim_evidence_pairs:
            # Sometimes elicit doesnt retrieve the full text:
            if "Full text was not retrieved." in claim:
                continue
            scores = await self.score_claim(claim, excerpts, cit_ids)
            recall_ct += scores["recall_ct"]
            precision_rt += scores["precision_ct"]
            total_attributions += len(excerpts)
        try:
            recall, precision = (
                recall_ct / len(claim_evidence_pairs),
                precision_rt / total_attributions,
            )
        except ZeroDivisionError:
            recall, precision = 0, 0
        f1 = (
            2 * (recall * precision) / (recall + precision) if recall + precision else 0
        )
        return {"citation_recall": recall, "citation_precision": precision, "f1": f1}

    def count_sentences(self, citation_group: str) -> int:
        nltk.download("punkt_tab", quiet=True)
        return len(nltk.sent_tokenize(citation_group))

    async def score_citation_group(
        self,
        citation_group: str,
        citations: List[Dict[str, str]],
        is_retrieverless: bool = False,
    ):
        if not citations:
            return {
                "n_attributable": 0,
                "n_extrapolatory": self.count_sentences(citation_group),
                "n_half_credit_claims": 0,
                "supporting_counts": [],
                "non_supporting_counts": [],
                "n_half_credit_citations": [],
            }, {}
        prompt = "You are a claim validator. For each claim made in the following text you will determine if it is supported by the quote from it's corresponding inline citations. As is typically done in academic writing, assume that consecutive sentences can share citations. Make sure to also include claims presented in table format. For references with only the title available (ie no quotes from the reference are included), judge them as `supporting` if the title indicates that the paper is likely relevant to the claim being considered. Return a JSON object with a single key `claims` which is a list of `claim` objects, one for each sentence in the text. Each `claim` object contains the claim itself (`text`), a list of `supporting` inline citations and `non_supporting` inline citations and finally a boolean `is_fully_supported` which indicates if the claim is entirely supported by the quotations in the associated citations. Each inline citation corresponding to that claim should appear in either `supporting` or `non_supporting`, but not both. Each claim made in the text should appear in your output, but you should skip sentences covering high level introductory information.\n\nText:\n{}\n\nReferences:\n{}"
        prompt = prompt.format(
            citation_group,
            "\n\n".join(a["id"] + ": " + a["snippets"] for a in citations),
        )
        config = GenerateConfig(
            response_schema=assess_citation_group_schema,
            temperature=self.temperature,
            top_p=self.top_p,
        )
        result, out, num_retries = await generate_with_retry(
            self.grader_model, prompt, config, desired_schema=AssessCitationGroup
        )
        claims = out["claims"]
        n_attributable, n_extrapolatory, n_half_credit_claims = 0, 0, 0
        supporting_counts, non_supporting_counts, n_half_credit_citations = [], [], []
        half_credit = [c["id"] for c in citations if just_has_a_title(c)]
        for claim in claims:
            supported = claim["is_fully_supported"] and claim["supporting"]
            n_attributable += 1 if supported else 0
            n_extrapolatory += 0 if supported else 1
            supporting_citations = claim.get("supporting", [])
            n_half_credit_citations.append(
                len(supporting_citations)
                if is_retrieverless
                else len(citation_intersection(supporting_citations, half_credit))
            )
            supporting_counts.append(len(supporting_citations))
            n_half_credit_claims += (
                1
                if supported
                and (
                    (supporting_counts[-1] == n_half_credit_citations[-1])
                    or is_retrieverless
                )
                else 0
            )
            non_supporting_counts.append(len(claim.get("non_supporting", [])))
        prompt_logs = {
            "step_name": "score_citation_group",
            "prompt_id": str(hash(prompt)),
            "text": citation_group,
            "citations": citations,
            "half_credit": half_credit,
            "prompt": prompt,
            "claims": claims,
            "temperature": self.temperature,
            "top_p": self.top_p,
            "timestamp": str(datetime.now()),
            "num_retries": num_retries,
        }
        return {
            "n_attributable": n_attributable,
            "n_extrapolatory": n_extrapolatory,
            "n_half_credit_claims": n_half_credit_claims,
            "supporting_counts": supporting_counts,
            "non_supporting_counts": non_supporting_counts,
            "n_half_credit_citations": n_half_credit_citations,
        }, prompt_logs

    async def evaluate_citation_groups(
        self, citation_group_cit_pairs, is_retrieverless: bool = False
    ) -> Tuple[Dict[str, float], Dict]:
        n_attributable, n_extrapolatory, n_half_credit = 0, 0, 0
        precisions = []
        prompt_logs = {
            "step_name": [],
            "text": [],
            "citations": [],
            "half_credit": [],
            "prompt": [],
            "claims": [],
            "temperature": [],
            "top_p": [],
            "timestamp": [],
            "num_retries": [],
        }
        for citation_group, citations in citation_group_cit_pairs:
            scores, meta = await self.score_citation_group(
                citation_group, citations, is_retrieverless=is_retrieverless
            )
            append_prompt_logs(prompt_logs, meta)
            n_attributable += scores["n_attributable"]
            n_extrapolatory += scores["n_extrapolatory"]
            n_half_credit += scores["n_half_credit_claims"]
            precisions.extend(
                [
                    (s - 0.5 * p) / (s + e)
                    for s, e, p in zip(
                        scores["supporting_counts"],
                        scores["non_supporting_counts"],
                        scores["n_half_credit_citations"],
                    )
                    if s + e
                ]
            )
        try:
            recall = (
                (n_attributable - 0.5 * n_half_credit)
                / (n_attributable + n_extrapolatory)
                if n_attributable + n_extrapolatory
                else 0
            )
            precision = sum(precisions) / len(precisions) if precisions else 0
        except ZeroDivisionError:
            recall, precision = 0, 0
        f1 = (
            2 * (recall * precision) / (recall + precision) if recall + precision else 0
        )
        metadata = {"prompt_logs": prompt_logs}
        return {
            "citation_recall": recall,
            "citation_precision": precision,
            "f1": f1,
        }, metadata
