#!/usr/bin/env python3
"""Original, standard-library behavior-cloning exercise. Python 3.12+."""
from __future__ import annotations

import argparse
import csv
import hashlib
import json
import math
import random
import sys
import tempfile
from pathlib import Path

DT = 0.2
STEPS = 30
EPISODES = 80
KINDS = ("affine", "knn-1", "knn-7")
ROLES = ("train", "validation", "final-test")
COLUMNS = ("episode_id", "step", "x", "v", "action", "x_next", "v_next")
VERSION = 1
REFERENCE_FILES = ("episodes.csv", "conditions.json", "dataset.json", "splits/train.json", "splits/validation.json", "splits/final-test.json",
                   "model.json", "selection.json", "evaluation.json", "action-errors.csv", "rollout-conditions.json", "rollout-results.csv", "rollouts.csv")


def canonical(value):
    return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False).encode()


def digest(value):
    return hashlib.sha256(canonical(value)).hexdigest()


def file_digest(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def rng(seed, namespace, episode_id):
    # Mix the entire tuple before seeding; nearby episode IDs do not imply nearby states.
    return random.Random(int.from_bytes(hashlib.sha256(canonical([seed, namespace, episode_id])).digest(), "big"))


def clip(value):
    return max(-1.0, min(1.0, value))


def expert(x, v):
    return clip(-2.0 * x - 1.8 * v)


def step(x, v, action, gain=1.0, disturbance=0.0):
    acceleration = gain * action + disturbance
    return x + DT * v + 0.5 * DT * DT * acceleration, v + DT * acceleration


def generate_dataset(seed=7):
    rows, conditions = [], []
    for number in range(EPISODES):
        episode_id = f"episode-{number:03d}"
        draws = rng(seed, "demonstration-reset-v1", episode_id)
        x, v = -0.8 + 1.6 * draws.random(), -0.15 + 0.3 * draws.random()
        conditions.append({"episode_id": episode_id, "x": x, "v": v, "gain": 1.0, "disturbance": 0.0})
        # Fixed-length collection deliberately continues after reaching the goal.
        for k in range(STEPS):
            action = expert(x, v)
            nx, nv = step(x, v, action)
            rows.append(dict(zip(COLUMNS, (episode_id, k, x, v, action, nx, nv))))
            x, v = nx, nv
    return rows, conditions


def grouped(rows):
    result = {}
    for row in rows:
        result.setdefault(row["episode_id"], []).append(row)
    return result


def validate_dataset(rows, conditions):
    if len(rows) != EPISODES * STEPS or len(conditions) != EPISODES:
        raise ValueError("Expected 80 demonstration episodes, each with 30 action rows")
    groups = grouped(rows)
    lookup = {c["episode_id"]: c for c in conditions}
    expected_ids = {f"episode-{n:03d}" for n in range(EPISODES)}
    if set(groups) != expected_ids or set(lookup) != expected_ids:
        raise ValueError("Missing, duplicate, or unknown episode ID")
    for eid, records in groups.items():
        if len(records) != STEPS or [r["step"] for r in records] != list(range(STEPS)):
            raise ValueError("Each episode must contain ordered, unique steps 0 through 29")
        condition = lookup[eid]
        if not all(isinstance(condition[key], (int, float)) and math.isfinite(condition[key]) for key in ("x", "v", "gain", "disturbance")):
            raise ValueError("Condition values must be finite numbers")
        if condition["gain"] != 1.0 or condition["disturbance"] != 0.0:
            raise ValueError("Demonstration conditions must use the nominal plant")
        previous = (condition["x"], condition["v"])
        for row in records:
            if not all(isinstance(row[key], (int, float)) and math.isfinite(row[key]) for key in COLUMNS[2:]):
                raise ValueError("CSV state and action values must be finite numbers")
            current = (row["x"], row["v"])
            if any(abs(a - b) > 1e-12 for a, b in zip(current, previous)):
                raise ValueError("Broken transition continuity within an episode")
            if abs(row["action"] - expert(*current)) > 1e-12:
                raise ValueError("Action label does not match this exercise's declared expert")
            expected = step(*current, row["action"])
            previous = (row["x_next"], row["v_next"])
            if any(abs(a - b) > 1e-12 for a, b in zip(expected, previous)):
                raise ValueError("Next state does not match the declared dynamics")


def make_splits(rows, seed=7):
    ids = sorted(grouped(rows), key=lambda eid: hashlib.sha256(canonical([seed, "split-v1", eid])).digest())
    return {"train": sorted(ids[:48]), "validation": sorted(ids[48:64]), "final-test": sorted(ids[64:])}


def validate_splits(rows, splits):
    if set(splits) != set(ROLES):
        raise ValueError("Need train, validation, and final-test episode manifests")
    all_ids = []
    for role, expected in zip(ROLES, (48, 16, 16)):
        ids = splits[role]
        if not isinstance(ids, list) or len(ids) != expected or len(set(ids)) != expected:
            raise ValueError(f"Invalid {role} manifest: expected {expected} unique episode IDs")
        all_ids.extend(ids)
    if len(set(all_ids)) != EPISODES:
        raise ValueError("Episode overlap across split manifests")
    if set(all_ids) != set(grouped(rows)):
        raise ValueError("Split manifests must partition every dataset episode exactly once")


def selected_rows(rows, ids):
    allowed = set(ids)
    return [row for row in rows if row["episode_id"] in allowed]


def solve(matrix, target):
    """Three-variable Gauss-Jordan solve with partial pivoting."""
    n = len(target)
    augmented = [list(row) + [y] for row, y in zip(matrix, target)]
    for col in range(n):
        pivot = max(range(col, n), key=lambda i: abs(augmented[i][col]))
        augmented[col], augmented[pivot] = augmented[pivot], augmented[col]
        value = augmented[col][col]
        if abs(value) < 1e-14:
            raise ValueError("Singular least-squares system")
        augmented[col] = [v / value for v in augmented[col]]
        for i in range(n):
            if i != col:
                factor = augmented[i][col]
                augmented[i] = [a - factor * b for a, b in zip(augmented[i], augmented[col])]
    return [row[-1] for row in augmented]


def fit_policy(train_rows, kind):
    if not train_rows or kind not in KINDS:
        raise ValueError("Use a nonempty training set and a declared candidate")
    means = [sum(row[key] for row in train_rows) / len(train_rows) for key in ("x", "v")]
    scales = [math.sqrt(sum((row[key] - mean) ** 2 for row in train_rows) / len(train_rows)) for key, mean in zip(("x", "v"), means)]
    if min(scales) <= 1e-12:
        raise ValueError("Training features must have nonzero variance")
    examples = [[(row["x"] - means[0]) / scales[0], (row["v"] - means[1]) / scales[1], row["action"]] for row in train_rows]
    policy = {"kind": kind, "means": means, "scales": scales}
    if kind == "affine":
        # Mean squared loss + 1e-6 times squared slopes; intercept is unpenalized.
        design = [[1.0, x, v] for x, v, _ in examples]
        matrix = [[sum(a[i] * a[j] for a in design) / len(design) + (1e-6 if i == j and i > 0 else 0) for j in range(3)] for i in range(3)]
        target = [sum(a[i] * row[2] for a, row in zip(design, examples)) / len(design) for i in range(3)]
        policy["coefficients"] = solve(matrix, target)
        policy["ridge"] = 1e-6
    else:
        policy["k"] = int(kind.split("-")[1])
        policy["examples"] = examples
    return policy


def predict(policy, x, v):
    if policy["kind"] == "expert":
        return expert(x, v)
    z = [(value - mean) / scale for value, mean, scale in zip((x, v), policy["means"], policy["scales"])]
    if policy["kind"] == "affine":
        a, b, c = policy["coefficients"]
        return clip(a + b * z[0] + c * z[1])
    # Equal distances use training CSV order, a declared deterministic tie rule.
    nearest = sorted(enumerate(policy["examples"]), key=lambda item: ((item[1][0] - z[0]) ** 2 + (item[1][1] - z[1]) ** 2, item[0]))[:policy["k"]]
    return clip(sum(example[2] for _, example in nearest) / len(nearest))


def action_mse(policy, rows):
    if not rows:
        raise ValueError("An action metric needs observations")
    return sum((predict(policy, row["x"], row["v"]) - row["action"]) ** 2 for row in rows) / len(rows)


def build_model(rows, splits):
    validate_splits(rows, splits)
    train = selected_rows(rows, splits["train"])
    validation = selected_rows(rows, splits["validation"])
    candidates = [fit_policy(train, kind) for kind in KINDS]
    scores = {policy["kind"]: action_mse(policy, validation) for policy in candidates}
    # Strict floating-point argmin, candidate declaration order breaks exact ties.
    winner = min(candidates, key=lambda policy: scores[policy["kind"]])
    return {"version": VERSION, "policy": winner, "selected_candidate": winner["kind"], "validation_mse": scores[winner["kind"]], "candidate_scores": scores,
            "train_episode_ids": list(splits["train"]), "validation_episode_ids": list(splits["validation"]),
            "training_records_sha256": digest(train), "validation_records_sha256": digest(validation),
            "selection_rule": "minimum validation action MSE; exact ties use affine, knn-1, knn-7 order"}


def outcome(x, v, streak, k):
    streak = streak + 1 if abs(x) <= 0.1 and abs(v) <= 0.15 else 0
    if abs(x) >= 1.25:
        return "boundary", streak
    if streak >= 3:
        return "success", streak
    return ("timeout" if k == STEPS else None), streak


def rollout(policy, condition):
    x, v, streak, trace = condition["x"], condition["v"], 0, []
    for k in range(STEPS):
        action = predict(policy, x, v)
        nx, nv = step(x, v, action, condition["gain"], condition["disturbance"])
        ending, streak = outcome(nx, nv, streak, k + 1)
        trace.append({"step": k, "time": (k + 1) * DT, "x": x, "v": v, "action": action, "x_next": nx, "v_next": nv, "streak": streak,
                      "terminated": ending in ("success", "boundary"), "truncated": ending == "timeout"})
        x, v = nx, nv
        if ending:
            return {"outcome": ending, "steps": k + 1, "duration": (k + 1) * DT, "final_x": x, "final_v": v, "trace": trace}
    raise AssertionError("Finite budget must produce an outcome")


def summarize_rollouts(results):
    counts = {ending: sum(result["outcome"] == ending for result in results) for ending in ("success", "boundary", "timeout")}
    successes = [result["duration"] for result in results if result["outcome"] == "success"]
    return {"attempted": len(results), **counts, "success_rate": counts["success"] / len(results),
            "mean_successful_seconds": sum(successes) / len(successes) if successes else None}


def evaluate_model(model, rows, splits, conditions, seed=7):
    validate_splits(rows, splits)
    if model["train_episode_ids"] != splits["train"] or model["validation_episode_ids"] != splits["validation"]:
        raise ValueError("Frozen model was fitted on different split IDs")
    for role, key in (("train", "training_records_sha256"), ("validation", "validation_records_sha256")):
        if model[key] != digest(selected_rows(rows, splits[role])):
            raise ValueError("Training or validation records changed after freezing the model")
    policy = model["policy"]
    final = selected_rows(rows, splits["final-test"])
    errors = [{"episode_id": row["episode_id"], "step": row["step"], "expert_action": row["action"], "predicted_action": predict(policy, row["x"], row["v"])} for row in final]
    for record in errors:
        record["squared_error"] = (record["predicted_action"] - record["expert_action"]) ** 2
    offline = {"episodes": len(splits["final-test"]), "actions": len(final), "action_mse": action_mse(policy, final),
               "first_five_actions_mse": action_mse(policy, [row for row in final if row["step"] < 5]),
               "per_episode_mse": {eid: action_mse(policy, records) for eid, records in grouped(final).items()}}
    lookup = {condition["episode_id"]: condition for condition in conditions}
    scenarios, all_results, all_conditions = {}, [], []
    for scenario in ("nominal", "shifted"):
        paired = {"learned": [], "expert": []}
        for eid in splits["final-test"]:
            condition = dict(lookup[eid])
            if scenario == "shifted":
                draws = rng(seed, "shifted-dynamics-v1", eid)
                condition.update(gain=0.6 + 0.8 * draws.random(), disturbance=-0.25 + 0.5 * draws.random())
            condition["scenario"] = scenario
            all_conditions.append(condition)
            for name, controller in (("learned", policy), ("expert", {"kind": "expert"})):
                result = rollout(controller, condition)
                paired[name].append(result)
                all_results.append({"scenario": scenario, "episode_id": eid, "controller": name, **result})
        counts = {"both_success": 0, "learned_only_success": 0, "expert_only_success": 0, "neither_success": 0}
        for a, b in zip(paired["learned"], paired["expert"]):
            success_a, success_b = a["outcome"] == "success", b["outcome"] == "success"
            key = "both_success" if success_a and success_b else "learned_only_success" if success_a else "expert_only_success" if success_b else "neither_success"
            counts[key] += 1
        scenarios[scenario] = {name: summarize_rollouts(results) for name, results in paired.items()}
        scenarios[scenario]["paired"] = counts
    return {"version": VERSION, "seed": seed, "selected_candidate": model["selected_candidate"], "offline": offline, "scenarios": scenarios,
            "conditions": all_conditions, "results": all_results, "action_errors": errors}


def write_json(path, value):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, indent=2, sort_keys=True, allow_nan=False) + "\n", encoding="utf-8")


def read_json(path):
    return json.loads(path.read_text(encoding="utf-8"), parse_constant=lambda value: (_ for _ in ()).throw(ValueError(f"Non-finite JSON number: {value}")))


def write_csv(path, columns, rows):
    with path.open("w", newline="", encoding="utf-8") as stream:
        writer = csv.DictWriter(stream, fieldnames=columns, lineterminator="\n")
        writer.writeheader()
        writer.writerows(rows)


def generate_files(out, seed):
    if out.exists() and any(out.iterdir()):
        raise ValueError("Output directory is not empty; use a fresh directory to preserve the earlier run")
    out.mkdir(parents=True, exist_ok=True)
    rows, conditions = generate_dataset(seed)
    write_csv(out / "episodes.csv", COLUMNS, rows)
    write_json(out / "conditions.json", conditions)
    write_json(out / "dataset.json", {"version": VERSION, "seed": seed, "episodes": EPISODES, "steps": STEPS,
                                    "episodes_sha256": file_digest(out / "episodes.csv"), "conditions_sha256": file_digest(out / "conditions.json")})


def load_dataset(out):
    metadata = read_json(out / "dataset.json")
    if metadata["version"] != VERSION or metadata["episodes_sha256"] != file_digest(out / "episodes.csv") or metadata["conditions_sha256"] != file_digest(out / "conditions.json"):
        raise ValueError("Dataset version or hash mismatch; regenerate in a fresh output directory")
    with (out / "episodes.csv").open(newline="", encoding="utf-8") as stream:
        reader = csv.DictReader(stream)
        if tuple(reader.fieldnames or ()) != COLUMNS:
            raise ValueError("Unexpected demonstration CSV columns")
        rows = [{**row, "step": int(row["step"]), **{key: float(row[key]) for key in COLUMNS[2:]}} for row in reader]
    conditions = read_json(out / "conditions.json")
    validate_dataset(rows, conditions)
    return rows, conditions, metadata


def split_files(out):
    rows, _, metadata = load_dataset(out)
    splits = make_splits(rows, metadata["seed"])
    validate_splits(rows, splits)
    for role in ROLES:
        write_json(out / "splits" / f"{role}.json", {"role": role, "episode_ids": splits[role], "episodes_sha256": metadata["episodes_sha256"]})
    return splits


def load_splits(out, rows, metadata):
    splits = {}
    for role in ROLES:
        record = read_json(out / "splits" / f"{role}.json")
        if record["role"] != role or record["episodes_sha256"] != metadata["episodes_sha256"]:
            raise ValueError("Split manifest role or dataset hash mismatch")
        splits[role] = record["episode_ids"]
    validate_splits(rows, splits)
    return splits


def train_files(out):
    rows, _, metadata = load_dataset(out)
    splits = load_splits(out, rows, metadata)
    model = build_model(rows, splits)
    write_json(out / "model.json", model)
    write_json(out / "selection.json", {"version": VERSION, "selected_candidate": model["selected_candidate"], "candidate_scores": model["candidate_scores"],
                                      "model_sha256": file_digest(out / "model.json"), "dataset_sha256": metadata["episodes_sha256"], "splits_sha256": digest(splits)})
    return model


def evaluate_files(out):
    rows, conditions, metadata = load_dataset(out)
    splits = load_splits(out, rows, metadata)
    selection = read_json(out / "selection.json")
    if selection["model_sha256"] != file_digest(out / "model.json") or selection["dataset_sha256"] != metadata["episodes_sha256"] or selection["splits_sha256"] != digest(splits):
        raise ValueError("Frozen model, dataset, or split manifest changed after selection")
    model = read_json(out / "model.json")
    evaluation = evaluate_model(model, rows, splits, conditions, metadata["seed"])
    write_json(out / "evaluation.json", {key: value for key, value in evaluation.items() if key not in ("conditions", "results", "action_errors")})
    write_json(out / "rollout-conditions.json", evaluation["conditions"])
    write_csv(out / "action-errors.csv", ("episode_id", "step", "expert_action", "predicted_action", "squared_error"), evaluation["action_errors"])
    result_columns = ("scenario", "episode_id", "controller", "outcome", "steps", "duration", "final_x", "final_v")
    write_csv(out / "rollout-results.csv", result_columns, [{key: result[key] for key in result_columns} for result in evaluation["results"]])
    traces = [{"scenario": result["scenario"], "episode_id": result["episode_id"], "controller": result["controller"], **point} for result in evaluation["results"] for point in result["trace"]]
    write_csv(out / "rollouts.csv", ("scenario", "episode_id", "controller", "step", "time", "x", "v", "action", "x_next", "v_next", "streak", "terminated", "truncated"), traces)
    return model, evaluation


def report(model, evaluation):
    lines = ["80 episodes; 2400 actions; split: 48 train / 16 validation / 16 final-test."]
    for kind in KINDS:
        lines.append(f"Validation MSE {kind}: {model['candidate_scores'][kind]:.8f} (m/s^2)^2")
    lines.append(f"Selected and frozen: {model['selected_candidate']}")
    offline = evaluation["offline"]
    lines.append(f"Final nominal action MSE: {offline['action_mse']:.8f}; first five actions: {offline['first_five_actions_mse']:.8f}")
    for scenario in ("nominal", "shifted"):
        for controller in ("learned", "expert"):
            summary = evaluation["scenarios"][scenario][controller]
            mean = "undefined" if summary["mean_successful_seconds"] is None else f"{summary['mean_successful_seconds']:.6f} s"
            lines.append(f"{scenario} {controller}: success={summary['success']}/16, boundary={summary['boundary']}, timeout={summary['timeout']}; mean successful time={mean}")
        paired = evaluation["scenarios"][scenario]["paired"]
        lines.append(f"  paired: both={paired['both_success']}, learned-only={paired['learned_only_success']}, expert-only={paired['expert_only_success']}, neither={paired['neither_success']}")
    lines.append("Synthetic diagnostic batch; final labels never select the model. No hardware or untouched-benchmark claim.")
    return "\n".join(lines) + "\n"


def checks():
    import copy
    x, v = step(-0.5, 0.1, 0.98, 1.0, 0.1)
    assert abs(x + 0.4584) < 1e-12 and abs(v - 0.316) < 1e-12
    rows, conditions = generate_dataset()
    validate_dataset(rows, conditions)
    splits = make_splits(rows)
    validate_splits(rows, splits)
    bad = copy.deepcopy(splits)
    bad["final-test"][0] = bad["train"][0]
    try:
        validate_splits(rows, bad)
    except ValueError:
        pass
    else:
        raise AssertionError("Episode overlap was accepted")
    model = build_model(rows, splits)
    poisoned = copy.deepcopy(rows)
    for row in poisoned:
        if row["episode_id"] in splits["final-test"]:
            row["action"] += 100.0
    assert build_model(poisoned, splits) == model, "Final labels influenced fitting or selection"
    train = selected_rows(rows, splits["train"])
    means = [sum(row[key] for row in train) / len(train) for key in ("x", "v")]
    assert model["policy"]["means"] == means
    assert outcome(0.0, 0.0, 2, STEPS)[0] == "success"
    assert outcome(1.25, 0.0, 2, STEPS)[0] == "boundary"
    assert outcome(0.0, 0.16, 2, STEPS)[0] == "timeout"
    assert summarize_rollouts([{"outcome": "timeout", "duration": 6.0}])["mean_successful_seconds"] is None
    evaluation = evaluate_model(model, rows, splits, conditions)
    for scenario in evaluation["scenarios"].values():
        for name in ("learned", "expert"):
            summary = scenario[name]
            assert summary["success"] + summary["boundary"] + summary["timeout"] == summary["attempted"] == 16
    with tempfile.TemporaryDirectory() as directory:
        out = Path(directory) / "run"
        generate_files(out, 7)
        split_files(out)
        train_files(out)
        disk_model, disk_evaluation = evaluate_files(out)
        assert disk_model == model and disk_evaluation == evaluation
        reference = Path(__file__).with_name("reference")
        if reference.exists():
            for name in REFERENCE_FILES:
                if (out / name).read_bytes() != (reference / name).read_bytes():
                    raise ValueError(f"Default generated artifact differs from supplied reference/{name}")
        # File integrity is a separate guard from the pure selection-invariance check above.
        with (out / "episodes.csv").open("a", encoding="utf-8") as stream:
            stream.write("tampered\n")
        try:
            load_dataset(out)
        except ValueError:
            pass
        else:
            raise AssertionError("Modified dataset was accepted")
    return model, evaluation


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    if sys.version_info < (3, 12):
        parser.error("Python 3.12 or later is required for the supplied numerical reference artifacts")
    parser.add_argument("command", choices=("generate", "split", "train", "evaluate", "check"))
    parser.add_argument("--out", type=Path, help="Artifact directory (default: run); not used by check")
    parser.add_argument("--seed", type=int, help="Dataset seed (default: 7); generate only")
    args = parser.parse_args()
    if args.command == "check" and args.out is not None:
        parser.error("check verifies a fresh default run; use evaluate --out PATH to verify saved artifacts")
    if args.command != "generate" and args.seed is not None:
        parser.error("--seed belongs to generate; later stages read the saved dataset seed")
    args.out = args.out or Path("run")
    args.seed = 7 if args.seed is None else args.seed
    try:
        if args.command == "generate":
            generate_files(args.out, args.seed)
            print("Generated 80 episodes / 2400 actions. Next: split.")
        elif args.command == "split":
            split_files(args.out)
            print("Saved disjoint episode manifests: 48 train / 16 validation / 16 final-test. Next: train.")
        elif args.command == "train":
            model = train_files(args.out)
            print(f"Selected {model['selected_candidate']} using validation action MSE; saved model.json and selection.json. Next: evaluate.")
        elif args.command == "evaluate":
            model, evaluation = evaluate_files(args.out)
            print(report(model, evaluation), end="")
        else:
            model, evaluation = checks()
            print("PASS: dynamics, episode guards, train-only normalization, final-label invariance, stopping rules, metrics, and artifact integrity.")
            print(report(model, evaluation), end="")
    except (ValueError, KeyError, TypeError, OSError, json.JSONDecodeError) as error:
        parser.exit(2, f"Error: {error}\n")


if __name__ == "__main__":
    main()