"""CFD interface for LegacyCelerisLab (pycuda_3_10 env).

All functions use the LegacyCelerisLab (old) CFD API via:
    from LegacyCelerisLab import FlowField

Must be run inside: conda run -n pycuda_3_10

NOTE: This module should be imported directly, not through SR_analysis.utils
because it requires pycuda. Other active utils (feature_builder, g_operator)
do NOT require pycuda and can be imported from the __init__.
"""
from __future__ import annotations

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

import numpy as np

# -- Import legacy CFD -------------------------------------------------------
# LegacyCelerisLab lives at the repo root; SR_analysis is at repo_root/src/SR_analysis.
_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
if _REPO not in sys.path:
    sys.path.insert(0, _REPO)
_SRC = os.path.join(_REPO, "src")
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

if TYPE_CHECKING:
    from LegacyCelerisLab import FlowField


def _legacy_cfd():
    """Import CUDA-backed legacy CFD only when an environment helper needs it."""
    from LegacyCelerisLab import FlowField
    from LegacyCelerisLab import utils as legacy_utils
    return FlowField, legacy_utils

# ---------------------------------------------------------------------------
# Action-smoothing constant (legacy run() internal)
# ---------------------------------------------------------------------------
ACTION_SMOOTH_WEIGHT = 0.1  # used by FlowField.run() internally


def nu_from_re(re_code: float, u0: float = 0.01, d_ref: float = 40.0) -> float:
    """Return kinematic viscosity for a given code Reynolds number.

    ``re_code`` uses reference length *2*D* = 40.0 (matching model file naming).
    """
    return u0 * d_ref / re_code


def load_legacy_configs(config_dir: str) -> Tuple[Any, Any]:
    """Load and return legacy (cuda_config, field_config) from *config_dir*."""
    _, legacy_utils = _legacy_cfd()
    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")
    )
    return cuda_cfg, field_cfg


# ---------------------------------------------------------------------------
# Environment helpers -- Karman cloak (disturbance cylinder + pinball)
# ---------------------------------------------------------------------------

def build_karman_cloak_env(
    flow_field: FlowField,
    *,
    u0: float,
    l0: float,
    sample_interval: int,
    fifo_len: int,
    data_type: type,
) -> Tuple[np.ndarray, dict]:
    """Phase 0-1: add dist-cylinder & 3 sensors, stabilize, record target.

    Steps (mirrors env_karman_cloak_standard.__init__):
    1. add dist_cylinder (id=0)
    2. add 3 sensors (id=1,2,3)
    3. stabilize run(4*NX/U0, zero-action[4])
    4. record FIFO_LEN x run(SAMPLE_INTERVAL, zero[4]), collect obs[2:8]

    Returns
    -------
    target_states : ndarray (FIFO_LEN, 6)  -- sensor0/1/2 ux,uy
    info : dict with n_objects, NX, NY
    """
    # dist cylinder
    center = (10.0 * l0, (flow_field.FIELD_SHAPE[1] - 1) / 2, 0.0)
    flow_field.add_cylinder(center, l0)

    # sensors
    for y_off in [2.0, 0.0, -2.0]:
        sc = (40.0 * l0, (flow_field.FIELD_SHAPE[1] - 1) / 2 + y_off * l0, 0.0)
        flow_field.add_sensor(sc, l0 / 4.0)

    n_obj = flow_field.obs.size // 2

    # stabilize
    stabilize_steps = int(4 * flow_field.FIELD_SHAPE[0] / u0)
    print(f"  stabilising ({stabilize_steps} steps)...")
    flow_field.run(stabilize_steps, np.zeros(n_obj, dtype=data_type))

    # record target (only sensor signals = obs[2:8])
    target_states = np.empty((0, 6), dtype=data_type)
    for _ in range(fifo_len):
        flow_field.run(sample_interval, np.zeros(n_obj, dtype=data_type))
        new_state = flow_field.obs.copy()[2:8]
        target_states = np.vstack((target_states, new_state))

    print(f"  target recorded: {target_states.shape}")
    return target_states, {"n_objects": n_obj, "NX": flow_field.FIELD_SHAPE[0],
                           "NY": flow_field.FIELD_SHAPE[1]}


def add_pinball(
    flow_field: FlowField,
    *,
    l0: float,
    u0: float,
    sample_interval: int,
    fifo_len: int,
    data_type: type,
    action_bias: Optional[Tuple[float, float, float]] = None,
    pinball_front_x: float = 30.0,
    pinball_rear_x: float = 31.3,
    obs_slice_start: int = 2,
    obs_slice_end: int = 14,
    n_objects_total: Optional[int] = None,
) -> dict:
    """Add pinball cylinders, stabilize, compute norm, bias rollout.

    Steps:
    1. add front, upper, lower cylinders
    2. stabilize run(4*NX/U0, zero-action)
    3. get_ddf() + save_ddf()  (checkpoint)
    4. FIFO_LEN x run(SAMPLE_INTERVAL, zero) -> compute norm
    5. apply_ddf() (restore pre-bias state)
    6. FIFO_LEN x run(SAMPLE_INTERVAL, bias-action) -> save_states
    7. apply_ddf()

    Parameters
    ----------
    pinball_front_x, pinball_rear_x : pinball geometry (L0 units).
        Default 30.0/31.3 for Karman; 19.0/20.3 for Illusion.
    obs_slice_start, obs_slice_end : slice of obs for norm.
        Default [2:14] for Karman (7 objects); [0:12] for Illusion (6 objects).
    n_objects_total : if provided, used for bias array length.
        Default: inferred from flow_field after adding cylinders.

    Returns dict with norm values.
    """
    if action_bias is None:
        action_bias = (0.0, -4.0, 4.0)

    u0_float = float(u0)

    # add 3 pinball cylinders
    ny = flow_field.FIELD_SHAPE[1]
    centers = [
        (pinball_front_x * l0, (ny - 1) / 2, 0.0),
        (pinball_rear_x * l0, (ny - 1) / 2 + 0.75 * l0, 0.0),
        (pinball_rear_x * l0, (ny - 1) / 2 - 0.75 * l0, 0.0),
    ]
    for c in centers:
        flow_field.add_cylinder(c, l0 / 2.0)

    n_obj = flow_field.obs.size // 2 if n_objects_total is None else n_objects_total
    print(f"  bodies after pinball: {n_obj}")

    # stabilize
    stabilize_steps = int(4 * flow_field.FIELD_SHAPE[0] / u0_float)
    print(f"  stabilising pinball ({stabilize_steps} steps)...")
    flow_field.run(stabilize_steps, np.zeros(n_obj, dtype=data_type))

    # checkpoint DDF
    flow_field.get_ddf()
    flow_field.save_ddf()

    # --- norm phase (zero-action) ---
    fifo = deque(maxlen=fifo_len)
    for _ in range(fifo_len):
        flow_field.run(sample_interval, np.zeros(n_obj, dtype=data_type))
        fifo.append(flow_field.obs.copy()[obs_slice_start:obs_slice_end])

    temp_states = np.array(fifo, dtype=data_type)
    # forces are at indices [6:12] relative to the slice end
    force_start = obs_slice_end - obs_slice_start - 6
    force_end = force_start + 6
    force_norm_fact = 6.0 * float(np.max(np.abs(temp_states[:, force_start:force_end])))
    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])))

    print(f"  norm: force_norm_fact={force_norm_fact:.6f}")
    print(f"  norm: sens_deviation={sens_deviation}")
    print(f"  norm: sens_norm_fact={sens_norm_fact}")

    # --- bias-action rollout ---
    flow_field.apply_ddf()
    bias = np.zeros(n_obj, dtype=data_type)
    bias[n_obj - 3] = float(action_bias[0] * u0_float)
    bias[n_obj - 2] = float(action_bias[1] * u0_float)
    bias[n_obj - 1] = float(action_bias[2] * u0_float)
    print(f"  bias action: {bias}")

    fifo.clear()
    for _ in range(fifo_len):
        flow_field.run(sample_interval, bias)
        fifo.append(flow_field.obs.copy()[obs_slice_start:obs_slice_end])

    save_states = np.array(list(fifo), dtype=data_type)
    # CRITICAL: save DDF again AFTER bias FIFO, so restore_ddf() goes
    # to the post-bias state (consistent with saved FIFO).
    # Without this, reset() restores to a bare stabilized state with
    # no bias history, invalidating the FIFO.
    flow_field.get_ddf()
    flow_field.save_ddf()
    flow_field.apply_ddf()

    return {
        "force_norm_fact": force_norm_fact,
        "sens_deviation": sens_deviation.tolist(),
        "sens_norm_fact": sens_norm_fact.tolist(),
        "action_bias": list(action_bias),
        "save_states": save_states,
    }


def build_observation(
    obs_slice: np.ndarray,
    norm: dict,
) -> np.ndarray:
    """Assemble normalised DRL observation (12-dim) from a single obs slice.

    ``obs_slice`` is 12-element: sensor[0:6] + force[6:12].

    Returns clipped 12-dim array in [-1, 1].
    """
    forces = obs_slice[6:12] / norm["force_norm_fact"]
    sens = (obs_slice[0:6] - norm["sens_deviation"]) / norm["sens_norm_fact"]
    obs = np.clip(np.hstack([forces, sens]), -1.0, 1.0).astype(np.float32)
    return obs


def action_to_physical(
    action_norm: np.ndarray,
    *,
    scale: float = 8.0,
    bias: Tuple[float, float, float] = (0.0, -4.0, 4.0),
    u0: float = 0.01,
) -> np.ndarray:
    """Convert normalized action [-1,1] to physical omega (lattice units).

    physical_omega[i] = (action_norm[i] * scale + bias[i]) * u0
    """
    action_norm = np.asarray(action_norm, dtype=np.float64).reshape(-1, 3)
    bias_arr = np.array(bias, dtype=np.float64)
    return (action_norm * scale + bias_arr) * u0


def scale_action(
    action_norm: np.ndarray,
    *,
    scale: float = 8.0,
    bias: Tuple[float, float, float] = (0.0, -4.0, 4.0),
    u0: float = 0.01,
    n_total_bodies: int = 7,
) -> np.ndarray:
    """Convert native [front, upper, lower] action to legacy CFD array.

    Returns array of length *n_total_bodies* with cylinders' omegas at the
    last 3 slots. Historical callers may label upper/lower as top/bottom; this
    function deliberately preserves native object order and never rearranges.
    """
    a = np.zeros(n_total_bodies, dtype=np.float32)
    omega = (np.array(action_norm, dtype=np.float32) * scale + np.array(bias, dtype=np.float32)) * u0
    a[n_total_bodies - 3:] = omega
    return a


# ---------------------------------------------------------------------------
# Vorticity & field export
# ---------------------------------------------------------------------------

def vorticity_from_ddf(flow_field: FlowField, u0: float) -> np.ndarray:
    """Compute z-vorticity from current DDF on host."""
    flow_field.get_ddf()
    ddf = flow_field.ddf.copy().reshape((9, flow_field.FIELD_SHAPE[1],
                                         flow_field.FIELD_SHAPE[0])).transpose(2, 1, 0)
    ux = (ddf[:, :, 1] + ddf[:, :, 5] + ddf[:, :, 8]
          - ddf[:, :, 3] - ddf[:, :, 6] - ddf[:, :, 7]) / u0
    uy = (ddf[:, :, 2] + ddf[:, :, 5] + ddf[:, :, 6]
          - ddf[:, :, 4] - ddf[:, :, 7] - ddf[:, :, 8]) / u0
    omega = np.gradient(uy, axis=0) - np.gradient(ux, axis=1)
    return omega.astype(np.float64)


def save_vorticity_png(path: str, omega: np.ndarray, title: str = ""):
    """Save vorticity field as a PNG with symmetric colour bar."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    abs_o = np.abs(omega[np.isfinite(omega)])
    vmax = float(np.percentile(abs_o, 99.5)) if abs_o.size > 0 else 1.0
    if vmax <= 0:
        vmax = 1.0

    ny, nx = omega.shape
    fig, ax = plt.subplots(figsize=(min(18, max(8, nx / 60)), min(10, max(3, ny / 40))))
    im = ax.imshow(omega, origin="lower", aspect="equal", cmap="RdBu_r",
                   vmin=-vmax, vmax=vmax, extent=(0, nx - 1, 0, ny - 1))
    ax.set_xlabel("x (lattice)")
    ax.set_ylabel("y (lattice)")
    if title:
        ax.set_title(title)
    fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label=r"$\omega_z$")
    fig.tight_layout()
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close(fig)


# ---------------------------------------------------------------------------
# DTW similarity
# ---------------------------------------------------------------------------

LEGACY_DTW_VERSION = "legacy_dtw_v1_abs_n_unclipped"


def calc_lag(target: np.ndarray, state: np.ndarray) -> int:
    """Find lag that maximises cross-correlation between two 1-D signals."""
    t = target - np.mean(target)
    s = state - np.mean(state)
    corr = np.correlate(t, s, mode="full")
    lags = np.arange(-len(target) + 1, len(target))
    return int(lags[np.argmax(corr)])


def calc_dtw_sim(target: np.ndarray, state: np.ndarray) -> float:
    """DTW-based similarity: 1 - (DTW distance / len(target))."""
    n, m = len(target), len(state)
    dtw = np.full((n + 1, m + 1), np.inf)
    dtw[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]))
            dtw[i, j] = cost + min(dtw[i - 1, j], dtw[i, j - 1], dtw[i - 1, j - 1])
    return float(1.0 - dtw[n, m] / n)


def _legacy_scene_similarity(
    target_states: np.ndarray,
    state_series: np.ndarray,
    conv_len: int,
    target_sensor_start: int,
) -> float:
    target = np.asarray(target_states)
    state = np.asarray(state_series)
    target_sensors = target[:, target_sensor_start:target_sensor_start + 6]
    if target_sensors.shape[1] != 6 or state.ndim != 2 or state.shape[1] < 6:
        raise ValueError("legacy DTW requires six target and state sensor channels")
    if len(target) < 2 * conv_len or len(state) < conv_len:
        raise ValueError("legacy DTW requires target >= 2*conv_len and state >= conv_len")

    ref = target_sensors[conv_len:2 * conv_len, 1]
    cur = state[-conv_len:, 1]
    lag = calc_lag(ref, cur)
    return float(sum(
        calc_dtw_sim(
            np.roll(target_sensors[:, i], -lag)[conv_len:2 * conv_len],
            state[-conv_len:, i],
        )
        for i in range(6)
    ) / 6.0)


def compute_similarity_karman_legacy_v1(
    target_states: np.ndarray,
    state_series: np.ndarray,
    conv_len: int = 30,
) -> float:
    """Karman ``legacy_dtw_v1_abs_n_unclipped`` similarity."""
    return _legacy_scene_similarity(target_states, state_series, conv_len, 0)


def compute_similarity_illusion_legacy_v1(
    target_states: np.ndarray,
    state_series: np.ndarray,
    conv_len: int = 36,
) -> float:
    """Illusion legacy similarity; target layout is [fx, fy, six sensors]."""
    return _legacy_scene_similarity(target_states, state_series, conv_len, 2)


def compute_similarity(
    target_states: np.ndarray,
    state_series: np.ndarray,
    conv_len: int,
) -> float:
    """Backward-compatible Karman wrapper for the versioned legacy metric."""
    return compute_similarity_karman_legacy_v1(target_states, state_series, conv_len)


# ---------------------------------------------------------------------------
# Dummy env for loading SB3 models
# ---------------------------------------------------------------------------

def create_dummy_env(s_dim: int = 12, a_dim: int = 3):
    """Return a gym.Env with correct observation/action spaces for model loading."""
    import gymnasium as gym
    from gymnasium import spaces

    class DummyEnv(gym.Env):
        def __init__(self):
            super().__init__()
            self.observation_space = spaces.Box(low=-1, high=1, shape=(s_dim,), dtype=np.float32)
            self.action_space = spaces.Box(low=-1, high=1, shape=(a_dim,), dtype=np.float32)

        def reset(self, seed=None):
            return np.zeros(s_dim, dtype=np.float32), {}

        def step(self, action):
            return np.zeros(s_dim, dtype=np.float32), 0.0, False, False, {}

        def render(self):
            pass

    return DummyEnv()


def load_ppo_model(model_path: str, device: str = "cuda:0", s_dim: int = 12, a_dim: int = 3):
    """Load a PPO model with Sin activation."""
    import torch
    from torch.nn import Module
    from stable_baselines3 import PPO

    class Sin(Module):
        def forward(self, x):
            return torch.sin(x)

    dummy_env = create_dummy_env(s_dim, a_dim)
    model = PPO.load(model_path, env=dummy_env, device=device)
    return model
