#!/usr/bin/env python3
"""Build plotting-ready SR datasets from immutable article artifacts."""
from __future__ import annotations

import argparse
import csv
import json
from copy import deepcopy
from pathlib import Path
from typing import Any, Mapping, Sequence

import numpy as np

from SR_analysis.configs import get_scene
from SR_analysis.stage_2_fit import prepare_scene
from SR_analysis.tools.telemetry_to_csv import export as export_trajectory
from SR_analysis.utils.data_contracts import CAUSAL_ALIGNMENT
from SR_analysis.utils.feature_state import StatefulFeatureBuilder
from SR_analysis.utils.formula_schema import (
    decompose_additive_terms,
    evaluate_formula,
    load_formula,
)
from SR_analysis.utils.g_operator import apply_G_alpha, apply_G_raw12, apply_G_target_force
from SR_analysis.utils.provenance import atomic_write_json, hash_file

TRAINING_SCENES = (
    "karman_re50",
    "karman_re100",
    "karman_re200",
    "karman_re400",
    "illusion_0.75L",
    "illusion_1L",
    "illusion_1.5L",
)
ACTION_NAMES = ("front", "upper", "lower")
FORMULA_ROOTS = {
    "karman": Path("src/SR_analysis/results/runs/article-refit-karman-topology-a-20260718/formulas"),
    "illusion": Path("src/SR_analysis/results/runs/article-refit-illusion-topology-a-20260718/formulas"),
}
DATA_ROOTS = {
    "karman": Path("src/SR_analysis/data/runs/article-joint-data-karman-20260718"),
    "illusion": Path("src/SR_analysis/data/runs/article-joint-data-illusion-20260718"),
}
LONG_RUNS = {
    "karman": Path("src/SR_analysis/results/runs/article2-long-karman-20260720"),
    "illusion": Path("src/SR_analysis/results/runs/article2-long-illusion-20260720"),
}


def _formula_paths(scene_id: str) -> tuple[Path, Path]:
    root = FORMULA_ROOTS[scene_id]
    return root / "joint_front.json", root / "joint_rear_shared_upper.json"


def _term_formula(formula: Mapping[str, Any], expression: str) -> dict[str, Any]:
    item = deepcopy(dict(formula))
    for key in ("source_path", "source_hash", "deployment_expression_hash"):
        item.pop(key, None)
    item["deployment_expression"] = expression
    return item


def _mapped_predictions(
    front: Mapping[str, Any],
    rear: Mapping[str, Any],
    X: np.ndarray,
    GX: np.ndarray,
) -> tuple[np.ndarray, dict[str, np.ndarray]]:
    front_x = np.asarray(evaluate_formula(front, X), dtype=float)
    front_g = np.asarray(evaluate_formula(front, GX), dtype=float)
    upper = np.asarray(evaluate_formula(rear, X), dtype=float)
    lower = -np.asarray(evaluate_formula(rear, GX), dtype=float)
    prediction = np.column_stack((0.5 * (front_x - front_g), upper, lower))
    contributions: dict[str, np.ndarray] = {}
    for index, term in enumerate(decompose_additive_terms(front["deployment_expression"], front["feature_names"])):
        term_formula = _term_formula(front, term)
        contribution = 0.5 * (
            np.asarray(evaluate_formula(term_formula, X), dtype=float)
            - np.asarray(evaluate_formula(term_formula, GX), dtype=float)
        )
        contributions[f"front_t{index}_{term}"] = contribution
    for index, term in enumerate(decompose_additive_terms(rear["deployment_expression"], rear["feature_names"])):
        term_formula = _term_formula(rear, term)
        contributions[f"upper_t{index}_{term}"] = np.asarray(evaluate_formula(term_formula, X), dtype=float)
        contributions[f"lower_t{index}_{term}"] = -np.asarray(evaluate_formula(term_formula, GX), dtype=float)
    return prediction, contributions


def _safe_name(value: str) -> str:
    return "".join(character if character.isalnum() else "_" for character in value).strip("_")


def _write_rows(path: Path, fieldnames: Sequence[str], rows: Sequence[Mapping[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)


def export_offline_predictions(repo_root: Path, output_dir: Path) -> dict[str, Any]:
    rows: list[dict[str, Any]] = []
    term_rows: list[dict[str, Any]] = []
    feature_names: set[str] = set()
    term_names: set[str] = set()
    for scene in TRAINING_SCENES:
        cfg = get_scene(scene)
        scene_id = str(cfg["scene_id"])
        front_path, rear_path = (repo_root / path for path in _formula_paths(scene_id))
        front, rear = load_formula(front_path), load_formula(rear_path)
        prepared = prepare_scene(
            scene,
            repo_root / DATA_ROOTS[scene_id],
            alignment=CAUSAL_ALIGNMENT,
            joint=True,
            feature_set="symmetry",
            feature_profile="actual_plus_target" if scene_id == "illusion" else "actual_only",
        )
        dataset = prepared.dataset
        mirrored = prepared.mirrored_dataset
        prediction, contributions = _mapped_predictions(front, rear, dataset.X, mirrored.X)
        feature_names.update(prepared.feature_names)
        term_names.update(_safe_name(name) for name in contributions)
        target_indices = np.asarray(dataset.metadata["target_indices"], dtype=int)
        for row_index, (state_index, target_index) in enumerate(zip(dataset.kept_indices, target_indices)):
            base: dict[str, Any] = {
                "objective": scene_id,
                "scene": scene,
                "state_step": int(state_index),
                "target_action_step": int(target_index),
                "t_D": float(state_index) * float(cfg["control_dt"]),
            }
            for column, name in enumerate(prepared.feature_names):
                base[f"feature_{name}"] = float(dataset.X[row_index, column])
            for channel, action_name in enumerate(ACTION_NAMES):
                actual = float(dataset.Y[row_index, channel])
                predicted = float(prediction[row_index, channel])
                base[f"actual_alpha_{action_name}"] = actual
                base[f"predicted_alpha_{action_name}"] = predicted
                base[f"residual_alpha_{action_name}"] = actual - predicted
            rows.append(base)
            term_base = {
                "objective": scene_id,
                "scene": scene,
                "source": "ppo_state_offline",
                "step": int(state_index),
                "t_D": float(state_index) * float(cfg["control_dt"]),
            }
            for name, values in contributions.items():
                term_base[f"term_{_safe_name(name)}"] = float(values[row_index])
            for channel, action_name in enumerate(ACTION_NAMES):
                term_base[f"predicted_alpha_{action_name}"] = float(prediction[row_index, channel])
            term_rows.append(term_base)

    offline_fields = [
        "objective", "scene", "state_step", "target_action_step", "t_D",
        *(f"feature_{name}" for name in sorted(feature_names)),
        *(f"actual_alpha_{name}" for name in ACTION_NAMES),
        *(f"predicted_alpha_{name}" for name in ACTION_NAMES),
        *(f"residual_alpha_{name}" for name in ACTION_NAMES),
    ]
    term_fields = [
        "objective", "scene", "source", "step", "t_D",
        *(f"term_{name}" for name in sorted(term_names)),
        *(f"predicted_alpha_{name}" for name in ACTION_NAMES),
    ]
    _write_rows(output_dir / "offline_predictions.csv", offline_fields, rows)
    _write_rows(output_dir / "term_contributions_ppo_states.csv", term_fields, term_rows)
    return {"offline_rows": len(rows), "offline_term_rows": len(term_rows)}


def _telemetry_path(run_root: Path, scene: str) -> Path:
    matches = sorted((run_root / "telemetry").glob(f"{scene.replace('.', 'p')}__pysr__*.npz"))
    if len(matches) != 1:
        raise FileNotFoundError(f"expected one SR telemetry for {scene} below {run_root}, found {matches}")
    return matches[0]


def export_closed_loop_terms(repo_root: Path, output_dir: Path) -> dict[str, Any]:
    rows: list[dict[str, Any]] = []
    all_terms: set[str] = set()
    for scene in TRAINING_SCENES:
        cfg = get_scene(scene)
        scene_id = str(cfg["scene_id"])
        front_path, rear_path = (repo_root / path for path in _formula_paths(scene_id))
        front, rear = load_formula(front_path), load_formula(rear_path)
        telemetry_path = _telemetry_path(repo_root / LONG_RUNS[scene_id], scene)
        with np.load(telemetry_path, allow_pickle=False) as data:
            sensors = np.asarray(data["sensors"], dtype=float)
            forces = np.asarray(data["forces"], dtype=float)
            actions = np.asarray(data["actions_omega"], dtype=float)
            target_forces = np.asarray(data["target_forces"], dtype=float)
        common = dict(
            mu=float(cfg["mu"]), dt_c=float(cfg["control_dt"]), u0=float(cfg["u0"]),
            d=20.0, action_mode="physical", include_mu=True, include_cos_sin=True,
            require_target=scene_id == "illusion",
        )
        builder = StatefulFeatureBuilder(**common)
        g_builder = StatefulFeatureBuilder(**common)
        initial_omega = np.asarray(cfg["policy_init_action"], dtype=float) * float(cfg["u0"])
        for _ in range(2):
            builder.observe_action(initial_omega)
            g_builder.observe_action(apply_G_alpha(initial_omega))
        front_terms = decompose_additive_terms(front["deployment_expression"], front["feature_names"])
        rear_terms = decompose_additive_terms(rear["deployment_expression"], rear["feature_names"])
        for step in range(len(sensors) - 1):
            # Telemetry state ``step`` is post-action ``step``; under the causal
            # contract it predicts the command applied at ``step + 1``.
            builder.observe_action(actions[step])
            g_builder.observe_action(apply_G_alpha(actions[step]))
            raw = np.hstack((sensors[step], forces[step]))
            target = target_forces[step] if len(target_forces) else None
            features, _ = builder.step(raw[:6], raw[6:], target_force=target)
            g_raw = apply_G_raw12(raw)
            g_target = None if target is None else apply_G_target_force(target)
            g_features, _ = g_builder.step(g_raw[:6], g_raw[6:], target_force=g_target)
            row: dict[str, Any] = {
                "objective": scene_id,
                "scene": scene,
                "source": "sr_closed_loop_400",
                "state_step": step,
                "target_action_step": step + 1,
                "t_D": step * float(cfg["control_dt"]),
            }
            predicted = np.zeros(3, dtype=float)
            for index, term in enumerate(front_terms):
                formula = _term_formula(front, term)
                value = 0.5 * (
                    float(evaluate_formula(formula, features))
                    - float(evaluate_formula(formula, g_features))
                )
                name = f"front_t{index}_{_safe_name(term)}"
                row[f"term_{name}"] = value
                all_terms.add(name)
                predicted[0] += value
            for index, term in enumerate(rear_terms):
                formula = _term_formula(rear, term)
                upper = float(evaluate_formula(formula, features))
                lower = -float(evaluate_formula(formula, g_features))
                upper_name = f"upper_t{index}_{_safe_name(term)}"
                lower_name = f"lower_t{index}_{_safe_name(term)}"
                row[f"term_{upper_name}"] = upper
                row[f"term_{lower_name}"] = lower
                all_terms.update((upper_name, lower_name))
                predicted[1] += upper
                predicted[2] += lower
            for channel, action_name in enumerate(ACTION_NAMES):
                row[f"predicted_alpha_{action_name}"] = float(predicted[channel])
                target_action = float(actions[step + 1, channel] / cfg["u0"])
                row[f"target_alpha_{action_name}"] = target_action
                row[f"residual_alpha_{action_name}"] = target_action - float(predicted[channel])
            rows.append(row)
    fields = [
        "objective", "scene", "source", "state_step", "target_action_step", "t_D",
        *(f"term_{name}" for name in sorted(all_terms)),
        *(f"predicted_alpha_{name}" for name in ACTION_NAMES),
        *(f"target_alpha_{name}" for name in ACTION_NAMES),
        *(f"residual_alpha_{name}" for name in ACTION_NAMES),
    ]
    _write_rows(output_dir / "term_contributions_sr_closed_loop_400.csv", fields, rows)
    return {"closed_loop_term_rows": len(rows)}


def export_ablation(repo_root: Path, output_dir: Path) -> dict[str, Any]:
    rows: list[dict[str, Any]] = []
    for run_dir in sorted((repo_root / "src/SR_analysis/results/runs").glob("article-ablation-L2-*-20260718")):
        for path in sorted((run_dir / "validations").glob("*.json")):
            record = json.loads(path.read_text(encoding="utf-8"))
            metric = record.get("metrics", {}).get("legacy_reference_cycle_vs_last_recorded_cycle")
            rows.append({
                "run_id": run_dir.name,
                "variant": run_dir.name.removeprefix("article-ablation-L2-").removesuffix("-20260718"),
                "scene": record.get("scene"),
                "steps": record.get("n_steps"),
                "status": record.get("status"),
                "similarity": None if metric is None else metric.get("similarity"),
                "lag": None if metric is None else metric.get("lag"),
                "termination_reason": record.get("termination_reason"),
                "formula_hash": record.get("formula_hash"),
            })
    fields = ("run_id", "variant", "scene", "steps", "status", "similarity", "lag", "termination_reason", "formula_hash")
    _write_rows(output_dir / "ablation_summary.csv", fields, rows)

    scaling_path = repo_root / "src/SR_analysis/results/runs/article-scaling-L1-v2-20260718/summary.json"
    scaling = json.loads(scaling_path.read_text(encoding="utf-8"))["rows"]
    scaling_rows = []
    for item in scaling:
        row = {key: item.get(key) for key in ("family", "scene", "variant", "failed", "n")}
        for metric in ("action_min", "action_max", "action_rms"):
            values = item.get(metric) or [None, None, None]
            for index, action_name in enumerate(ACTION_NAMES):
                row[f"{metric}_{action_name}"] = values[index]
        scaling_rows.append(row)
    scaling_fields = [
        "family", "scene", "variant", "failed", "n",
        *(f"{metric}_{name}" for metric in ("action_min", "action_max", "action_rms") for name in ACTION_NAMES),
    ]
    _write_rows(output_dir / "scaling_summary.csv", scaling_fields, scaling_rows)
    return {"ablation_rows": len(rows), "scaling_rows": len(scaling_rows)}


def export_steady(repo_root: Path, output_dir: Path) -> dict[str, Any]:
    count = 0
    for magnitude in range(7):
        run = repo_root / f"src/SR_analysis/results/runs/article2-steady-sweep-a{magnitude}-20260720"
        telemetry = run / "telemetry/steady__constant__f-none__si-800.npz"
        export_trajectory(
            telemetry,
            output_dir / "steady_timeseries" / f"alpha_{magnitude}",
            "steady",
            "sr",
            800,
            30,
        )
        count += 1
    return {"steady_timeseries_cases": count}


def build_manifest(repo_root: Path, output_dir: Path, summary: Mapping[str, Any]) -> None:
    artifacts = []
    for path in sorted(
        candidate
        for pattern in ("*.csv", "*.png", "*.pdf", "*.md", "*.json", "*.npz")
        for candidate in output_dir.rglob(pattern)
        if candidate.name != "manifest.json"
    ):
        artifacts.append({"path": str(path.relative_to(repo_root)), "sha256": hash_file(path)})
    atomic_write_json(
        output_dir / "manifest.json",
        {
            "schema_version": "sr-plotting-package-v1",
            "source_policy": "immutable article and Article2 artifacts; no scientific refit",
            "summary": dict(summary),
            "artifacts": artifacts,
        },
    )


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--repo-root", type=Path, default=Path(__file__).resolve().parents[3])
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--replace", action="store_true", help="replace derived files in output-dir")
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    repo_root = args.repo_root.resolve()
    output_dir = args.output_dir if args.output_dir.is_absolute() else repo_root / args.output_dir
    if output_dir.exists() and any(output_dir.iterdir()) and not args.replace:
        raise FileExistsError(f"refusing to overwrite non-empty plotting package: {output_dir}")
    summary: dict[str, Any] = {}
    summary.update(export_offline_predictions(repo_root, output_dir))
    summary.update(export_closed_loop_terms(repo_root, output_dir))
    summary.update(export_ablation(repo_root, output_dir))
    summary.update(export_steady(repo_root, output_dir))
    build_manifest(repo_root, output_dir, summary)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
