import argparse
import pandas as pd

parser = argparse.ArgumentParser()
parser.add_argument("--input", default="afghanistan_local_knowledge_scored.csv")
parser.add_argument("--output", default="evaluation_summary.csv")
args = parser.parse_args()

df = pd.read_csv(args.input)
required = {
    "condition", "answerability", "human_score",
    "unsupported_specific_claim"
}
missing = required - set(df.columns)
if missing:
    raise ValueError(f"Missing required columns: {sorted(missing)}")

conditions = ["closed_book", "uncertainty_instruction", "source_supplied"]
rows = []

for condition in conditions:
    part = df[df["condition"] == condition]
    a = part[part["answerability"] == "answerable_from_source"]
    u = part[part["answerability"] == "not_in_source"]

    rows.append({
        "condition": condition,
        "total_responses": len(part),
        "answerable_questions": len(a),
        "exact_correct": int((a["human_score"] == "correct").sum()),
        "partially_correct": int((a["human_score"] == "partially_correct").sum()),
        "appropriate_abstentions": int(
            (u["human_score"] == "appropriate_abstention").sum()
        ),
        "unsupported_or_contradicted_flags": int(
            (part["unsupported_specific_claim"].astype(str).str.lower() == "yes").sum()
        )
    })

result = pd.DataFrame(rows)
result.to_csv(args.output, index=False)
print(result.to_string(index=False))
print(f"\\nSaved summary to {args.output}")