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

Copied and adapted from SR_analysis/utils/cfd_interface.py (verified working).
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 utils/__init__.py,
because it requires pycuda. Other utils (resampling) do NOT require pycuda.
"""
from __future__ import annotations

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

import numpy as np

_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)

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

ACTION_SMOOTH_WEIGHT = 0.1  # used by FlowField.run() internally


def load_legacy_configs(config_dir: str) -> Tuple[Any, Any]:
    """Load and return legacy (cuda_config, field_config)."""
    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


# ---------------------------------------------------------------------------
# Karman cloak env builder (disturbance cylinder + 3 sensors)
# ---------------------------------------------------------------------------

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]:
    """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, info_dict).
    """
    cy = (flow_field.FIELD_SHAPE[1] - 1) / 2.0
    flow_field.add_cylinder((10.0 * l0, cy, 0.0), l0)
    for y_off in [2.0, 0.0, -2.0]:
        flow_field.add_sensor((40.0 * l0, cy + y_off * l0, 0.0), l0 / 4.0)

    n_obj = flow_field.obs.size // 2
    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))

    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))
        target_states = np.vstack((target_states, flow_field.obs.copy()[2:8]))

    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]}


# ---------------------------------------------------------------------------
# Pinball adder + norm computation (configurable for Karman / Illusion)
# ---------------------------------------------------------------------------

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, preset-action FIFO init.

    Parameters
    ----------
    pinball_front_x, pinball_rear_x : pinball geometry (L0 units).
    obs_slice_start, obs_slice_end : slice of obs for norm.

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

    u0_float = float(u0)
    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_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))
    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 the last 6 positions of the slice
    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}")

    # ---- preset-action FIFO init ----
    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)

    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)
    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,
    }


# ---------------------------------------------------------------------------
# Observation builder and action helpers
# ---------------------------------------------------------------------------

def build_observation(obs_slice: np.ndarray, norm: dict) -> np.ndarray:
    """Assemble normalised DRL observation (12-dim) from 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 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 normalised action ([-1,1]^3) to legacy CFD action array.

    Returns array of length n_total_bodies with cylinders' omegas at the last 3 slots.
    """
    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 get_velocity_field(flow_field: FlowField, u0: float = 0.01):
    """Extract ux, uy fields from DDF on host. Returns (ux, uy) each (NX, NY)."""
    flow_field.get_ddf()
    NX = flow_field.FIELD_SHAPE[0]
    NY = flow_field.FIELD_SHAPE[1]
    ddf = flow_field.ddf.copy().reshape((9, NY, NX)).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
    return ux.astype(np.float32), uy.astype(np.float32)


def vorticity_from_ddf(flow_field: FlowField, u0: float) -> np.ndarray:
    """Compute z-vorticity from current DDF on host.

    Returns omega_z with shape (NY, NX) for direct imshow use.
    """
    ux, uy = get_velocity_field(flow_field, u0)
    # ux, uy have shape (NX, NY), transpose to (NY, NX) for correct axes
    ux_t = ux.T.astype(np.float64)
    uy_t = uy.T.astype(np.float64)
    # omega_z = duy/dx - dux/dy
    omega = np.gradient(uy_t, axis=1) - np.gradient(ux_t, axis=0)
    return omega


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
# ---------------------------------------------------------------------------

def calc_lag(target: np.ndarray, state: np.ndarray) -> int:
    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:
    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 compute_similarity(target_states: np.ndarray, state_series: np.ndarray, conv_len: int) -> float:
    """Lag-compensated DTW similarity over conv_len window."""
    ref = target_states[conv_len:2 * conv_len, 1]
    cur = state_series[-conv_len:, 1]
    lag = calc_lag(ref, cur)

    sim_sum = 0.0
    for i in range(6):
        target_seq = np.roll(target_states[:, i], -lag)[conv_len:2 * conv_len]
        state_seq = state_series[-conv_len:, i]
        sim_sum += calc_dtw_sim(target_seq, state_seq) / 6.0
    return float(sim_sum)


# ---------------------------------------------------------------------------
# PPO model loading (DummyEnv + Sin activation)
# ---------------------------------------------------------------------------

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
