import pandas as pd
import inspect_ai.log as ialog
from astabench.helpers.prompt_logs import PromptLogs, build_df_from_prompt_logs
import re

logs = [
    "logs/2025-05-08T20-53-49+00-00_cqa-debug_VgrNSrCewCmjCAKaYgBABt.eval",
    "logs/2025-05-08T20-53-49+00-00_cqa-debug_23rPEqmz2Y9RMMUyDLYFPd.eval",
    "logs/2025-05-08T20-53-50+00-00_cqa-debug_7kFcw4FUGibK4wXquNaJfn.eval",
    "logs/2025-05-08T20-53-50+00-00_cqa-debug_J8UqiBiYnPTZjfnABcQYfJ.eval",
    "logs/2025-05-08T20-53-50+00-00_cqa-debug_5MAjYRLmnx8WyV9H8DtYpQ.eval",
    "logs/2025-05-08T20-53-50+00-00_cqa-debug_L2vqTMPV6cx3bG35VWjf66.eval",
    "logs/2025-05-08T20-53-50+00-00_cqa-debug_HPQkbF4mAPnJeUQ5DmGB3R.eval",
    "logs/2025-05-08T20-53-50+00-00_cqa-debug_YUY3W8hhDNStYMqUhsZsc7.eval",
    "logs/2025-05-08T20-53-50+00-00_cqa-debug_AfusbAYfAmcEb2gKmJ8fDv.eval",
    "logs/2025-05-08T20-53-50+00-00_cqa-debug_oWgMkFxeZrnHZJFhXn29cV.eval",
    "logs/2025-05-08T20-53-50+00-00_cqa-debug_aKw8Rbrq4ZgnYk99Qf57Nz.eval",
    "logs/2025-05-08T20-53-50+00-00_cqa-debug_iP9LWg69LPvChiHzt7nUiT.eval",
    "logs/2025-05-08T20-53-50+00-00_cqa-debug_6xePw2j4TrJcuRMgjkG77r.eval",
    "logs/2025-05-08T20-53-50+00-00_cqa-debug_fbkaSNBTXoZRLcWTHNRjf3.eval",
    "logs/2025-05-08T20-53-50+00-00_cqa-debug_XgmtfwgwyqnDBaTG5kkoTv.eval",
    "logs/2025-05-08T20-53-50+00-00_cqa-debug_8DqvBP7V6JanVDzYbYovXt.eval",
    "logs/2025-05-08T20-53-51+00-00_cqa-debug_2kUjmTHWVo3csci9peq8mQ.eval",
    "logs/2025-05-08T20-53-51+00-00_cqa-debug_PGF3gQn9tCEcSy3ACBfrPy.eval",
    "logs/2025-05-08T20-53-51+00-00_cqa-debug_Qnn6hJ3qdoNFMDKNfnJAJL.eval",
    "logs/2025-05-08T20-53-51+00-00_cqa-debug_nMvxDNeFE7YLyVsDhHzgeG.eval",
    "logs/2025-05-08T20-53-51+00-00_cqa-debug_4Raw23Sf59v64mtpnazA7e.eval",
    "logs/2025-05-08T20-53-51+00-00_cqa-debug_cCQoSuwJ5j8BiAfpTookcC.eval",
    "logs/2025-05-08T20-53-51+00-00_cqa-debug_RFTcKudhiUz56speq6TP3N.eval",
    "logs/2025-05-08T20-53-51+00-00_cqa-debug_Q9XJWXfnxbJ3PgVwNtPY2A.eval",
    "logs/2025-05-08T20-53-51+00-00_cqa-debug_QVGGRUrgjQkG26ssoqfr9W.eval",
    "logs/2025-05-08T20-53-51+00-00_cqa-debug_X7VGcje4SPuJZ9wPwanbxP.eval",
    "logs/2025-05-08T20-53-51+00-00_cqa-debug_7ZmKrF7Huimq72sCL6sDWp.eval",
    "logs/2025-05-08T20-53-51+00-00_cqa-debug_h4W5Q9dNLxALstwPWXVLDe.eval",
    "logs/2025-05-08T20-53-51+00-00_cqa-debug_ETFo8QawnHt4oEnjaHRdZR.eval",
    "logs/2025-05-08T20-53-51+00-00_cqa-debug_ZQjEkAPH64YkKccpZVMfwZ.eval",
    "logs/2025-05-08T20-53-51+00-00_cqa-debug_M7CAqc9HgyHbiGAkj9GDKJ.eval",
    "logs/2025-05-08T20-53-51+00-00_cqa-debug_22k7XP2tE5p9AsKhvjLSoZ.eval",
    "logs/2025-05-08T20-53-51+00-00_cqa-debug_8FKH8hQN7idPyw539ARhLJ.eval",
    "logs/2025-05-08T20-53-51+00-00_cqa-debug_73a5vW3KfqS9khRAF3HkrU.eval",
    "logs/2025-05-08T20-53-51+00-00_cqa-debug_T8WU6aBymfi3cx3UrPVCtr.eval",
    "logs/2025-05-08T20-53-51+00-00_cqa-debug_ZH3AALrpXDz3dtMR82iscL.eval",
    "logs/2025-05-08T20-53-51+00-00_cqa-debug_A8BtuZQ4YBnT7Z4FXdKSaF.eval",
    "logs/2025-05-08T20-53-52+00-00_cqa-debug_VavajxjAZApx9DeQPizbTB.eval",
    "logs/2025-05-08T20-53-52+00-00_cqa-debug_nSCPjAHqbwHoAYbwMT6zzg.eval",
    "logs/2025-05-08T20-53-52+00-00_cqa-debug_RV69iB33YJAXUPACMwtRvF.eval",
]

# logs = [
#     "2025-05-07T18-16-12+00-00_cqa-debug_DENpHcKDEKyKAUJGXziCPj.eval",
#     "2025-05-07T18-16-12+00-00_cqa-debug_akWCziRXx8UMPWCSW3wRr4.eval",
#     "2025-05-07T18-16-12+00-00_cqa-debug_hpuvNrE9TfSw2bGS2nEv8K.eval",
#     "2025-05-07T18-16-12+00-00_cqa-debug_fttEibv9EVV9GB4GFUdzGm.eval",
#     "2025-05-07T18-16-12+00-00_cqa-debug_8XkUnuKVVEEiLYj9sHfao4.eval",
#     "2025-05-07T18-16-12+00-00_cqa-debug_3hqzGUHzFTZo3LihDbEzfN.eval",
#     "2025-05-07T18-16-12+00-00_cqa-debug_fy9arLEoXCUmd5vTdy3JfT.eval",
#     "2025-05-07T18-16-12+00-00_cqa-debug_AaoP7LPLzupX6FHya2qurd.eval",
#     "2025-05-07T18-16-12+00-00_cqa-debug_aT6EqbdagJCdbi2iUhQA8d.eval",
#     "2025-05-07T18-16-12+00-00_cqa-debug_JSgjZWjzEqdEd6i2jssamP.eval",
#     "2025-05-07T18-16-12+00-00_cqa-debug_kEB8S6tMREaGAHTTqRjCm7.eval",
#     "2025-05-07T18-16-12+00-00_cqa-debug_UtLjQ4SC9pia44VXsDfptG.eval",
#     "2025-05-07T18-16-13+00-00_cqa-debug_39xCkjQkWERcFFvhcop5H2.eval",
#     "2025-05-07T18-16-13+00-00_cqa-debug_Ks9KDCTnx5s4T4bsc73WQ6.eval",
#     "2025-05-07T18-16-13+00-00_cqa-debug_VoAAtQJVXPtwdqPJ2imELq.eval",
#     "2025-05-07T18-16-13+00-00_cqa-debug_PDKAyrBwn5BwNEqxpSuLP3.eval",
#     "2025-05-07T18-16-13+00-00_cqa-debug_jKFnhYUTHrJ27qVzVLGnK2.eval",
#     "2025-05-07T18-16-14+00-00_cqa-debug_ZicgeZnBFkNiRB3HDqAaTw.eval",
#     "2025-05-07T18-16-14+00-00_cqa-debug_bLFrgQ8UrZoJZY2GTmsLi5.eval",
#     "2025-05-07T18-16-14+00-00_cqa-debug_bSKJSMJ2kdAMuYytrtVzSp.eval",
# ]


def calc_frac_majority(x):
    count = x.value_counts(sort=True)
    if len(count) == 0:
        return 0
    else:
        return count.iloc[0] / len(x)


def clean_citation(c):
    pattern = r"\((.*?)\)"
    match = re.search(pattern, c)
    if match:
        return match.group(1)
    else:
        return c


def frac_citations_which_never_flipped(x):
    all_non_supporting = set()
    all_supporting = set()
    for _, i in x.iterrows():
        all_non_supporting.update(i["non_supporting"])
        all_supporting.update(i["supporting"])
    if len(all_non_supporting) == 0 and len(all_supporting) == 0:
        return pd.NA
    return 1 - len(all_non_supporting.intersection(all_supporting)) / len(
        all_non_supporting.union(all_supporting)
    )


def main():
    counts = {}
    dfs = []
    for p in logs:
        l = ialog.read_eval_log(p)
        for s in l.samples:
            if "score_citation" not in s.scores:
                continue
            df = build_df_from_prompt_logs(
                PromptLogs(**s.scores["score_citation"].metadata)
            )
            df["question"] = s.metadata["question"]
            df["solver"] = l.plan.steps[0].solver
            df["metric"] = "score_citation"
            temp = [
                a.params["temperature"]
                for a in l.results.scores
                if a.scorer == "score_citation"
            ][0]
            if s.score.value:
                counts[temp] = counts.get(temp, 0) + 1
            df["scorer_model"] = [
                a.params["model"]
                for a in l.results.scores
                if a.scorer == "score_citation"
            ][0]
            df["new_eval"] = [
                a.params["all_at_once"]
                for a in l.results.scores
                if a.scorer == "score_citation"
            ][0]
            df["log_file"] = p
            dfs.append(df)
    df = pd.concat(dfs)

    c = df.explode("claims")
    c["claim"] = c.claims.apply(lambda a: a["text"])
    c["non_supporting"] = c.claims.apply(
        lambda a: [clean_citation(c) for c in a["non_supporting"]]
    )
    c["supporting"] = c.claims.apply(
        lambda a: [clean_citation(c) for c in a["supporting"]]
    )
    c["is_supported"] = c.claims.apply(
        lambda a: a["is_fully_supported"] & (len(a["supporting"]) > 0)
    )

    def _fn(x):
        citations = []
        for _, i in x.iterrows():
            citations.extend(i["supporting"])
            citations.extend(i["non_supporting"])
        return len(citations) > 0

    claims_with_at_least_1_citation = c.groupby(["question", "claim"]).apply(
        _fn, include_groups=False
    )
    c = (
        c.set_index(["question", "claim"])
        .loc[claims_with_at_least_1_citation]
        .reset_index()
    )
    c = c.sort_values(["question", "claim"]).reset_index(drop=True).copy()

    print(
        "frac mostly agree:",
        (
            c.groupby(["temperature", "top_p", "question", "claim"]).is_supported.apply(
                calc_frac_majority
            )
            >= 0.75
        )
        .groupby(["top_p", "temperature"])
        .mean(),
    )
    print(
        "frac full agreement:",
        (
            c.groupby(["temperature", "top_p", "question", "claim"]).is_supported.apply(
                calc_frac_majority
            )
            == 1
        )
        .groupby(["top_p", "temperature"])
        .mean(),
    )
    # not_all_agree = (
    #     c.groupby(["question", "claim"]).is_supported.apply(calc_frac_majority) < 0.6
    # )
    # c_disagree = (
    #     c.set_index(["question", "claim"])
    #     .loc[not_all_agree]
    #     .reset_index()
    #     .sort_values("claim")
    # )

    print(
        "fraction of citations that never flipped:",
        c.groupby(["temperature", "top_p", "question", "claim"])
        .apply(frac_citations_which_never_flipped, include_groups=False)
        .groupby(["top_p", "temperature"])
        .mean(),
    )
    # flipped = (
    #     c.groupby(["question", "claim"]).apply(
    #         frac_citations_which_never_flipped, include_groups=False
    #     )
    #     < 1
    # )
    # c_flipped = (
    #     c.set_index(["question", "claim"])
    #     .loc[flipped]
    #     .reset_index()
    #     .sort_values(["question", "claim"])
    # )


if __name__ == "__main__":
    main()
