"""Single fail-fast stage runner for the OID v2 baseline."""
from __future__ import annotations

import argparse
import json
import os
import shutil
import subprocess
import sys
from pathlib import Path
from typing import Any

import numpy as np

from OID_analysis.configs import get_scene
from .alignment import exact_time_pair, phase_conditioned_correction, valid_delayed_pairs
from .analysis import fit_pls_type, fit_pod, fit_regularized_cca, modal_assurance, physical_pls_modes, principal_angles
from .artifacts import atomic_save_npz, atomic_write_json, artifact_record, validate_field_artifact
from .layout import RUNS_ROOT, RunPaths, make_run_id
from CCD_analysis.utils.run_contract import preflight_disk
from .contracts import validate_realization_contract
from .manifest import build_inventory, environment_record, git_state, sha256_file, validate_execution_environment

REPO = Path(__file__).resolve().parents[3]
OID_ROOT = Path(__file__).resolve().parents[1]
ALLOWED_SCENES = {"karman_re100", "illusion_1.0L", "synthetic"}


def _require_scene(scene: str) -> None:
    if scene not in ALLOWED_SCENES:
        raise ValueError(f"OID v2 scene must be one of {sorted(ALLOWED_SCENES)}")


def inventory_paths() -> list[Path]:
    patterns = (
        "src/OID_analysis/data/**/*.npz", "src/OID_analysis/data/**/*.npy",
        "src/OID_analysis/data/**/replay_verify.json", "src/OID_analysis/data/**/*.json",
        "src/OID_analysis/data_li22b/**/*.npz", "src/OID_analysis/data_li22b/**/*.npy",
        "src/OID_analysis/data_li22b/**/*.json",
        "src/CCD_analysis/data/karman/karman_re100/*",
        "src/CCD_analysis/data/illusion/illusion_1.0L/*",
        "src/CCD_analysis/data/target_cylinder/target_cylinder_1.0L/*",
        "src/CCD_analysis/data/pinball/pinball/*",
        "src/CCD_analysis/data/runs/**/replay_verify.json",
        "models/old/d1a3o12_re100.zip",
        "models/250525/d1a3o14_250525_imit_1L_2U_600S.zip",
    )
    result: set[Path] = set()
    for pattern in patterns:
        result.update(path for path in REPO.glob(pattern) if path.is_file())
    return sorted(result)


def run_inventory(output: Path, *, dry_run: bool) -> dict[str, Any]:
    paths = inventory_paths()
    if dry_run:
        return {"dry_run": True, "file_count": len(paths), "would_write": str(output)}
    if output.exists():
        raise FileExistsError(output)
    inventory = build_inventory(paths, base=REPO)
    inventory["classification"] = {
        "src/OID_analysis/data": "legacy-unverified-read-only",
        "src/OID_analysis/data/derived": "legacy-unverified-derived-read-only",
        "src/OID_analysis/data_li22b": "archive-read-only",
        "src/CCD_analysis/data": "canonical-read-only",
    }
    inventory["contract"] = {
        "legacy_roots": ["src/OID_analysis/data", "src/OID_analysis/data_li22b", "src/CCD_analysis/data"],
        "new_runs_root": "src/OID_analysis/runs",
        "overwrite": False,
    }
    atomic_write_json(output, inventory)
    return inventory


def _flatten(ux: np.ndarray, uy: np.ndarray) -> np.ndarray:
    if ux.shape != uy.shape or ux.ndim != 3:
        raise ValueError("fields must be aligned (N, ny, nx)")
    return np.concatenate((ux.reshape(len(ux), -1), uy.reshape(len(uy), -1)), axis=1)


def run_synthetic(paths: RunPaths, config: dict[str, Any]) -> dict[str, Any]:
    rng = np.random.default_rng(int(config.get("seed", 7)))
    n, ny, nx = 120, 5, 7
    steps = np.arange(10, 10 * (n + 1), 10, dtype=np.int64)
    phi = rng.normal(size=(2 * ny * nx, 3))
    phi = np.linalg.qr(phi)[0]
    coefficients = rng.normal(size=(n, 3)) * np.array([3.0, 1.5, 0.5])
    correction = coefficients @ phi.T
    ux = correction[:, :ny * nx].reshape(n, ny, nx).astype(np.float32)
    uy = correction[:, ny * nx:].reshape(n, ny, nx).astype(np.float32)
    timeline = {"lattice_steps": steps, "control_indices": (steps - 1) // 20,
                "control_offsets": (steps - 1) % 20 + 1, "control_interval": np.asarray(20, dtype=np.int64)}
    raw_path = atomic_save_npz(paths.child("raw/synthetic_fields.npz"), ux=ux, uy=uy, **timeline)
    with np.load(raw_path, allow_pickle=False) as archive:
        raw_validation = validate_field_artifact(archive, expected_control_interval=20)
    pair = valid_delayed_pairs(correction, coefficients[:, :1], steps, delay_steps=20)
    pod = fit_pod(pair.left, rank=3)
    pls = fit_pls_type(pod.coefficients, pair.right)
    modes = physical_pls_modes(pod.modes, pls["a_scale"], pls["U"])
    cca = fit_regularized_cca(pod.coefficients, pair.right, 1e-3)
    analysis_path = atomic_save_npz(paths.child("analysis/synthetic_decomposition.npz"),
        pod_mean=pod.mean, pod_modes=pod.modes, pod_S=pod.singular_values,
        pod_energy=pod.energy, U=pls["U"], S=pls["S"], Vt=pls["Vt"],
        a_mean=pls["a_mean"], a_scale=pls["a_scale"], physical_modes=modes,
        cca_S=cca["S"], valid_sample_ids=pair.sample_ids)
    report = {"valid_n": len(pair.sample_ids), "dropped": pair.dropped,
              "energy_sum_retained": float(pod.energy.sum()),
              "sign_invariant_mac": modal_assurance(modes[:, 0], -modes[:, 0]),
              "self_principal_angles_deg": np.degrees(principal_angles(modes[:, :2], modes[:, :2])).tolist()}
    atomic_write_json(paths.child("validation/synthetic_report.json"), report)
    return {"artifacts": [artifact_record(raw_path, raw_validation),
                          artifact_record(analysis_path, {"finite": True}, parent_ids=[sha256_file(raw_path)])],
            "validation": report}


def analyze_field_pair(paths: RunPaths, config: dict[str, Any]) -> dict[str, Any]:
    ctl_path, blk_path = Path(config["controlled_fields"]), Path(config["baseline_fields"])
    observable_path = Path(config["observables"])
    for item in (ctl_path, blk_path, observable_path):
        if not item.is_file():
            raise FileNotFoundError(item)
    with np.load(ctl_path, allow_pickle=False) as ctl, np.load(blk_path, allow_pickle=False) as blk:
        validate_field_artifact(ctl); validate_field_artifact(blk)
        ux_pair = exact_time_pair(ctl["ux"], ctl["lattice_steps"], blk["ux"], blk["lattice_steps"])
        uy_pair = exact_time_pair(ctl["uy"], ctl["lattice_steps"], blk["uy"], blk["lattice_steps"])
        correction = _flatten(ux_pair.left - ux_pair.right, uy_pair.left - uy_pair.right)
        steps = ux_pair.sample_ids
    with np.load(observable_path, allow_pickle=False) as obs_archive:
        if "values" not in obs_archive or "lattice_steps" not in obs_archive:
            raise ValueError("observable artifact requires values and lattice_steps")
        if not np.array_equal(obs_archive["lattice_steps"], steps):
            raise ValueError("observable timeline must exactly match field timeline")
        observables = obs_archive["values"]
    delay = int(config.get("delay_steps", 0))
    pairing = valid_delayed_pairs(correction, observables, steps, delay_steps=delay,
                                  segment_ids=config.get("segment_ids"))
    split = int(config.get("development_count", 0))
    if split < 2 or split >= len(pairing.left) - 1:
        raise ValueError("development_count must leave at least two validation samples")
    ranks = tuple(config.get("ranks", [6, 10, 16]))
    results = {}
    for rank in ranks:
        if rank > min(pairing.left[:split].shape):
            continue
        pod = fit_pod(pairing.left[:split], rank)
        pls = fit_pls_type(pod.coefficients, pairing.right[:split])
        physical = physical_pls_modes(pod.modes, pls["a_scale"], pls["U"])
        validation_a = pod.transform(pairing.left[split:])
        validation_a_std = (validation_a - pls["a_mean"]) / pls["a_scale"]
        validation_y = (pairing.right[split:] - pls["y_mean"]) / pls["y_scale"]
        covariance = validation_a_std.T @ validation_y / len(validation_y)
        retention = float(np.linalg.norm(covariance, "fro") / (np.linalg.norm(pls["cross_covariance"], "fro") + 1e-30))
        cca = {str(eps): fit_regularized_cca(pod.coefficients, pairing.right[:split], eps)["S"] for eps in (1e-4, 1e-3, 1e-2)}
        output = atomic_save_npz(paths.child(f"analysis/rank_{rank}.npz"), pod_mean=pod.mean,
            pod_modes=pod.modes, pod_S=pod.singular_values, pod_energy=pod.energy,
            U=pls["U"], S=pls["S"], Vt=pls["Vt"], a_mean=pls["a_mean"],
            a_scale=pls["a_scale"], y_mean=pls["y_mean"], y_scale=pls["y_scale"],
            physical_modes=physical, valid_sample_ids=pairing.sample_ids,
            cca_eps_1e4=cca["0.0001"], cca_eps_1e3=cca["0.001"], cca_eps_1e2=cca["0.01"])
        results[str(rank)] = {"artifact": str(output.relative_to(paths.partial)), "validation_cross_covariance_retention": retention,
                              "retained_energy": float(pod.energy.sum())}
    report = {"valid_n": len(pairing.left), "dropped": pairing.dropped, "development_n": split,
              "validation_n": len(pairing.left) - split, "delay_steps": delay, "ranks": results}
    atomic_write_json(paths.child("validation/decomposition_report.json"), report)
    return {"validation": report, "inputs": [str(ctl_path), str(blk_path), str(observable_path)]}


def execute_run(scene: str, purpose: str, config: dict[str, Any], *, run_id: str | None, dry_run: bool, device: int | None) -> dict[str, Any]:
    _require_scene(scene)
    validate_execution_environment(device=device, require_gpu=False)
    preflight_disk(RUNS_ROOT, int(config.get("required_bytes", 0)),
                   reserve_bytes=int(config.get("reserve_bytes", 1024**3)))
    if purpose != "synthetic" and config.get("formal", False):
        validate_realization_contract(config.get("realizations", []))
    git = git_state(REPO)
    identifier = run_id or make_run_id(scene, purpose, config, git["commit"])
    paths = RunPaths.create(identifier, dry_run=dry_run)
    if dry_run:
        return {"dry_run": True, "run_id": identifier, "would_create": str(paths.final)}
    try:
        atomic_write_json(paths.child("config.json"), config)
        if purpose == "synthetic":
            payload = run_synthetic(paths, config)
        elif purpose == "decompose":
            payload = analyze_field_pair(paths, config)
        else:
            raise ValueError(f"unsupported CPU purpose: {purpose}")
        manifest = {"schema_version": 1, "run_id": identifier, "scene": scene, "purpose": purpose,
                    "environment": environment_record(repo=REPO, command=sys.argv, device=device), **payload}
        atomic_write_json(paths.child("manifest.json"), manifest)
        atomic_write_json(paths.child("_SUCCESS"), {"run_id": identifier, "manifest_sha256": sha256_file(paths.child("manifest.json"))})
        final = paths.commit()
        return {"run_id": identifier, "path": str(final), "validation": payload.get("validation", {})}
    except BaseException:
        paths.quarantine("failed")
        raise


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="OID v2 immutable stage runner")
    sub = parser.add_subparsers(dest="command", required=True)
    inv = sub.add_parser("inventory")
    inv.add_argument("--output", default=str(OID_ROOT / "legacy_inventory.json"))
    inv.add_argument("--dry-run", action="store_true")
    run = sub.add_parser("run")
    run.add_argument("--scene", required=True)
    run.add_argument("--purpose", choices=("synthetic", "decompose"), required=True)
    run.add_argument("--config")
    run.add_argument("--run-id")
    run.add_argument("--dry-run", action="store_true")
    run.add_argument("--device", type=int)
    args = parser.parse_args(argv)
    if args.command == "inventory":
        result = run_inventory(Path(args.output), dry_run=args.dry_run)
    else:
        config = {} if args.config is None else json.loads(Path(args.config).read_text())
        result = execute_run(args.scene, args.purpose, config, run_id=args.run_id, dry_run=args.dry_run, device=args.device)
    print(json.dumps(result, indent=2, allow_nan=False))
    return 0


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