"""Unified feature builder for all cloak scenes.

Produces dimensionless features with consistent G-equivariant structure.
All scenes (Karman, steady, vortex, erase) use this same builder.
"""
from __future__ import annotations

from typing import Dict, List, Optional, Tuple

import numpy as np


# -- Physical constants ------------------------------------------------------
U0 = 0.01       # inlet velocity (lattice units)
D_CYL = 20.0    # cylinder diameter (lattice)


# -- Dimensionless conversion ------------------------------------------------

def compute_dimensionless(
    sensors: np.ndarray,  # (T, 6) raw lattice [s0_ux,s0_uy, s1_ux,s1_uy, s2_ux,s2_uy]
    forces: np.ndarray,   # (T, 6) raw lattice [f0_fx,f0_fy, f1_fx,f1_fy, f2_fx,f2_fy]
    u0: float = U0,
    d: float = D_CYL,
    rho: float = 1.0,
) -> Dict[str, np.ndarray]:
    """Convert raw lattice CFD data to dimensionless quantities.

    Sensor order: [s0_ux,s0_uy, s1_ux,s1_uy, s2_ux,s2_uy]
        where s0=top(y=+2L0), s1=mid(y=0), s2=bottom(y=-2L0)
    Force order: [front_fx,front_fy, bottom_fx,bottom_fy, top_fx,top_fy]

    Returns:
        u_hat_B, u_hat_C, u_hat_T: nondim streamwise velocity (bottom/centre/top)
        v_hat_B, v_hat_C, v_hat_T: nondim crosswise velocity
        Cd_F, Cd_T, Cd_B: drag coefficient per cylinder
        Cl_F, Cl_T, Cl_B: lift coefficient per cylinder
    """
    s = np.asarray(sensors, dtype=np.float64)
    f = np.asarray(forces, dtype=np.float64)

    # Sensor positions: s0=top, s1=centre, s2=bottom
    # Convention: B=bottom=s2, C=centre=s1, T=top=s0
    return {
        "u_hat_T": s[:, 0] / u0,
        "v_hat_T": s[:, 1] / u0,
        "u_hat_C": s[:, 2] / u0,
        "v_hat_C": s[:, 3] / u0,
        "u_hat_B": s[:, 4] / u0,
        "v_hat_B": s[:, 5] / u0,
        "Cd_F": 2.0 * f[:, 0] / (rho * u0**2 * d),
        "Cl_F": 2.0 * f[:, 1] / (rho * u0**2 * d),
        "Cd_B": 2.0 * f[:, 2] / (rho * u0**2 * d),
        "Cl_B": 2.0 * f[:, 3] / (rho * u0**2 * d),
        "Cd_T": 2.0 * f[:, 4] / (rho * u0**2 * d),
        "Cl_T": 2.0 * f[:, 5] / (rho * u0**2 * d),
    }


# -- G operator (corrected) --------------------------------------------------

def apply_G_alpha(alpha: np.ndarray) -> np.ndarray:
    """Apply mirror G to action: [aF, aT, aB] -> [-aF, -aB, -aT]."""
    return np.array([-alpha[0], -alpha[2], -alpha[1]], dtype=alpha.dtype)


def apply_G_x(dim: Dict[str, np.ndarray],
              a_prev: np.ndarray,
              a_prev2: np.ndarray) -> Tuple[Dict, np.ndarray, np.ndarray]:
    """Apply G to dimensionless state.

    Returns (G_dim, G_a_prev, G_a_prev2) with corrected sign rules.
    """
    G_dim = {
        "u_hat_B": dim["u_hat_T"], "u_hat_C": dim["u_hat_C"], "u_hat_T": dim["u_hat_B"],
        "v_hat_B": -dim["v_hat_T"], "v_hat_C": -dim["v_hat_C"], "v_hat_T": -dim["v_hat_B"],
        "Cd_F": dim["Cd_F"], "Cd_T": dim["Cd_B"], "Cd_B": dim["Cd_T"],
        "Cl_F": -dim["Cl_F"], "Cl_T": -dim["Cl_B"], "Cl_B": -dim["Cl_T"],
    }
    G_a_prev = np.column_stack([-a_prev[:, 0], -a_prev[:, 2], -a_prev[:, 1]])
    G_a_prev2 = np.column_stack([-a_prev2[:, 0], -a_prev2[:, 2], -a_prev2[:, 1]])
    return G_dim, G_a_prev, G_a_prev2


# -- Feature key definitions -------------------------------------------------

CORE_FEAT_KEYS = [
    "u_m", "u_a", "u_c",
    "v_a",
    "Cd_tot", "Cd_rear",
    "Cl_tot", "Cl_diff",
    "sin_ua", "cos_ua",
    "aF_lag1", "aB_lag1", "aT_lag1",
    "daF", "daB", "daT",
]

MU_FEAT_KEYS = ["mu", "mu_u_a", "mu_v_a", "mu_Cd_tot", "mu_Cl_diff"]

ALL_FEAT_KEYS = CORE_FEAT_KEYS + MU_FEAT_KEYS


# -- Feature computation -----------------------------------------------------

def compute_features(
    dim: Dict[str, np.ndarray],
    actions_prev: np.ndarray,   # (T, 3) physical omega(t-1) or nondim alpha(t-1)
    actions_prev2: np.ndarray,  # (T, 3) physical omega(t-2)
    mu: float,
    alpha_mode: bool = False,  # if True, actions_prev are already nondim alpha
    include_mu: bool = True,
) -> Dict[str, np.ndarray]:
    """Compute unified feature dictionary from dimensionless primitives.

    Args:
        dim: from compute_dimensionless()
        actions_prev: lagged actions (physical omega or nondim alpha)
        actions_prev2: twice-lagged actions
        mu: 1/Re_D
        alpha_mode: if True, actions are already nondim; else convert
        include_mu: include mu modulation terms

    Returns dict with all features as (T,) or (T,3) arrays.
    """
    T = actions_prev.shape[0]
    u_B, u_C, u_T = dim["u_hat_B"], dim["u_hat_C"], dim["u_hat_T"]
    v_B, v_C, v_T = dim["v_hat_B"], dim["v_hat_C"], dim["v_hat_T"]
    Cd_F, Cd_T, Cd_B = dim["Cd_F"], dim["Cd_T"], dim["Cd_B"]
    Cl_F, Cl_T, Cl_B = dim["Cl_F"], dim["Cl_T"], dim["Cl_B"]

    # If actions are in physical omega, convert to nondim alpha
    if alpha_mode:
        a = actions_prev.astype(np.float64)
        a2 = actions_prev2.astype(np.float64)
    else:
        a = actions_prev.astype(np.float64) / U0
        a2 = actions_prev2.astype(np.float64) / U0

    sym = {}

    # Sensor combinations (nondim)
    sym["u_m"] = (u_B + u_C + u_T) / 3.0
    sym["u_a"] = (u_T - u_B) / 2.0
    sym["u_c"] = u_C.copy()
    sym["v_a"] = (v_T - v_B) / 2.0

    # Force combinations (dimensionless Cd/Cl)
    sym["Cd_tot"] = Cd_F + Cd_T + Cd_B
    sym["Cd_rear"] = Cd_T + Cd_B
    sym["Cl_tot"] = Cl_F + Cl_T + Cl_B
    sym["Cl_diff"] = Cl_T - Cl_B

    # Phase
    sym["sin_ua"] = np.sin(np.pi * sym["u_a"])
    sym["cos_ua"] = np.cos(np.pi * sym["u_a"])

    # Memory (nondim alpha)
    sym["aF_lag1"] = a[:, 0]
    sym["aB_lag1"] = a[:, 1]
    sym["aT_lag1"] = a[:, 2]
    sym["daF"] = a[:, 0] - a2[:, 0]
    sym["daB"] = a[:, 1] - a2[:, 1]
    sym["daT"] = a[:, 2] - a2[:, 2]

    # Mu modulation
    if include_mu:
        sym["mu"] = np.full(T, mu, dtype=np.float64)
        sym["mu_u_a"] = sym["u_a"] * mu
        sym["mu_v_a"] = sym["v_a"] * mu
        sym["mu_Cd_tot"] = sym["Cd_tot"] * mu
        sym["mu_Cl_diff"] = sym["Cl_diff"] * mu

    return sym


def build_feature_matrix(
    sym: Dict[str, np.ndarray],
    feat_keys: List[str],
    add_bias: bool = True,
) -> np.ndarray:
    """Build feature matrix (T, N) from symbol dict."""
    cols = []
    if add_bias:
        cols.append(np.ones(sym[feat_keys[0]].shape[0], dtype=np.float64))
    for k in feat_keys:
        if k in sym:
            cols.append(np.asarray(sym[k], dtype=np.float64))
        else:
            # Missing key (e.g. mu terms when include_mu=False) -> zero
            T = sym.get("u_m", np.ones(1)).shape[0]
            cols.append(np.zeros(T, dtype=np.float64))
    return np.column_stack(cols)


def get_feature_names(feat_keys: List[str], add_bias: bool = True) -> List[str]:
    """Get feature names matching build_feature_matrix output."""
    names = []
    if add_bias:
        names.append("bias")
    names.extend(feat_keys)
    return names


# -- Scene metadata ----------------------------------------------------------

def get_scene_metadata(scene: str) -> dict:
    """Return default metadata for a cloak scene."""
    meta = {
        "karman": {"scene_id": "karman", "control_interval": 800, "target_type": "periodic"},
        "steady": {"scene_id": "steady", "control_interval": 800, "target_type": "steady"},
        "vortex_lamb": {"scene_id": "vortex_lamb", "control_interval": 800, "target_type": "transient"},
        "vortex_taylor": {"scene_id": "vortex_taylor", "control_interval": 800, "target_type": "transient"},
        "erase": {"scene_id": "erase", "control_interval": 600, "target_type": "periodic"},
    }
    return meta.get(scene, {"scene_id": scene, "control_interval": 800, "target_type": "unknown"})