"""Phase 3 label audit.

Samples questions, hides the generated label, asks the user to label each one,
then reports agreement between human and generated labels.

Usage:
    python scripts/audit_labels.py
    python scripts/audit_labels.py --n 50 --source data/raw/questions.jsonl
"""

import argparse
import json
import os
import random

ORDERED_LABELS = ["Surface", "Mechanistic", "Critical"]
INPUT_MAP = {
    "s": "Surface", "surface": "Surface",
    "m": "Mechanistic", "mechanistic": "Mechanistic",
    "c": "Critical", "critical": "Critical",
}

DEFINITIONS = {
    "Surface": "recall a fact/definition, or explain/interpret a single concept",
    "Mechanistic": "apply a method, or analyse/compare/trace mechanism",
    "Critical": "evaluate, critique, weigh trade-offs, or design something new",
}


def load_records(source):
    """Load all records from a JSONL file."""
    records = []
    with open(source, "r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if line:
                records.append(json.loads(line))
    return records


def stratified_sample(records, n, seed):
    """Sample n records with roughly equal counts per class."""
    random.seed(seed)
    by_class = {label: [] for label in ORDERED_LABELS}
    for record in records:
        label = record.get("bloom_class")
        if label in by_class:
            by_class[label].append(record)

    per_class = max(1, n // len(ORDERED_LABELS))
    sample = []
    for label in ORDERED_LABELS:
        pool = by_class[label]
        random.shuffle(pool)
        sample.extend(pool[:per_class])
    random.shuffle(sample)
    return sample


def prompt_one(question, index, total):
    """Show a question and collect a blind label, or None to quit."""
    print(f"\n[{index}/{total}] {question}")
    print("  Which cognitive level?")
    for label in ORDERED_LABELS:
        print(f"    {label[0].lower()} = {label:12s} ({DEFINITIONS[label]})")
    while True:
        choice = input("  your label [s/m/c, or q to stop]: ").strip().lower()
        if choice == "q":
            return None
        if choice in INPUT_MAP:
            return INPUT_MAP[choice]
        print("  please enter s, m, c, or q")


def summarise(annotations):
    """Print and return agreement stats."""
    if not annotations:
        print("\nNo items labelled.")
        return None

    total = len(annotations)
    agree = sum(1 for a in annotations if a["human_label"] == a["true_label"])
    overall = agree / total

    print("\n" + "=" * 46)
    print(f"Label fidelity audit  (n = {total})")
    print(f"  overall agreement : {overall:.1%}")
    print(f"\n  {'class':12s} {'agree':>6s} {'n':>4s}")
    per_class = {}
    for label in ORDERED_LABELS:
        subset = [a for a in annotations if a["true_label"] == label]
        if subset:
            hits = sum(1 for a in subset if a["human_label"] == label)
            rate = hits / len(subset)
            per_class[label] = {"agreement": rate, "n": len(subset)}
            print(f"  {label:12s} {rate:6.1%} {len(subset):4d}")
        else:
            per_class[label] = {"agreement": None, "n": 0}
    print("=" * 46)
    return {"overall_agreement": overall, "per_class": per_class, "n": total}


def main():
    """Run the interactive audit and save results."""
    parser = argparse.ArgumentParser(description="Label-fidelity audit")
    parser.add_argument("--source", default="data/raw/questions.jsonl",
                        help="JSONL to sample from")
    parser.add_argument("--n", type=int, default=50, help="sample size")
    parser.add_argument("--seed", type=int, default=7, help="sampling seed")
    parser.add_argument("--output", default="data/outputs/label_fidelity.json")
    args = parser.parse_args()

    if not os.path.exists(args.source):
        fallback = "data/processed/train.jsonl"
        print(f"{args.source} not found, falling back to {fallback}")
        args.source = fallback

    records = load_records(args.source)
    sample = stratified_sample(records, args.n, args.seed)
    print(f"Sampled {len(sample)} questions from {args.source}.")
    print("Label each by cognitive intent. The generated label is hidden.\n")

    annotations = []
    for i, record in enumerate(sample, start=1):
        human = prompt_one(record["question"], i, len(sample))
        if human is None:
            print("Stopping early.")
            break
        annotations.append({
            "question": record["question"],
            "true_label": record["bloom_class"],
            "human_label": human,
            "agree": human == record["bloom_class"],
        })

    summary = summarise(annotations)
    if summary is not None:
        os.makedirs(os.path.dirname(args.output), exist_ok=True)
        with open(args.output, "w", encoding="utf-8") as handle:
            json.dump({"summary": summary, "annotations": annotations},
                      handle, indent=2, ensure_ascii=False)
        print(f"\nSaved -> {args.output}")


if __name__ == "__main__":
    main()