# drl_pinball/legacy_env/legacy_karman_env.py
"""
Standalone Karman cloak re100 environment using LegacyCelerisLab.

This module exactly reproduces env_karman_cloak_standard.py but exposes
ALL intermediate data for validation against the new CelerisLab API.

Usage::

    from legacy_karman_env import legacy_build_re100, legacy_infer_re100

    data = legacy_build_re100(device_id=0)
    # data['target_states'], data['norm'], data['flow_field'], ...
    
    results = legacy_infer_re100(data['flow_field'], model, data['target_states'], data['norm'], n_steps=50)
    # results['sensors'], results['forces'], results['obs'], results['actions'], results['rewards']
"""

from __future__ import annotations

import os
import sys
from collections import deque
from typing import Any, Dict, Optional, Tuple

import numpy as np

# Add repo root for LegacyCelerisLab import
_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
if _REPO not in sys.path:
    sys.path.insert(0, _REPO)

from LegacyCelerisLab import FlowField  # noqa: E402
from LegacyCelerisLab import utils as legacy_utils  # noqa: E402

# ---------------------------------------------------------------------------
# Constants — matching env_karman_cloak_standard.py EXACTLY
# ---------------------------------------------------------------------------
CONFIG_DIR = os.path.join(_REPO, "configs", "legacy_configs")

S_DIM, A_DIM = 12, 3
U0 = 0.01
T0 = 1000
SAMPLE_INTERVAL = 800
FIFO_LEN = 150
CONV_LEN = 30
MAX_STEPS = 500
DATA_TYPE = np.float32
ACTION_SMOOTH_WEIGHT = 0.1  # legacy run() uses internal exponential smoothing

# Geometry constants (in lattice units)
L0 = 20.0
# Disturbance cylinder (id=0 in env order)
DIST_CENTER = (10.0 * L0, None, 0.0)  # y will be set at runtime
DIST_RADIUS = 1.0 * L0

# Sensors (ids=1,2,3)
SENSOR_RADIUS = L0 / 4.0

# Pinball cylinders (ids=4,5,6)
PINBALL_RADIUS = L0 / 2.0
FRONT_CENTER = (30.0 * L0, None, 0.0)  # y = CENTER_Y
BOTTOM_CENTER = (31.3 * L0, None, 0.0)  # y = CENTER_Y - 0.75*L0
TOP_CENTER = (31.3 * L0, None, 0.0)  # y = CENTER_Y + 0.75*L0


def _center_y(ff: FlowField) -> float:
    """Return the center y-coordinate of the flow field."""
    return (ff.FIELD_SHAPE[1] - 1) / 2.0


def _fill_y(cfg: Tuple[float, Optional[float], float], cy: float) -> tuple:
    """Replace None y with actual center y."""
    x, y, z = cfg
    if y is None:
        y = cy
    return (x, y, z)


# ---------------------------------------------------------------------------
# Phase 1a: Build reference dataset (reproduces env.__init__)
# ---------------------------------------------------------------------------

def legacy_build_re100(
    device_id: int = 0,
    viscosity: Optional[float] = None,
) -> Dict[str, Any]:
    """Reproduce env_karman_cloak_standard.__init__() exactly.
    
    Parameters
    ----------
    device_id : int
        GPU device ID.
    viscosity : float, optional
        Override viscosity (default: 0.004 for Re=100).
    
    Returns
    -------
    dict with keys:
        flow_field : FlowField — the CFD instance (state at end of init)
        target_states : ndarray (FIFO_LEN, 6) — target sensor signals
        norm : dict with force_norm_fact, sens_deviation, sens_norm_fact
        config : dict with all runtime parameters
        fifo_after_bias : deque — FIFO state after bias rollout
    """
    # Default viscosity for Re=100 (code Re, using 2D reference)
    if viscosity is None:
        viscosity = 0.004

    # Load legacy configs
    cuda_cfg = legacy_utils.load_cuda_config(
        os.path.join(CONFIG_DIR, "config_cuda.json")
    )
    field_cfg = legacy_utils.load_flow_field_config(
        os.path.join(CONFIG_DIR, "config_flowfield.json")
    )
    # Override viscosity
    field_cfg = field_cfg._replace(viscosity=float(viscosity))

    # -- Step 0: Create FlowField ------------------------------------------
    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
    cy = _center_y(ff)
    NX = ff.FIELD_SHAPE[0]
    NY = ff.FIELD_SHAPE[1]

    # -- Step 1: Add disturbance cylinder + sensors -------------------------
    # Order matters for obs indexing: dist_cyl(0), sensor0(1), sensor1(2), sensor2(3)
    ff.add_cylinder(_fill_y(DIST_CENTER, cy), DIST_RADIUS)
    for y_off in [2.0, 0.0, -2.0]:
        sc = (40.0 * L0, cy + y_off * L0, 0.0)
        ff.add_sensor(sc, SENSOR_RADIUS)

    n_obj_phase1 = ff.obs.size // 2  # 4 objects
    assert n_obj_phase1 == 4, f"Expected 4 objects after sensors, got {n_obj_phase1}"

    # -- Step 2: Stabilize --------------------------------------------------
    stabilize_steps = int(4 * NX / U0)
    ff.run(stabilize_steps, np.zeros(n_obj_phase1, dtype=DATA_TYPE))

    # -- Step 3: Record target signals --------------------------------------
    target_states = np.empty((0, 6), dtype=DATA_TYPE)
    for _ in range(FIFO_LEN):
        ff.run(SAMPLE_INTERVAL, np.zeros(n_obj_phase1, dtype=DATA_TYPE))
        new_state = ff.obs.copy()[2:8]  # sensors only skip dist_cyl
        target_states = np.vstack((target_states, new_state))

    # -- Step 4: Add pinball cylinders (ids=4,5,6) -------------------------
    ff.add_cylinder(_fill_y(FRONT_CENTER, cy), PINBALL_RADIUS)
    ff.add_cylinder(_fill_y(BOTTOM_CENTER, cy - 0.75 * L0), PINBALL_RADIUS)
    ff.add_cylinder(_fill_y(TOP_CENTER, cy + 0.75 * L0), PINBALL_RADIUS)

    n_obj_total = ff.obs.size // 2  # 7 objects
    assert n_obj_total == 7, f"Expected 7 objects, got {n_obj_total}"

    # -- Step 5: Stabilize with pinball -------------------------------------
    ff.run(stabilize_steps, np.zeros(n_obj_total, dtype=DATA_TYPE))

    # -- Step 6: Checkpoint DDF (steady pinball + disturbance state) --------
    ff.get_ddf()
    ff.save_ddf()

    # -- Step 7: Zero-action norm collection --------------------------------
    fifo = deque(maxlen=FIFO_LEN)
    for _ in range(FIFO_LEN):
        ff.run(SAMPLE_INTERVAL, np.zeros(n_obj_total, dtype=DATA_TYPE))
        fifo.append(ff.obs.copy()[2:14])  # sensor[6] + force[6]

    temp_states = np.array(fifo, dtype=DATA_TYPE)
    force_norm_fact = 6.0 * float(np.max(np.abs(temp_states[:, 6:12])))
    sens_deviation = np.mean(temp_states[:, 0:6], axis=0).astype(DATA_TYPE)
    sens_norm_fact = np.zeros(6, dtype=DATA_TYPE)
    for i in range(6):
        sens_norm_fact[i] = 5.0 * float(np.max(np.abs(temp_states[:, i] - sens_deviation[i])))

    # -- Step 8: Bias-action rollout (for FIFO init in controlled runs) -----
    ff.apply_ddf()  # restore pre-bias state

    # Action bias: front=0, bottom=-4*U0, top=4*U0
    bias_arr = np.zeros(n_obj_total, dtype=DATA_TYPE)
    bias_arr[n_obj_total - 3] = float((0.0 * 8.0 + 0.0) * U0)     # front = 0
    bias_arr[n_obj_total - 2] = float((0.0 * 8.0 + (-4.0)) * U0)  # bottom = -4*U0
    bias_arr[n_obj_total - 1] = float((0.0 * 8.0 + 4.0) * U0)     # top = 4*U0

    fifo.clear()
    for _ in range(FIFO_LEN):
        ff.run(SAMPLE_INTERVAL, bias_arr)
        fifo.append(ff.obs.copy()[2:14])

    save_states = np.array(list(fifo), dtype=DATA_TYPE)

    # -- Step 9: Restore to steady state (ready for reset) ------------------
    ff.apply_ddf()

    norm = {
        "force_norm_fact": force_norm_fact,
        "sens_deviation": sens_deviation.tolist(),
        "sens_norm_fact": sens_norm_fact.tolist(),
        "save_states": save_states,
        "action_bias": [0.0, -4.0, 4.0],
        "n_obj_total": n_obj_total,
    }

    config = {
        "device_id": device_id,
        "viscosity": viscosity,
        "u0": U0,
        "sample_interval": SAMPLE_INTERVAL,
        "fifo_len": FIFO_LEN,
        "conv_len": CONV_LEN,
        "nx": NX,
        "ny": NY,
        "n_obj_total": n_obj_total,
        "action_scale": 8.0,
        "action_bias": [0.0, -4.0, 4.0],
    }

    return {
        "flow_field": ff,
        "target_states": target_states,
        "norm": norm,
        "config": config,
        "fifo_after_bias": fifo,
    }


# ---------------------------------------------------------------------------
# Phase 1b: Inference — reproduces env.step() exactly
# ---------------------------------------------------------------------------

def legacy_infer_re100(
    flow_field: FlowField,
    model: Any,
    target_states: np.ndarray,
    norm: Dict[str, Any],
    n_steps: int = 50,
    *,
    use_deterministic: bool = True,
) -> Dict[str, np.ndarray]:
    """Run Karman cloak re100 controlled inference with legacy CFD.
    
    This follows the exact pattern in env_karman_cloak_standard.step() and
    analysis_crossre/scripts/phase1_infer.py.
    
    Parameters
    ----------
    flow_field : FlowField
        Initialized flow field (must be in steady pinball state).
    model : PPO
        Trained PPO model (with Sin activation).
    target_states : ndarray (FIFO_LEN, 6)
        Target sensor signals.
    norm : dict
        Normalization factors from legacy_build_re100().
    n_steps : int
        Number of inference steps (each = SAMPLE_INTERVAL LBM steps).
    use_deterministic : bool
        Use deterministic action (True) or stochastic (False).
    
    Returns
    -------
    dict with keys:
        sensors : (n_steps, 6) raw sensor velocities
        forces : (n_steps, 6) raw forces (all 6 force components)
        obs : (n_steps, 12) normalised DRL observations
        actions : (n_steps, 3) normalised DRL actions in [-1, 1]
        rewards : (n_steps,) step rewards
        reward_cd : (n_steps,)
        reward_cl : (n_steps,)
        reward_sim : (n_steps,)
        similarities : (n_steps,) DTW similarity scores
    """
    n_obj_total = norm.get("n_obj_total", 7)
    action_scale = 8.0
    action_bias = np.array(norm.get("action_bias", [0.0, -4.0, 4.0]), dtype=np.float32)
    force_norm_fact = float(norm["force_norm_fact"])
    sens_deviation = np.array(norm["sens_deviation"], dtype=np.float32)
    sens_norm_fact = np.array(norm["sens_norm_fact"], dtype=np.float32)

    # Restore steady state
    flow_field.restore_ddf()
    flow_field.apply_ddf()

    # Bias-action FIFO init (reproduces env.__init__ bias rollout)
    fifo = deque(maxlen=FIFO_LEN)
    bias_arr = np.zeros(n_obj_total, dtype=DATA_TYPE)
    bias_arr[n_obj_total - 3] = float(action_bias[0] * U0)
    bias_arr[n_obj_total - 2] = float(action_bias[1] * U0)
    bias_arr[n_obj_total - 1] = float(action_bias[2] * U0)

    for _ in range(FIFO_LEN):
        flow_field.run(SAMPLE_INTERVAL, bias_arr)
        fifo.append(flow_field.obs.copy()[2:14])

    # Inference loop
    sens_list, forc_list, obs_list = [], [], []
    action_list, reward_list = [], []
    reward_cd_list, reward_cl_list, reward_sim_list = [], [], []
    sim_list = []

    obs = np.zeros(S_DIM, dtype=np.float32)

    for step in range(n_steps):
        # --- PPO action ---
        action, _states = model.predict(obs, deterministic=use_deterministic)
        action = action.astype(np.float32).flatten()
        action_list.append(action.copy())

        # --- Convert to legacy action array ---
        action_arr = np.zeros(n_obj_total, dtype=DATA_TYPE)
        omega = (action * action_scale + action_bias) * U0
        action_arr[n_obj_total - 3:] = omega

        # --- Run CFD ---
        # Context management: push/pop to avoid PyTorch CUDA context conflicts
        flow_field.context.push()
        try:
            flow_field.run(SAMPLE_INTERVAL, action_arr)
        finally:
            flow_field.context.pop()

        # --- Read telemetry ---
        obs_slice = flow_field.obs.copy()[2:14]
        fifo.append(obs_slice)
        sens_list.append(obs_slice[0:6].copy())
        forc_list.append(obs_slice[6:12].copy())

        # --- Build normalised observation ---
        forces_norm = obs_slice[6:12] / force_norm_fact
        sens_norm = (obs_slice[0:6] - sens_deviation) / sens_norm_fact
        obs = np.clip(np.hstack([forces_norm, sens_norm]), -1.0, 1.0).astype(np.float32)
        obs_list.append(obs)

        # --- Compute reward (exactly matching env logic) ---
        states_arr = np.array(fifo, dtype=np.float32)
        if len(states_arr) >= CONV_LEN:
            forces = states_arr[-1, 6:12] / force_norm_fact
            cd = float((forces[0] + forces[2] + forces[4]) / 3.0)
            cl = float((forces[1] + forces[3] + forces[5]) / 3.0)

            # Similarity = lag-compensated DTW over all 6 sensor channels
            # calc_lag on middle sensor (index 1 = sensor1_uy)
            target_seq = target_states[CONV_LEN:2 * CONV_LEN, 1]
            state_seq = states_arr[-CONV_LEN:, 1]
            lag = _calc_lag(target_seq, state_seq)

            sim_sum = 0.0
            for i in range(6):
                t_seq = np.roll(target_states[:, i], -lag)[CONV_LEN:2 * CONV_LEN]
                s_seq = states_arr[-CONV_LEN:, i]
                sim_sum += _calc_dtw_sim(t_seq, s_seq) / 6.0
            similarities = float(sim_sum)
            sim_list.append(similarities)

            r_cd = float(np.exp(-abs(cd * 20.0)))
            r_cl = float(np.exp(-abs(cl * 80.0)))
            r_sim = float(np.exp(-10.0 * abs(similarities - 1.0)))
            reward = float(min(0.3 * r_cd + 0.4 * r_cl + 0.3 * r_sim, 1.0))
        else:
            reward = 0.0
            r_cd = 0.0
            r_cl = 0.0
            r_sim = 0.0
            similarities = 0.0

        reward_list.append(reward)
        reward_cd_list.append(r_cd)
        reward_cl_list.append(r_cl)
        reward_sim_list.append(r_sim)

    return {
        "sensors": np.array(sens_list, dtype=np.float32),
        "forces": np.array(forc_list, dtype=np.float32),
        "obs": np.array(obs_list, dtype=np.float32),
        "actions": np.array(action_list, dtype=np.float32),
        "rewards": np.array(reward_list, dtype=np.float32),
        "reward_cd": np.array(reward_cd_list, dtype=np.float32),
        "reward_cl": np.array(reward_cl_list, dtype=np.float32),
        "reward_sim": np.array(reward_sim_list, dtype=np.float32),
        "similarities": np.array(sim_list, dtype=np.float32) if sim_list else np.zeros(n_steps, dtype=np.float32),
    }


# ---------------------------------------------------------------------------
# Helper: uncontrolled rollout (zero action)
# ---------------------------------------------------------------------------

def legacy_uncontrolled_re100(
    flow_field: FlowField,
    n_steps: int = 50,
) -> Dict[str, np.ndarray]:
    """Run uncontrolled Karman cloak re100 inference.
    
    Parameters
    ----------
    flow_field : FlowField
        Must be in saved DDF state (steady pinball + disturbance).
    n_steps : int
        Number of steps.
    
    Returns
    -------
    dict with sensors, forces.
    """
    n_obj_total = 7

    flow_field.restore_ddf()
    flow_field.apply_ddf()

    sens_list, forc_list = [], []

    for _ in range(n_steps):
        flow_field.run(SAMPLE_INTERVAL, np.zeros(n_obj_total, dtype=DATA_TYPE))
        obs_slice = flow_field.obs.copy()[2:14]
        sens_list.append(obs_slice[0:6].copy())
        forc_list.append(obs_slice[6:12].copy())

    return {
        "sensors": np.array(sens_list, dtype=np.float32),
        "forces": np.array(forc_list, dtype=np.float32),
    }


# ---------------------------------------------------------------------------
# DTW helpers (exact copies from env_karman_cloak_standard.py)
# ---------------------------------------------------------------------------

def _calc_lag(target: np.ndarray, state: np.ndarray) -> int:
    target_mean = float(np.mean(target))
    state_mean = float(np.mean(state))
    correlation = np.correlate(
        target - target_mean,
        state - state_mean,
        mode="full",
    )
    lags = np.arange(-len(target) + 1, len(target))
    return int(lags[np.argmax(correlation)])


def _calc_dtw_sim(target: np.ndarray, state: np.ndarray) -> float:
    n = len(target)
    m = len(state)
    dtw_matrix = np.full((n + 1, m + 1), np.inf)
    dtw_matrix[0, 0] = 0.0
    for i in range(1, n + 1):
        for j in range(1, m + 1):
            cost = abs(float(target[i - 1]) - float(state[j - 1]))
            last_min = min(
                dtw_matrix[i - 1, j],
                dtw_matrix[i, j - 1],
                dtw_matrix[i - 1, j - 1],
            )
            dtw_matrix[i, j] = cost + last_min
    return float(1.0 - dtw_matrix[n, m] / n)


# ---------------------------------------------------------------------------
# CLI entry point for reference dataset generation
# ---------------------------------------------------------------------------

def main():
    """Generate reference dataset for Karman re100 validation."""
    import argparse
    import json

    ap = argparse.ArgumentParser(description="Legacy Karman re100 reference dataset")
    ap.add_argument("--device", type=int, default=0, help="GPU device ID")
    ap.add_argument("--out", type=str, default="output/validate_re100", help="Output directory")
    args = ap.parse_args()

    out_dir = os.path.abspath(args.out)
    os.makedirs(out_dir, exist_ok=True)

    print("=== Building legacy Karman re100 reference dataset ===")
    print(f"Output: {out_dir}")

    # Build env
    data = legacy_build_re100(device_id=args.device)
    ff = data["flow_field"]

    # Save target and norm
    np.savez(os.path.join(out_dir, "target.npz"),
             target_states=data["target_states"])

    norm_json = {
        "force_norm_fact": float(data["norm"]["force_norm_fact"]),
        "sens_deviation": list(float(x) for x in data["norm"]["sens_deviation"]),
        "sens_norm_fact": list(float(x) for x in data["norm"]["sens_norm_fact"]),
        "action_bias": data["norm"]["action_bias"],
    }
    with open(os.path.join(out_dir, "norm.json"), "w") as f:
        json.dump(norm_json, f, indent=2)

    np.savez(os.path.join(out_dir, "save_states.npz"),
             save_states=data["norm"]["save_states"])

    # Save config
    with open(os.path.join(out_dir, "config.json"), "w") as f:
        json.dump({k: str(v) if isinstance(v, (np.integer, np.floating)) else v
                   for k, v in data["config"].items()}, f, indent=2)

    # Uncontrolled rollout
    print("  uncontrolled rollout (50 steps)...")
    unc = legacy_uncontrolled_re100(ff, n_steps=50)
    np.savez(os.path.join(out_dir, "uncontrolled.npz"),
             sensors=unc["sensors"], forces=unc["forces"])
    
    print("  Reference dataset saved.")
    print(f"    force_norm_fact = {norm_json['force_norm_fact']:.6f}")
    print(f"    sens_deviation  = {norm_json['sens_deviation']}")
    print(f"    sens_norm_fact  = {norm_json['sens_norm_fact']}")

    # Cleanup
    del ff
    print("Done.")


if __name__ == "__main__":
    main()
