#!/usr/bin/env python3
"""Verify the legacy fluidic-pinball geometry/action/force channel contract."""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any, Mapping, Sequence

import numpy as np

REPO_ROOT = Path(__file__).resolve().parents[3]
SRC_ROOT = REPO_ROOT / "src"
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))
if str(SRC_ROOT) not in sys.path:
    sys.path.insert(0, str(SRC_ROOT))

from SR_analysis.configs import LEGACY_CFG_DIR, get_scene
from SR_analysis.utils.cfd_interface import load_legacy_configs
from SR_analysis.utils.provenance import atomic_write_json

CONTRACT_VERSION = "legacy-pinball-order-v1"
BODY_NAMES = ("front", "upper", "lower")


def expected_ledger(scene: str) -> dict[str, Any]:
    cfg = get_scene(scene)
    body_ids = tuple(range(int(cfg["n_objects_env"]) - 3, int(cfg["n_objects_env"])))
    return {
        "contract_version": CONTRACT_VERSION,
        "scene": scene,
        "geometry_order": [
            {"name": "front", "y_offset_L": 0.0},
            {"name": "upper", "y_offset_L": 0.75},
            {"name": "lower", "y_offset_L": -0.75},
        ],
        "body_object_ids": dict(zip(BODY_NAMES, body_ids)),
        "sensor_layout": list(cfg["sensor_layout"]),
        "force_layout": list(cfg["force_layout"]),
        "action_layout": list(cfg["action_layout"]),
        "ppo_decoder": {
            "alpha": "normalized_action * action_scale + action_bias",
            "action_scale": float(cfg["action_scale"]),
            "action_bias": list(cfg["action_bias"]),
            "native_slots": list(body_ids),
        },
        "kernel_binding": "action[id_obj] and obs[2*id_obj:2*id_obj+2] share id_obj",
    }


def _centers_by_id(objects: Mapping[Any, Mapping[str, Any]]) -> list[tuple[float, float, float]]:
    return [tuple(float(value) for value in objects[key]["center"]) for key in sorted(objects, key=int)]


def verify_runtime_objects(flow_field: Any, scene: str) -> dict[str, Any]:
    cfg = get_scene(scene)
    ledger = expected_ledger(scene)
    centers = _centers_by_id(flow_field.objects)
    body_ids = [ledger["body_object_ids"][name] for name in BODY_NAMES]
    body_centers = [centers[index] for index in body_ids]
    cy = (float(flow_field.FIELD_SHAPE[1]) - 1.0) / 2.0
    l0 = 20.0
    expected = [
        (float(cfg["pinball_front_x"]) * l0, cy, 0.0),
        (float(cfg["pinball_rear_x"]) * l0, cy + 0.75 * l0, 0.0),
        (float(cfg["pinball_rear_x"]) * l0, cy - 0.75 * l0, 0.0),
    ]
    np.testing.assert_allclose(body_centers, expected)
    ledger["runtime_centers"] = {name: list(center) for name, center in zip(BODY_NAMES, body_centers)}
    ledger["runtime_verified"] = True
    return ledger


def impulse_commands(n_objects: int, epsilon: float) -> dict[str, np.ndarray]:
    if n_objects < 3:
        raise ValueError("pinball impulse test requires at least three objects")
    commands: dict[str, np.ndarray] = {}
    for offset, name in enumerate(BODY_NAMES, start=n_objects - 3):
        command = np.zeros(n_objects, dtype=np.float32)
        command[offset] = float(epsilon)
        commands[name] = command
    return commands


def build_runtime_environment(scene: str, device: int) -> Any:
    from LegacyCelerisLab import FlowField

    cfg = get_scene(scene)
    cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(cfg["nu"]))
    ff = FlowField(field_cfg, cuda_cfg, device_id=device)
    cy = (ff.FIELD_SHAPE[1] - 1) / 2.0
    l0 = 20.0
    if cfg["scene_id"] == "karman":
        ff.add_cylinder((10.0 * l0, cy, 0.0), l0)
    for y in (cy + 2 * l0, cy, cy - 2 * l0):
        ff.add_sensor((float(cfg["sensor_x"]) * l0, y, 0.0), 5.0)
    ff.add_cylinder((float(cfg["pinball_front_x"]) * l0, cy, 0.0), l0 / 2.0)
    ff.add_cylinder((float(cfg["pinball_rear_x"]) * l0, cy + 0.75 * l0, 0.0), l0 / 2.0)
    ff.add_cylinder((float(cfg["pinball_rear_x"]) * l0, cy - 0.75 * l0, 0.0), l0 / 2.0)
    return ff


def run_impulse_test(scene: str, device: int, epsilon: float, steps: int) -> dict[str, Any]:
    ff = build_runtime_environment(scene, device)
    ledger = verify_runtime_objects(ff, scene)
    n_objects = len(ff.objects)
    zero = np.zeros(n_objects, dtype=np.float32)
    warmup = int(4 * ff.FIELD_SHAPE[0] / float(get_scene(scene)["u0"]))
    ff.run(warmup, zero)
    ff.get_ddf()
    ff.save_ddf()
    responses: dict[str, Any] = {}
    body_ids = ledger["body_object_ids"]
    for name, command in impulse_commands(n_objects, epsilon).items():
        ff.restore_ddf()
        ff.apply_ddf()
        ff.run(steps, command)
        force_pairs = np.asarray(ff.obs, dtype=np.float64).reshape(-1, 2)
        responses[name] = {
            "nonzero_action_slots": np.flatnonzero(command).tolist(),
            "command": float(command[body_ids[name]]),
            "force_at_commanded_body": force_pairs[body_ids[name]].tolist(),
            "all_body_forces": {
                body: force_pairs[index].tolist() for body, index in body_ids.items()
            },
        }
    del ff
    ledger["impulse_test"] = {
        "epsilon": float(epsilon),
        "steps": int(steps),
        "responses": responses,
        "interpretation": "slot identity is exact; force magnitude is diagnostic and is not used to infer causality",
    }
    return ledger


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--scene", choices=("karman_re100", "illusion_1L"), required=True)
    parser.add_argument("--device", type=int, default=2)
    parser.add_argument("--epsilon", type=float, default=1e-4)
    parser.add_argument("--steps", type=int, default=10)
    parser.add_argument("--output", type=Path)
    parser.add_argument("--source-only", action="store_true")
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    document = expected_ledger(args.scene) if args.source_only else run_impulse_test(
        args.scene, args.device, args.epsilon, args.steps
    )
    if args.output:
        atomic_write_json(args.output.resolve(), document)
    print(json.dumps(document, indent=2, sort_keys=True))
    return 0


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