"""Create training-seed-aware summaries for priority-reveal diagnostics."""

from __future__ import annotations

import argparse
import csv
import json
import math
from collections import defaultdict
from pathlib import Path

import numpy as np


METRICS = (
    "mean_event_task_probability",
    "immediate_target_assignments",
    "immediate_policy_switches",
    "departures_from_fragile_sources",
    "executed_target_assignments",
    "executed_switches",
    "executed_non_target_switches",
    "executed_fragile_departures",
    "success",
    "task_completion_ratio",
    "priority_weighted_completion",
    "event_task_completed",
    "event_task_response_delay",
    "remaining_demand",
    "episode_length",
    "team_reward_after_event",
)


def _read(path):
    with Path(path).open(encoding="utf-8", newline="") as stream:
        return list(csv.DictReader(stream))


def _write(path, rows):
    path = Path(path)
    if not rows:
        path.write_text("", encoding="utf-8")
        return
    with path.open("w", encoding="utf-8", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)


def _number(row, key):
    value = row.get(key, "")
    return float(value) if value not in (None, "") else math.nan


def _paired_rows(rows):
    index = {
        (r["model"], int(r["scenario_id"]), r["condition"]): r for r in rows
    }
    result = []
    contrasts = (
        ("priority_high_minus_low", "high", "low"),
        (
            "react_minus_keep_current",
            "high",
            "high_keep_current_at_event",
        ),
    )
    for (model, scenario, condition), high in sorted(index.items()):
        if condition != "high":
            continue
        for label, left_name, right_name in contrasts:
            left = index.get((model, scenario, left_name))
            right = index.get((model, scenario, right_name))
            if right is None:
                continue
            item = {
                "model": model,
                "model_role": left["model_role"],
                "training_condition": left["training_condition"],
                "ablation": left["ablation"],
                "train_seed": int(left["train_seed"]),
                "scenario_id": scenario,
                "contrast": label,
                "pre_event_state_hash_equal": int(
                    left["pre_event_state_hash"] == right["pre_event_state_hash"]
                ),
            }
            for metric in METRICS:
                if metric.startswith("executed_"):
                    source = {
                        "executed_target_assignments": "immediate_target_assignments",
                        "executed_switches": "immediate_policy_switches",
                        "executed_non_target_switches": None,
                        "executed_fragile_departures": "departures_from_fragile_sources",
                    }[metric]
                    if metric == "executed_non_target_switches":
                        left_value = _number(
                            left, "immediate_policy_switches"
                        ) - _number(left, "immediate_target_assignments")
                        right_value = (
                            0.0
                            if label == "react_minus_keep_current"
                            else _number(right, "immediate_policy_switches")
                            - _number(right, "immediate_target_assignments")
                        )
                    else:
                        left_value = _number(left, source)
                        right_value = (
                            0.0
                            if label == "react_minus_keep_current"
                            else _number(right, source)
                        )
                    item[metric] = left_value - right_value
                else:
                    item[metric] = _number(left, metric) - _number(right, metric)
            result.append(item)
    return result


def _aggregate_training_seeds(paired):
    by_seed = defaultdict(list)
    for row in paired:
        key = (
            row["training_condition"],
            row["model_role"],
            row["ablation"],
            row["train_seed"],
            row["contrast"],
        )
        by_seed[key].append(row)

    seed_rows = []
    for key, items in sorted(by_seed.items()):
        condition, role, ablation, seed, contrast = key
        record = {
            "training_condition": condition,
            "model_role": role,
            "ablation": ablation,
            "train_seed": seed,
            "contrast": contrast,
            "eval_scenarios": len(items),
            "paired_state_check": int(
                all(int(x["pre_event_state_hash_equal"]) == 1 for x in items)
            ),
        }
        for metric in METRICS:
            values = np.asarray([float(x[metric]) for x in items], dtype=float)
            record[metric] = float(np.nanmean(values))
        seed_rows.append(record)

    grouped = defaultdict(list)
    for row in seed_rows:
        key = (
            row["training_condition"],
            row["model_role"],
            row["ablation"],
            row["contrast"],
        )
        grouped[key].append(row)

    aggregate = []
    for key, items in sorted(grouped.items()):
        condition, role, ablation, contrast = key
        record = {
            "training_condition": condition,
            "model_role": role,
            "ablation": ablation,
            "contrast": contrast,
            "training_seeds": len(items),
            "eval_scenarios_per_seed": min(int(x["eval_scenarios"]) for x in items),
            "paired_state_check": int(all(x["paired_state_check"] for x in items)),
        }
        for metric in METRICS:
            values = np.asarray([float(x[metric]) for x in items], dtype=float)
            mean = float(np.mean(values))
            sd = float(np.std(values, ddof=1)) if len(values) > 1 else math.nan
            # With three training seeds this is descriptive uncertainty.  Keep
            # the individual seed rows for every paper-facing interpretation.
            half = 4.303 * sd / math.sqrt(len(values)) if len(values) == 3 else math.nan
            record[f"{metric}_mean"] = mean
            record[f"{metric}_training_seed_sd"] = sd
            record[f"{metric}_t95_low"] = mean - half
            record[f"{metric}_t95_high"] = mean + half
        aggregate.append(record)
    return seed_rows, aggregate


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("output_dir")
    args = parser.parse_args()
    output = Path(args.output_dir).resolve()
    rows = _read(output / "all_results.csv")
    paired = _paired_rows(rows)
    seed_rows, aggregate = _aggregate_training_seeds(paired)
    _write(output / "paired_effects_checked.csv", paired)
    _write(output / "per_training_seed_summary.csv", seed_rows)
    _write(output / "across_training_seed_summary.csv", aggregate)
    audit = {
        "raw_rows": len(rows),
        "paired_rows": len(paired),
        "models": len({r["model"] for r in rows}),
        "all_within_model_pairs_share_pre_event_state": all(
            int(r["pre_event_state_hash_equal"]) == 1 for r in paired
        ),
        "uncertainty_unit": "independent training seed",
        "eval_scenarios_are_paired_repeated_measurements": True,
    }
    (output / "analysis_audit.json").write_text(
        json.dumps(audit, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )


if __name__ == "__main__":
    main()