#!/usr/bin/env python3
"""LockedIn Labs: compare authored routing reports, not model performance.

Python 3.11+, standard library, no account/network. Predictions are fixed teaching
fixtures, not model executions. References are joined only by the evaluator.
"""
import argparse
import json

REFERENCE = ("handoff", "guide") * 5
# Fixed authored output sequences. Neither report generator receives references.
PREDICTIONS = {
    "baseline": {
        "routine": ("handoff", "guide", "handoff", "guide", "handoff", "guide", "handoff", "guide", "guide", "handoff"),
        "complex": ("handoff", "guide", "guide", "handoff", "guide", "handoff", "guide", "handoff", "guide", "handoff"),
    },
    "candidate": {
        "routine": ("handoff", "guide", "handoff", "guide", "handoff", "guide", "guide", "handoff", "handoff", "handoff"),
        "complex": ("handoff", "handoff", "guide", "handoff", "guide", "handoff", "guide", "handoff", "guide", "handoff"),
    },
}


def inputs(counts, prefix=""):
    return [{"case_id": f"{prefix}{kind[0].upper()}{i:02d}",
             "slice": kind, "slot": (i - 1) % 10}
            for kind, count in counts.items() for i in range(1, count + 1)]


def roster_for(cases):
    return [{"case_id": row["case_id"], "slice": row["slice"],
             "reference": REFERENCE[row["slot"]]} for row in cases]


def authored_report(cases, version):
    return [{"case_id": row["case_id"],
             "prediction": PREDICTIONS[version][row["slice"]][row["slot"]]}
            for row in cases]


def index_unique(rows, required):
    if not isinstance(rows, list):
        raise ValueError("Rows must be a list")
    indexed = {}
    for row in rows:
        if (not isinstance(row, dict) or not isinstance(row.get("case_id"), str)
                or not row["case_id"].strip() or any(key not in row for key in required)):
            raise ValueError("Each row needs a nonempty case_id and required fields")
        if row["case_id"] in indexed:
            raise ValueError(f"Duplicate case_id: {row['case_id']}")
        indexed[row["case_id"]] = row
    return indexed


def summarize(roster, reports):
    expected = index_unique(roster, ("slice", "reference"))
    actual = index_unique(reports, ("prediction",))
    if not expected or set(actual) - set(expected):
        raise ValueError("Reports need a nonempty declared roster and no unknown cases")
    if any(row["reference"] not in REFERENCE or row["slice"] not in ("routine", "complex") for row in roster):
        raise ValueError("Use the teaching reference labels and declared slices")
    if any(row["prediction"] not in REFERENCE for row in reports):
        raise ValueError("Use the teaching prediction labels")
    correct = {key: actual[key]["prediction"] == expected[key]["reference"] for key in actual}
    def counts(keys):
        values = [correct[key] for key in keys]
        return {"correct": sum(values), "executed": len(values),
                "accuracy": sum(values) / len(values) if values else None}
    return {"planned": len(expected), **counts(actual),
            "missing_ids": sorted(set(expected) - set(actual)),
            "slices": {kind: counts(key for key in actual if expected[key]["slice"] == kind)
                       for kind in ("routine", "complex")}}, correct


def match_direction(lower, upper):
    """Conclusions true across every completion of this fixed reference roster.

    Bounds are integer candidate-minus-baseline match counts, not statistical
    confidence limits. This says nothing about critical-error costs or release.
    """
    if type(lower) is not int or type(upper) is not int or lower > upper:
        raise ValueError("Use ordered integer match-count bounds")
    if upper < 0:
        return "candidate-fewer-matches"
    if lower > 0:
        return "candidate-more-matches"
    if lower == upper == 0:
        return "same-match-count"
    if lower == 0:
        return "candidate-at-least-as-many"
    if upper == 0:
        return "candidate-at-most-as-many"
    return "direction-undetermined"


def paired(roster, baseline, candidate):
    before, base = summarize(roster, baseline)
    after, cand = summarize(roster, candidate)
    base_predictions = {row["case_id"]: row["prediction"] for row in baseline}
    candidate_predictions = {row["case_id"]: row["prediction"] for row in candidate}
    common = sorted(set(base) & set(cand))
    improved = [key for key in common if not base[key] and cand[key]]
    regressed = [key for key in common if base[key] and not cand[key]]
    count = len(roster)
    # These are completion possibilities over this fixed roster, not confidence
    # limits. No missing prediction is treated as an observed failure or success.
    lower_matches = after["correct"] - before["correct"] - len(before["missing_ids"])
    upper_matches = after["correct"] + len(after["missing_ids"]) - before["correct"]
    low, high = lower_matches / count, upper_matches / count
    return {
        "baseline": before, "candidate": after,
        "matched": len(common), "improved_ids": improved, "regressed_ids": regressed,
        "same_correctness_count": len(common) - len(improved) - len(regressed),
        "matched_baseline_correct": sum(base[key] for key in common),
        "matched_candidate_correct": sum(cand[key] for key in common),
        "matched_accuracy_delta": (len(improved) - len(regressed)) / len(common) if common else None,
        "all_planned_pairs_observed": len(common) == count,
        "possible_full_roster_delta": {"lower": low, "upper": high},
        "possible_full_roster_match_difference": {"lower": lower_matches, "upper": upper_matches},
        "full_roster_match_direction": match_direction(lower_matches, upper_matches),
        "cases": [{"case_id": row["case_id"], "slice": row["slice"],
                   "reference": row["reference"],
                   "baseline_prediction": base_predictions.get(row["case_id"]),
                   "candidate_prediction": candidate_predictions.get(row["case_id"]),
                   "baseline_correct": base.get(row["case_id"]),
                   "candidate_correct": cand.get(row["case_id"])} for row in roster],
    }


def experiments(missing_ids=None):
    old = inputs({"routine": 10, "complex": 90}, "old-")
    new = inputs({"routine": 90, "complex": 10}, "new-")
    old_summary, _ = summarize(roster_for(old), authored_report(old, "baseline"))
    new_summary, _ = summarize(roster_for(new), authored_report(new, "candidate"))
    work = inputs({"routine": 10, "complex": 10})
    roster = roster_for(work)
    base = authored_report(work, "baseline")
    cand = authored_report(work, "candidate")
    missing = ["R07", "R08", "C02"] if missing_ids is None else missing_ids
    if len(set(missing)) != len(missing) or set(missing) - {row["case_id"] for row in work}:
        raise ValueError("Use unique case IDs from the planned paired roster")
    try:
        paired(roster, base, cand + [dict(cand[8])])
    except ValueError as error:
        duplicate = {"status": "refused", "error": str(error), "paired_report": None}
    else:
        raise AssertionError("Duplicate-row control was not refused")
    return [
        {"id": "workload-mix", "baseline": old_summary, "candidate": new_summary,
         "matched": len({row["case_id"] for row in old} & {row["case_id"] for row in new}),
         "paired_report": None,
         "limit": "Different unrandomized rosters and changed slice weights; the headline difference cannot be attributed to the candidate."},
        {"id": "same-cases", **paired(roster, base, cand)},
        {"id": "missing-cases", **paired(roster, base, [row for row in cand if row["case_id"] not in missing])},
        {"id": "duplicate-row", **duplicate},
    ]


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("case", nargs="?", default="all", choices=("all", "workload-mix", "same-cases", "missing-cases", "duplicate-row"))
    parser.add_argument("--missing-id", action="append", help="Replace the missing-cases default; repeat for multiple IDs")
    args = parser.parse_args()
    if args.missing_id is not None and args.case != "missing-cases":
        parser.error("--missing-id applies only to missing-cases")
    try:
        reports = experiments(args.missing_id)
    except ValueError as error:
        parser.error(str(error))
    print(json.dumps({"schema": "lockedin.paired-evaluation/v1",
                      "scope": "Original synthetic routing reports; no model executed, production decision or learner result established. Completion bounds are not confidence intervals.",
                      "experiments": [row for row in reports if args.case in ("all", row["id"])]}, indent=2))
