#!/usr/bin/env python3
"""Run the formal CPU-only Phase 3--5 matrix with no-clobber publication."""
from __future__ import annotations

import argparse
import hashlib
import json
import os
from pathlib import Path
import platform
import subprocess
import sys

THREAD_ENV = {"OPENBLAS_NUM_THREADS": "2", "OMP_NUM_THREADS": "2", "MKL_NUM_THREADS": "2", "NUMEXPR_NUM_THREADS": "2"}
for key, value in THREAD_ENV.items():
    os.environ[key] = value

ROOT = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(ROOT / "src"))

import numpy as np
from threadpoolctl import threadpool_info

from CCD_analysis.utils.run_contract import deterministic_json_bytes, sha256_file
from OID_analysis.v2.oid import ANALYSIS_SCHEMA, _hash, run_phase_3_5

SCENES = ("karman_re100", "illusion_1.0L")
OBSERVABLES = ("delta_forces", "delta_sensors")
OBSERVABLE_CHANNELS = {"delta_forces": 6, "delta_sensors": 6}
OBSERVABLE_COORDINATE_UNITS = {
    "delta_forces": ("LBM force units",) * 6,
    "delta_sensors": ("velocity normalized by U0",) * 6,
}
RANKS = (4, 8, 12)
LAGS = (0, 1, 2)
SURROGATES = 999
NEGATIVE_CONTROL_MARGIN = 0.01
RCOND = 1e-12
D, U0 = 20.0, 0.01


def _seed(scene: str, observable: str) -> int:
    digest = hashlib.sha256(f"oid-phase3-5-formal/v1:{scene}:{observable}".encode()).digest()
    return int.from_bytes(digest[:8], "little")


def _code_record() -> dict:
    paths = [
        Path(__file__).resolve(),
        ROOT / "src/OID_analysis/v2/oid.py",
        ROOT / "src/OID_analysis/v2/analysis.py",
    ]
    return {str(path.relative_to(ROOT)): sha256_file(path) for path in paths}


def _verify(destination: Path) -> dict:
    report_path, model_path = destination / "phase_3_5_report.json", destination / "phase_3_5_models.npz"
    report = json.loads(report_path.read_text(encoding="utf-8"))
    if report.get("schema_id") != ANALYSIS_SCHEMA or sha256_file(model_path) != report["model_artifact"]["sha256"]:
        raise RuntimeError(f"report/model binding verification failed: {destination}")
    with np.load(model_path, allow_pickle=False) as archive:
        actual = {name: _hash(archive[name]) for name in archive.files}
    if actual != report["model_artifact"]["array_sha256"]:
        raise RuntimeError(f"model array verification failed: {destination}")
    return {"schema_id": "oid-phase3-5-verification/v1", "report_sha256": sha256_file(report_path),
            "model_sha256": sha256_file(model_path), "model_array_sha256": actual,
            "report_json_verified": True, "model_npz_verified": True}


def _run(scene: str, observable: str, canonical: Path, analysis: Path) -> dict:
    destination = analysis / scene / observable
    if destination.exists():
        raise FileExistsError(f"refusing to clobber analysis destination: {destination}")
    products = [canonical / scene / realization for realization in ("r1", "r2", "r3")]
    if any(not (path / "paired_product.npz").is_file() for path in products):
        raise FileNotFoundError(f"canonical products incomplete for {scene}")
    pools = threadpool_info()
    if not pools or any(item.get("num_threads") != 2 for item in pools if item.get("user_api") == "blas"):
        raise RuntimeError(f"BLAS thread contract is not 2: {pools}")
    config = {"runner_schema_id": "oid-phase3-5-formal-run/v2", "scene": scene, "observable": observable,
              "ranks": list(RANKS), "sample_lags": list(LAGS), "surrogate_count": SURROGATES,
              "tested_statistic": "metric_weighted_r2",
              "negative_control_margin": NEGATIVE_CONTROL_MARGIN,
              "rcond": RCOND, "numerical_status": "truncated", "regularization": None,
              "observable_metric": np.ones(OBSERVABLE_CHANNELS[observable]).tolist(),
              "observable_metric_coordinates": "physical incremental observable channel coordinates",
              "observable_coordinate_units": list(OBSERVABLE_COORDINATE_UNITS[observable]),
              "selection_adjusted": True, "seed": _seed(scene, observable), "cpu_only": True,
              "gpu_or_cfd_execution": False, "process_concurrency_limit": 3,
              "thread_environment": {key: os.environ.get(key) for key in THREAD_ENV},
              "threadpool_info": pools, "python": sys.version, "executable": sys.executable,
              "platform": platform.platform(), "numpy": np.__version__, "code_sha256": _code_record()}
    config["config_sha256"] = hashlib.sha256(deterministic_json_bytes(config)).hexdigest()
    report = run_phase_3_5(products, destination, ranks=RANKS, lags=LAGS, observable=observable,
                           surrogate_count=SURROGATES, seed=config["seed"],
                           negative_control_margin=NEGATIVE_CONTROL_MARGIN,
                           rcond=RCOND,
                           observable_metric=np.ones(OBSERVABLE_CHANNELS[observable]),
                           observable_coordinate_units=OBSERVABLE_COORDINATE_UNITS[observable],
                           execution_metadata=config,
                           convective_scale={"D_lattice_units": D, "U0_lattice_units_per_lbm_step": U0,
                                             "D_over_U0_lbm_steps": D / U0,
                                             "stored_times_units": "LBM lattice-step units",
                                             "interpretation": "candidate lags are reported in stored LBM units and as fractions of D/U0"})
    verification = _verify(destination)
    verification_path = destination / "verification.json"
    if verification_path.exists():
        raise FileExistsError(f"refusing to overwrite {verification_path}")
    verification_path.write_bytes(deterministic_json_bytes(verification))
    return {"scene": scene, "observable": observable, "destination": str(destination),
            "report_sha256": verification["report_sha256"], "model_sha256": verification["model_sha256"],
            "verification_sha256": sha256_file(verification_path),
            "wall_time_seconds": report["execution_metadata"]["analysis_wall_time_seconds"],
            "peak_rss_kib": report["execution_metadata"]["peak_rss_kib"]}


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--canonical", type=Path, default=ROOT / "src/OID_analysis/runs/active-two-scene-v2/canonical")
    parser.add_argument("--analysis", type=Path, default=ROOT / "src/OID_analysis/runs/active-two-scene-v2/analysis")
    parser.add_argument("--scene", choices=SCENES)
    parser.add_argument("--observable", choices=OBSERVABLES)
    parser.add_argument("--all", action="store_true", help="launch all four jobs with at most three child processes")
    args = parser.parse_args()
    if args.all:
        if args.scene or args.observable:
            parser.error("--all cannot be combined with --scene/--observable")
        jobs = [(scene, observable) for scene in SCENES for observable in OBSERVABLES]
        running, completed = [], []
        while jobs or running:
            while jobs and len(running) < 3:
                scene, observable = jobs.pop(0)
                command = [sys.executable, str(Path(__file__).resolve()), "--scene", scene, "--observable", observable,
                           "--canonical", str(args.canonical), "--analysis", str(args.analysis)]
                running.append((scene, observable, subprocess.Popen(command, env={**os.environ, **THREAD_ENV})))
            scene, observable, process = running.pop(0)
            code = process.wait()
            if code:
                raise RuntimeError(f"formal job failed ({code}): {scene}/{observable}")
            completed.append({"scene": scene, "observable": observable})
        print(json.dumps(completed, indent=2))
        return 0
    if not args.scene or not args.observable:
        parser.error("provide --scene and --observable, or --all")
    print(json.dumps(_run(args.scene, args.observable, args.canonical.resolve(), args.analysis.resolve()), indent=2))
    return 0


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