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

old_name = "2025-05-12T17-20-47+00-00_cqa_NsEj8qE2HpgdeDVqxHYDSC.eval"
name = "2025-05-12T19-33-41+00-00_cqa_njrg9NnaKvPuDCzfbsUGSq.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["initial_prompt"]
            df["solver"] = l.plan.steps[0].solver
            df["metric"] = "score_citation"
            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["supporting_citations"] = c.apply(
        lambda row: [
            a for a in row["citations"] if clean_citation(a["id"]) in row["supporting"]
        ],
        axis=1,
    )
    c["non_supporting_citations"] = c.apply(
        lambda row: [
            a
            for a in row["citations"]
            if clean_citation(a["id"]) in row["non_supporting"]
        ],
        axis=1,
    )
    c["is_supported"] = c.claims.apply(
        lambda a: a["is_fully_supported"] & (len(a["supporting"]) > 0)
    )
    return c


def main():
    df = build_df_for_logs(["logs/" + name])
    full_df = df
    df = df[df.half_credit.apply(len) > 0]
    claims = to_claims_df(df)
    is_half_credit_claim = claims.apply(
        lambda row: len(
            set(row["supporting"]).intersection(
                set(clean_citation(c) for c in row["half_credit"])
            )
        )
        > 0,
        axis=1,
    )
    cites_title_claim = claims.apply(
        lambda row: len(
            set(row["supporting"] + row["non_supporting"]).intersection(
                set(clean_citation(c) for c in row["half_credit"])
            )
        )
        > 0,
        axis=1,
    )
    print(
        "Fraction of all claims which cite a paper with only a title:",
        cites_title_claim.sum() / len(full_df.explode("claims")),
    )
    # 16%
    print(
        "Fraction of all claims which get half credit:",
        is_half_credit_claim.sum() / len(full_df.explode("claims")),
    )
    # 14%
    claims["str_supporting"] = claims.supporting.apply(lambda a: "; ".join(a))
    claims["str_supporting_citations"] = claims.supporting_citations.apply(
        lambda a: "\n".join(["{}: {}".format(b["id"], b["snippets"]) for b in a])
    )
    claims["str_non_supporting"] = claims.non_supporting.apply(lambda a: "; ".join(a))
    claims["str_non_supporting_citations"] = claims.non_supporting_citations.apply(
        lambda a: "\n".join(["{}: {}".format(b["id"], b["snippets"]) for b in a])
    )
    claims[is_half_credit_claim][
        ["question", "claim_text", "str_supporting", "str_supporting_citations"]
    ].to_csv("half_credit.csv", index=False)
    claims[(~is_half_credit_claim) & (cites_title_claim)][
        ["question", "claim_text", "str_non_supporting", "str_non_supporting_citations"]
    ].to_csv("possibly_half_credit_but_not.csv", index=False)

    # compare with old
    df = build_df_for_logs(["logs/" + name])
    claims = to_claims_df(df)
    df_old = build_df_for_logs(["logs/" + old_name])
    old_claims = to_claims_df(df_old)

    joined = old_claims.set_index(["question", "short_claim"]).join(
        claims.set_index(["question", "short_claim"]),
        lsuffix="_old",
        rsuffix="_new",
        how="inner",
    )

    print((joined.is_supported_old == joined.is_supported_new).mean())
    # agreement 90%


if __name__ == "__main__":
    main()
