"""Rebuild selected, provenance-linked handoff evidence from local experiment files."""
import csv
import hashlib
import json
import re
from collections import defaultdict
from pathlib import Path
from statistics import mean, stdev
import yaml

ROOT = Path(__file__).resolve().parents[2]
OUT = Path(__file__).resolve().parent / "evidence"
OUT.mkdir(exist_ok=True)
sources = []


def load(path):
    path = ROOT / path
    data = path.read_bytes()
    sources.append({"path": str(path.relative_to(ROOT)), "sha256": hashlib.sha256(data).hexdigest()})
    return json.loads(data)


def clean(value):
    if isinstance(value, dict):
        return {k: clean(v) for k, v in value.items() if k not in {"pid", "log", "config_file", "models_file", "device"}}
    if isinstance(value, list):
        return [clean(v) for v in value]
    if isinstance(value, str):
        value = value.replace(str(ROOT) + "/", "")
        if value.startswith("/home/") or value.startswith("/tmp/"):
            return Path(value).name
    return value


def save(name, value):
    (OUT / name).write_text(json.dumps(clean(value), ensure_ascii=False, indent=2) + "\n")


def csv_save(name, rows):
    keys = list(dict.fromkeys(k for r in rows for k in r))
    with (OUT / name).open("w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=keys)
        w.writeheader()
        w.writerows(clean(rows))


metrics = ["success", "task_completion_ratio", "sent_percentage", "discovered_ratio", "episode_length", "episode_reward", "reconfiguration_decisions"]
formal = []
for prefix, label in [("20260818-0040", "formal-v20"), ("20260820-1505", "geometry"),
                      ("20260822-1223", "v33aa"), ("20260903-0036", "schema-pilot"),
                      ("20260903-1846", "schema-bridge"), ("20260904-1243", "schema-scale"),
                      ("20260905-1200", "padding"), ("20260906-2158", "v34")]:
    r = next((ROOT / "runs").glob(prefix + "*"))
    q = next(p for p in r.glob("staged_evaluation_gpu*") if (p / "final_test").exists())
    for p in sorted(q.glob("final_test/*/*/*/summary.json")):
        d = load(p.relative_to(ROOT))
        job = p.parent.parent.name
        sm = re.search(r"-s(\d+)(?=_|$)", job)
        family = re.sub(r"-s\d+(?=_|$)", "", job).split("_")[0]
        formal.append(dict(experiment=label, method=family, train_seed=int(sm[1]) if sm else 1,
                           selector=p.parent.name, episodes=d["num_episodes"],
                           source=str(p.relative_to(ROOT)),
                           **{k: d[k] for k in metrics if k in d}))
csv_save("formal_by_seed.csv", formal)


def aggregate(rows, group_keys):
    groups = defaultdict(list)
    for r in rows:
        groups[tuple(r[k] for k in group_keys)].append(r)
    result = []
    for key, ds in sorted(groups.items()):
        row = dict(zip(group_keys, key), n=len(ds))
        for k in metrics:
            if all(k in d for d in ds):
                row[k] = mean(d[k] for d in ds)
                row[k + "_train_sd"] = stdev(d[k] for d in ds) if len(ds) > 1 else None
        result.append(row)
    return result


ag = aggregate(formal, ["experiment", "method", "selector"])
csv_save("formal_aggregate.csv", ag)
save("formal_aggregate.json", ag)
common3 = [r for r in formal if r["train_seed"] <= 3 and r["experiment"] in ("formal-v20", "geometry", "v33aa", "v34")]
save("common_three_seeds.json", aggregate(common3, ["experiment", "method", "selector"]))

old = []
for prefix, label in [("20260725-2214", "detached"), ("20260729-1311", "coverage-rnn"), ("20260813-1841", "qmix-vdn")]:
    p = next((ROOT / "runs").glob(prefix + "*/eval_results/*/aggregate_all.json"))
    d = load(p.relative_to(ROOT))
    for job, ds in d.items():
        for x in ds:
            sel = x["name"].split("_")[-1]
            if sel not in {"best", "last", "p90"}:
                continue
            old.append(dict(experiment=label, method=re.sub(r"-s\d+(?=_|$)", "", job).split("_")[0],
                            train_seed=int(re.search(r"-s(\d+)", job)[1]), selector=sel,
                            source=str(p.relative_to(ROOT)), **{k: x[k] for k in metrics if k in x}))
csv_save("legacy_by_seed.csv", old)
save("legacy_aggregate.json", aggregate(old, ["experiment", "method", "selector"]))

cap = load("runs/20260812-2342_delivcoal_capcoal-midscale-legacy-completed-independent-recovery/aggregate_all.json")
save("capcoal_rows.json", cap)
save("capcoal_aggregate.json", load("runs/20260812-2342_delivcoal_capcoal-midscale-legacy-completed-independent-recovery/aggregate_by_condition.json"))
pcfa = load("runs/20260828-1330_delivcoal_20260828-pcfa-delivcoal-reproduction/pcfa-delivcoal_heldout-128/pcfa_results.json")
save("pcfa_results.json", pcfa)

aa = next((ROOT / "runs/20260822-1223_delivcoal_hatt-v33aa-focused/causal_aa_validation_best").glob("results/*/summary.json"))
save("aa_frozen_interventions.json", load(aa.relative_to(ROOT)))

external = []
for prefix in ("20260830-2132", "20260831-1333_cfm-long", "20260831-1333_mrt2a-cap-lr", "20260831-1333_mrt2a-cap-entropy"):
    r = next((ROOT / "runs").glob(prefix + "*"))
    for p in sorted(r.glob("evaluation/final_test/*/*/*/summary.json")):
        d = load(p.relative_to(ROOT))
        external.append(dict(run=r.name, selector=p.parts[-4], source=str(p.relative_to(ROOT)), data=d))
save("external_adaptations.json", external)
for name, rel in [
    ("cfm_official.json", "runs/20260828_heteromrta-cfm-official-eval-50/summary.json"),
    ("cfm_small_learning.json", "runs/20260828-1435_small-baseline-convergence/heteromrta_cfm_small/heldout_eval/summary.json"),
    ("mrt2a_small_learning.json", "runs/20260828-1435_small-baseline-convergence/marl_mrt2a_small/heldout_eval/summary.json"),
]:
    save(name, load(rel))

# Publish only research fields, not machine/runtime configuration.
fields = "env_name scenario_name map_size num_uav num_ugv target_num max_steps episode_time_limit cap_dim coalition_task_feature_mode explore_action_mode num_explore_actions explore_step_size explore_random_jitter explore_coverage_features explore_coverage_grid_resolution explore_observer_relative_geometry delivery_tasks_initially_discovered mask_explore_actions idle_time_cost_mode reconfiguration_mode trainer_algorithm_name graph_model use_gat hidden_size layer_N hatt_hidden_dim hatt_num_heads hatt_shared_task_obs hatt_candidate_aa_mode hatt_aa_task_context hatt_use_aa_attention hatt_use_coalition_context use_recurrent_policy max_train_episodes num_env_steps n_rollout_threads episode_length ppo_epoch num_mini_batch lr critic_lr entropy_coef clip_param gamma gae_lambda use_gae use_valuenorm use_linear_lr_decay ppo_target_kl ppo_max_clip_fraction max_grad_norm seed".split()
settings = []
seen = set()
for row in formal:
    if row["selector"] != "validation-best":
        continue
    raw = load(row["source"])
    config_path = Path(raw["checkpoint"]).parent.parent / "config.yaml"
    if config_path in seen:
        continue
    seen.add(config_path)
    blob = config_path.read_bytes()
    cfg = yaml.safe_load(blob)
    sources.append(dict(path=str(config_path.relative_to(ROOT)), sha256=hashlib.sha256(blob).hexdigest()))
    settings.append(dict(experiment=row["experiment"], method=row["method"], train_seed=row["train_seed"],
                         source=str(config_path.relative_to(ROOT)), **{k: cfg.get(k) for k in fields}))
csv_save("research_settings.csv", settings)

for rel in ["analysis/priority_reconfiguration_20260905/across_training_seed_summary.csv",
            "analysis/priority_reconfiguration_20260905/per_training_seed_summary.csv",
            "analysis/gymrt2a_reproduction_20260724/official_summary.csv",
            "analysis/gymrt2a_hatt_eval_20260725/heldout_summary.csv",
            "runs/20260808-1622_delivcoal_coverage-slot-permutation/deliver/summary.csv",
            "runs/20260808-1622_delivcoal_coverage-slot-permutation/explore/summary.csv"]:
    p = ROOT / rel
    with p.open() as f:
        rows = list(csv.DictReader(f))
    sources.append(dict(path=rel, sha256=hashlib.sha256(p.read_bytes()).hexdigest()))
    name = p.parent.name + "_" + p.name
    csv_save(name, rows)

# Final-window averages are computed per training seed before averaging seeds.
window = []
wr = ROOT / "runs/20260821-1045_delivcoal_final-window-evaluation"
for p in sorted(wr.glob("*/*/results/*/summary.json")):
    d = load(p.relative_to(ROOT))
    ds = d["models"]
    job = p.parents[2].name
    window.append(dict(method=re.sub(r"-s\d+(?=_|$)", "", job).split("_")[0],
                       train_seed=int(re.search(r"-s(\d+)", job)[1]),
                       checkpoints=len(ds), eval_episodes_per_checkpoint=ds[0]["num_episodes"],
                       source=str(p.relative_to(ROOT)),
                       **{k: mean(x[k] for x in ds) for k in metrics if all(k in x for x in ds)}))
csv_save("final_window_by_seed.csv", window)
save("final_window_aggregate.json", aggregate(window, ["method"]))
for filename in ['periodic_validation.csv', 'selection_and_windows.csv', 'final_by_seed.csv', 'paired_scenarios.csv']:
    p = ROOT / 'analysis/hatt_v34_taskctx_20260907' / filename
    with p.open() as f:
        rows = list(csv.DictReader(f))
    csv_save('v34_' + filename, rows)
    sources.append(dict(path=str(p.relative_to(ROOT)), sha256=hashlib.sha256(p.read_bytes()).hexdigest()))
save("sources.json", sorted({s["path"]: s for s in sources}.values(), key=lambda s: s["path"]))
print(json.dumps(dict(formal_rows=len(formal), legacy_rows=len(old), window_jobs=len(window),
                      source_files=len({s['path'] for s in sources})), ensure_ascii=False))