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 = [
    "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",
]
new_logs = [
    "logs/2025-05-08T17-46-15+00-00_cqa-debug_8RPSCUmTbjhbjx6weNXiv8.eval",
    "logs/2025-05-08T17-46-15+00-00_cqa-debug_XvnBJ2tfGHpiz3XeMnREra.eval",
    "logs/2025-05-08T17-46-15+00-00_cqa-debug_9YNgwa8zvPnhABpvc2LRZQ.eval",
    "logs/2025-05-08T17-46-15+00-00_cqa-debug_BoUPcvKsd9VVkD4ZG7cpn5.eval",
    "logs/2025-05-08T17-46-15+00-00_cqa-debug_kFByxvyy7fM24MRQa3EPjS.eval",
]


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


def build_df_for_logs(paths):
    dfs = []
    for p in paths:
        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]
            df["temperature"] = temp
            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)
    return pd.concat(dfs) if dfs else pd.DataFrame()


def to_claims_df(df):
    c = df.explode("claims")
    c["claim_text"] = c.claims.apply(lambda a: a["text"])
    c["short_claim"] = c.claim_text.str.slice(0, 50)
    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)
    )
    return c


def median_num_claims_per_log_file(claims_df):
    claims_df.groupby(["log_file", "claim_text"]).claim_text.count()


def main():
    orig_paths = ["logs/" + p for p in logs]
    df1 = build_df_for_logs(orig_paths)
    df1 = df1[df1.temperature == 1].reset_index(drop=True).copy()
    df2 = build_df_for_logs(new_logs)

    claims_before = to_claims_df(df1)
    claims_after = to_claims_df(df2)

    claim_counts_before = claims_before.groupby(
        ["question", "short_claim"]
    ).short_claim.count()
    claims_which_occurred_5_times_before = (
        claims_before.set_index(["question", "short_claim"])
        .loc[claim_counts_before[claim_counts_before >= 5].index]
        .reset_index()
    )

    claim_counts_after = claims_after.groupby(
        ["question", "short_claim"]
    ).short_claim.count()
    claims_which_occurred_3_times_after = (
        claims_after.set_index(["question", "short_claim"])
        .loc[claim_counts_after[claim_counts_after >= 3].index]
        .reset_index()
    )

    joined = (
        claims_which_occurred_5_times_before.drop_duplicates(
            ["question", "short_claim"]
        )
        .set_index(["question", "short_claim"])
        .join(
            claims_which_occurred_3_times_after.drop_duplicates(
                ["question", "short_claim"]
            ).set_index(["question", "short_claim"]),
            lsuffix="_before",
            rsuffix="_after",
            how="left",
        )
    )

    just_missing = joined[joined.claim_text_after.isna()]
    just_missing.reset_index("short_claim", drop=True).claim_text_before.to_csv(
        "just_missing.csv"
    )


if __name__ == "__main__":
    main()
