"""Create independently initialized policy-consistent parent trajectories for OID v2.

Independence is obtained without changing the scene distribution: every realization
uses identical geometry, fluid parameters, policy, and zero/bias actions, but advances
the deterministic CFD state by an explicitly recorded number of whole control windows
after the common zero-action warm-up.  The resulting phase-window offset changes the
checkpoint while preserving the autonomous scene's invariant distribution.  Formal
realizations must use distinct ``realization_id`` and ``warmup_offset_windows`` values.
"""
from __future__ import annotations

import argparse
from contextlib import ExitStack
from collections import deque
import hashlib
import json
import os
import sys
from pathlib import Path
from typing import Any

import numpy as np

try:  # Keep CLI/contract helpers importable in CPU-only test environments.
    import pycuda.driver as cuda
    from LegacyCelerisLab import FlowField
    from CCD_analysis.utils.cfd_interface import load_legacy_configs, load_ppo_model
except ImportError:  # pragma: no cover - exercised through explicit runtime guard
    cuda = None
    FlowField = Any
    load_legacy_configs = load_ppo_model = None

from CCD_analysis.utils.resampling import gen_target_states_at
from CCD_analysis.utils.run_contract import ContractError, array_sha256, global_cfd_lease, preflight_disk
from OID_analysis.configs import CENTER_Y, LEGACY_CFG_DIR, get_scene, model_path_for_scene
from .artifacts import atomic_save_npz, atomic_write_json
from .contracts import LOGICAL_GPU, PHYSICAL_GPU, SHARED_CFD_LEASE
from .manifest import sha256_file, validate_execution_environment

L0 = 20.0
FIFO_LEN = 150
PARENT_SCENES = ("karman_re100", "illusion_1.0L")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--scene", choices=PARENT_SCENES, required=True)
    parser.add_argument("--realization-id", required=True,
                        help="Unique formal realization identifier (for example r2 or r3)")
    parser.add_argument("--warmup-offset-windows", type=int, required=True,
                        help="Distinct non-negative deterministic phase-window offset")
    parser.add_argument("--device", type=int, choices=(PHYSICAL_GPU,), required=True,
                        help="Physical GPU; CUDA_VISIBLE_DEVICES=2 maps it to logical GPU 0")
    parser.add_argument("--count", type=int, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--norm", type=Path, required=True)
    parser.add_argument("--harmonics", type=Path,
                        help="Required only for illusion_1.0L")
    parser.add_argument("--compare-parent", type=Path, action="append", default=[],
                        help="Existing parent directory; repeat for distinctness checks")
    return parser


def validate_parent_request(*, scene: str, realization_id: str, count: int,
                            warmup_offset_windows: int, harmonics_path: Path | None) -> None:
    if scene not in PARENT_SCENES:
        raise ContractError(f"parent scene must be one of {PARENT_SCENES}")
    if not realization_id.strip():
        raise ContractError("realization_id must be a non-empty string")
    if count <= 0:
        raise ContractError("count must be positive")
    if warmup_offset_windows < 0:
        raise ContractError("warmup_offset_windows must be non-negative")
    if scene == "illusion_1.0L" and harmonics_path is None:
        raise ContractError("illusion_1.0L requires --harmonics")
    if scene == "karman_re100" and harmonics_path is not None:
        raise ContractError("karman_re100 does not use --harmonics")


def _copy_device(allocation: Any, template: np.ndarray) -> np.ndarray:
    result = np.empty_like(template)
    cuda.memcpy_dtoh(result, allocation)
    return result


def _run(ff: FlowField, steps: int, action: np.ndarray) -> None:
    ff.context.push()
    try:
        ff.run(steps, action)
    finally:
        ff.context.pop()


def _add_upstream_geometry(ff: FlowField, cfg: dict) -> None:
    if cfg.get("has_disturbance", False):
        ff.add_cylinder((10.0 * L0, CENTER_Y, 0.0), L0)
    for y_offset in (2.0, 0.0, -2.0):
        ff.add_sensor((cfg["sensor_x"] * L0, CENTER_Y + y_offset * L0, 0.0), L0 / 4.0)


def _add_pinball(ff: FlowField, cfg: dict) -> None:
    ff.add_cylinder((cfg["pinball_front_x"] * L0, CENTER_Y, 0.0), L0 / 2.0)
    ff.add_cylinder((cfg["pinball_rear_x"] * L0, CENTER_Y + 0.75 * L0, 0.0), L0 / 2.0)
    ff.add_cylinder((cfg["pinball_rear_x"] * L0, CENTER_Y - 0.75 * L0, 0.0), L0 / 2.0)


def _distinctness_diagnostics(manifest: dict, compare_parents: list[Path]) -> list[dict]:
    """Compare identity hashes against existing parents and fail on duplication."""
    diagnostics = []
    for directory in compare_parents:
        path = directory / "parent_manifest.json" if directory.is_dir() else directory
        existing = json.loads(path.read_text(encoding="utf-8"))
        checks = {
            "realization_id_distinct": existing.get("realization_id") != manifest["realization_id"],
            "checkpoint_distinct": existing.get("bundle_sha256") != manifest["bundle_sha256"],
            "trajectory_distinct": existing.get("trajectory_sha256") != manifest["trajectory_sha256"],
            "pre_bias_state_distinct": existing.get("pre_bias_ddf_sha256") != manifest["pre_bias_ddf_sha256"],
        }
        record = {
            "compared_parent_manifest": str(path),
            "compared_parent_manifest_sha256": sha256_file(path),
            "compared_realization_id": existing.get("realization_id"),
            **checks,
        }
        diagnostics.append(record)
        if not all(checks.values()):
            failed = sorted(key for key, value in checks.items() if not value)
            raise ContractError(f"parent is not distinct from {path}: {failed}")
    return diagnostics


def create_parent(*, scene: str, realization_id: str, warmup_offset_windows: int,
                  device: int, count: int, output_dir: Path, norm_path: Path,
                  harmonics_path: Path | None = None,
                  compare_parents: list[Path] | None = None) -> dict:
    validate_parent_request(scene=scene, realization_id=realization_id, count=count,
                            warmup_offset_windows=warmup_offset_windows,
                            harmonics_path=harmonics_path)
    if cuda is None:
        raise RuntimeError("parent CFD generation requires pycuda_3_10 with LegacyCelerisLab")
    execution = validate_execution_environment(device=device, require_gpu=True)
    if execution["logical_device"] != LOGICAL_GPU:
        raise RuntimeError("isolated physical GPU 2 must be addressed as logical GPU 0")
    preflight_disk(output_dir, 0, reserve_bytes=1024**3)
    cfg = get_scene(scene)
    if output_dir.exists():
        raise FileExistsError(output_dir)
    partial = output_dir.with_name(f"{output_dir.name}.partial.{os.getpid()}")
    partial.mkdir(parents=True, exist_ok=False)
    ff = model = None
    lease_stack = ExitStack()
    try:
        lease_stack.enter_context(global_cfd_lease(
            SHARED_CFD_LEASE, run_id=f"parent-{scene}-{realization_id}", scene=scene,
            physical_gpu=PHYSICAL_GPU, command=sys.argv,
        ))
        norm = json.loads(norm_path.read_text(encoding="utf-8"))
        harmonics = (json.loads(harmonics_path.read_text(encoding="utf-8"))
                     if harmonics_path is not None else None)
        cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR)
        field_cfg = field_cfg._replace(viscosity=float(cfg["nu"]), velocity=float(cfg["u0"]))
        ff = FlowField(field_cfg, cuda_cfg, device_id=LOGICAL_GPU)
        _add_upstream_geometry(ff, cfg)
        common_warmup_steps = int(4 * 1280 / cfg["u0"])
        # Preserve the established Karman construction: develop the disturbance and
        # sensors before introducing the pinball. Illusion has no separate phase.
        upstream_object_count = 4 if cfg.get("has_disturbance", False) else 3
        if cfg.get("has_disturbance", False):
            _run(ff, common_warmup_steps, np.zeros(upstream_object_count, np.float32))
        _add_pinball(ff, cfg)
        n_objects = int(cfg["n_objects_env"])
        zero = np.zeros(n_objects, np.float32)
        _run(ff, common_warmup_steps, zero)
        offset_steps = warmup_offset_windows * int(cfg["sample_interval"])
        if offset_steps:
            _run(ff, offset_steps, zero)
        ff.get_ddf()
        pre_bias = ff.ddf.copy()

        fifo: deque[np.ndarray] = deque(maxlen=FIFO_LEN)
        bias = zero.copy()
        bias[-3:] = np.asarray(cfg["action_bias"], np.float32) * cfg["u0"]
        obs_start, obs_end = cfg["obs_slice"]
        for _ in range(FIFO_LEN):
            _run(ff, cfg["sample_interval"], bias)
            fifo.append(ff.obs.copy()[obs_start:obs_end])
        current = _copy_device(ff.ddf_gpu, ff.ddf)
        temp = _copy_device(ff.temp_gpu, ff.ddf)
        bundle_path = atomic_save_npz(partial / "checkpoint_bundle.npz", current=current, temp=temp,
                                      action=ff.action.copy(), fifo=np.asarray(fifo, np.float32))

        raw = ff.obs.copy()[obs_start:obs_end]
        force_scale = float(norm["force_norm_fact"])
        sensor_mean = np.asarray(norm["sens_deviation"], np.float32)
        sensor_scale = np.asarray(norm["sens_norm_fact"], np.float32)
        obs12 = np.clip(np.hstack((raw[6:12] / force_scale,
            (raw[:6] - sensor_mean) / sensor_scale)), -1.0, 1.0).astype(np.float32)
        obs = obs12
        if harmonics is not None:
            target = gen_target_states_at(0, harmonics)
            obs = np.clip(np.hstack((obs12, target[:2] / force_scale)), -1.0, 1.0).astype(np.float32)
        model_file = Path(model_path_for_scene(scene))
        model = load_ppo_model(str(model_file), device=f"cuda:{LOGICAL_GPU}",
                               s_dim=int(cfg["s_dim"]), a_dim=3)
        model.set_random_seed(19)
        sensors, forces, actions = [], [], []
        for step in range(count):
            action, _ = model.predict(obs, deterministic=True)
            action = np.asarray(action, np.float32).reshape(-1)
            applied = zero.copy()
            applied[-3:] = (action * cfg["action_scale"]
                            + np.asarray(cfg["action_bias"], np.float32)) * cfg["u0"]
            _run(ff, cfg["sample_interval"], applied)
            raw = ff.obs.copy()[obs_start:obs_end]
            sensors.append(raw[:6]); forces.append(raw[6:12]); actions.append(action)
            obs12 = np.clip(np.hstack((raw[6:12] / force_scale,
                (raw[:6] - sensor_mean) / sensor_scale)), -1.0, 1.0).astype(np.float32)
            obs = obs12
            if harmonics is not None:
                target = gen_target_states_at(step + 1, harmonics)
                obs = np.clip(np.hstack((obs12, target[:2] / force_scale)), -1.0, 1.0).astype(np.float32)
        trajectory_arrays = {
            "sensors": np.asarray(sensors, np.float32),
            "forces": np.asarray(forces, np.float32),
            "actions": np.asarray(actions, np.float32),
            "control_end_lattice_steps": np.arange(1, count + 1, dtype=np.int64) * cfg["sample_interval"],
            "lattice_steps": np.arange(1, count + 1, dtype=np.int64) * cfg["sample_interval"],
            "sample_ids": np.arange(count, dtype=np.int64),
            "times": np.arange(1, count + 1, dtype=np.float64) * cfg["sample_interval"],
        }
        trajectory_path = atomic_save_npz(partial / "diagnostic_trajectory.npz", **trajectory_arrays)
        geometry = {
            "disturbance_x": 10.0 if cfg.get("has_disturbance") else None,
            "sensor_x": cfg["sensor_x"], "front_x": cfg["pinball_front_x"],
            "rear_x": cfg["pinball_rear_x"],
            "object_order": ((["disturbance"] if cfg.get("has_disturbance") else [])
                             + ["sensor_upper", "sensor_center", "sensor_lower", "front", "upper", "lower"]),
        }
        manifest = {
            "schema_version": "oid-v2-parent-v2", "scene": scene,
            "realization_id": realization_id, "parent_id": output_dir.name, "count": count,
            "device_physical": device, "device_logical": LOGICAL_GPU,
            "initialization": {
                "method": "deterministic-whole-control-window-phase-offset",
                "common_zero_warmup_lattice_steps": common_warmup_steps,
                "warmup_offset_windows": warmup_offset_windows,
                "warmup_offset_lattice_steps": offset_steps,
                "bias_fifo_windows": FIFO_LEN,
                "distribution_preservation": "geometry, fluid parameters, policy, and actions are unchanged; only deterministic autonomous warm-up duration differs",
            },
            "phase": {"origin": "post-offset bias-FIFO checkpoint", "target_phase_index": 0},
            "geometry": geometry, "control_interval": cfg["sample_interval"],
            # Parent diagnostics are intentionally control-end telemetry. Formal
            # branch observables use the separate field-level telemetry contract.
            "diagnostic_telemetry_clock": "control_end",
            "telemetry_array_hashes": {
                key: array_sha256(value) for key, value in trajectory_arrays.items()
            },
            "bundle_sha256": sha256_file(bundle_path),
            "trajectory_sha256": sha256_file(trajectory_path),
            "model_sha256": sha256_file(model_file), "norm_sha256": sha256_file(norm_path),
            "pre_bias_ddf_sha256": hashlib.sha256(pre_bias.tobytes()).hexdigest(),
        }
        if harmonics_path is not None:
            manifest["harmonics_sha256"] = sha256_file(harmonics_path)
        manifest["distinctness_diagnostics"] = _distinctness_diagnostics(
            manifest, list(compare_parents or []))
        manifest_path = partial / "parent_manifest.json"
        atomic_write_json(manifest_path, manifest)
        manifest_sha256 = sha256_file(manifest_path)
        atomic_write_json(partial / "_SUCCESS", {
            "parent_manifest_sha256": manifest_sha256,
            "checkpoint_bundle_sha256": manifest["bundle_sha256"],
            "trajectory_sha256": manifest["trajectory_sha256"],
        })
        os.replace(partial, output_dir)
        return {**manifest, "parent_manifest_sha256": manifest_sha256}
    except BaseException:
        if partial.exists():
            os.replace(partial, partial.with_name(partial.name.replace(".partial.", ".quarantine.")))
        raise
    finally:
        if ff is not None:
            try:
                ff.context.detach()
            except Exception:
                pass
        del model, ff
        lease_stack.close()


# Backward-compatible function name, now requiring an explicit independence contract.
def create_illusion_parent(**kwargs: Any) -> dict:
    return create_parent(scene="illusion_1.0L", **kwargs)


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    validate_parent_request(scene=args.scene, realization_id=args.realization_id,
                            count=args.count, warmup_offset_windows=args.warmup_offset_windows,
                            harmonics_path=args.harmonics)
    result = create_parent(scene=args.scene, realization_id=args.realization_id,
        warmup_offset_windows=args.warmup_offset_windows, device=args.device,
        count=args.count, output_dir=args.output_dir, norm_path=args.norm,
        harmonics_path=args.harmonics, compare_parents=args.compare_parent)
    print(json.dumps(result, indent=2))
    return 0


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