"""Unified feature builder for all cloak scenes.

Produces dimensionless features with consistent G-equivariant structure.
Active Karman, Illusion, and steady-calibration scenes use this same builder.

Copy of analysis_cloak/common/feature_builder.py -- kept as canonical source.
"""
from __future__ import annotations

from typing import Dict, List, Optional, Tuple

import numpy as np

from .g_operator import apply_G_alpha


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

    Native sensor order is ``[upper_ux, upper_uy, front_ux, front_uy,
    lower_ux, lower_uy]``.  The historical feature aliases are intentionally
    preserved: ``T=upper``, ``C=front``, and ``B=lower``.

    Native force order is ``[front_fx, front_fy, upper_fx, upper_fy,
    lower_fx, lower_fy]`` because LegacyCelerisLab force slots follow object
    insertion order. Historical coefficient aliases remain ``F=front``,
    ``T=upper``, and ``B=lower``.

    Returns dimensionless velocity and force-coefficient arrays under those
    historical ``T/C/B`` and ``F/T/B`` keys.
    """
    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_T": 2.0 * f[:, 2] / (rho * u0**2 * d),
        "Cl_T": 2.0 * f[:, 3] / (rho * u0**2 * d),
        "Cd_B": 2.0 * f[:, 4] / (rho * u0**2 * d),
        "Cl_B": 2.0 * f[:, 5] / (rho * u0**2 * d),
    }


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

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

RAW_COMPLETE_KEYS = [
    "u_hat_T", "v_hat_T", "u_hat_C", "v_hat_C", "u_hat_B", "v_hat_B",
    "Cd_F", "Cl_F", "Cd_T", "Cl_T", "Cd_B", "Cl_B",
]

SYMMETRY_KEYS = [
    "u_s", "u_a", "u_c", "v_s", "v_a", "v_c",
    "Cd_F", "Cd_rear_s", "Cd_rear_a", "Cd_tot",
    "Cl_F", "Cl_rear_s", "Cl_rear_a", "Cl_tot",
]

PHYSICS_REDUCED_KEYS = [
    "u_m", "u_a", "u_c", "v_a",
    "Cd_tot", "Cd_rear", "Cl_tot", "Cl_diff",
]

FEATURE_SET_KEYS = {
    "raw_complete": tuple(RAW_COMPLETE_KEYS),
    "symmetry": tuple(SYMMETRY_KEYS),
    "physics_reduced": tuple(PHYSICS_REDUCED_KEYS),
}

# Original feature set (includes sin_ua/cos_ua)
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",
]

# V2 core features: no sin_ua/cos_ua, no mu (for single-scene fitting)
CORE_FEAT_KEYS_V2 = [
    "u_m", "u_a", "u_c",
    "v_a",
    "Cd_tot", "Cd_rear",
    "Cl_tot", "Cl_diff",
    "aF_lag1", "aB_lag1", "aT_lag1",
    "daF", "daB", "daT",
]

# Time-explicit features: da/dt_c (for cross-scene coefficient comparison)
TIME_FEAT_KEYS = [
    "daF_dt", "daB_dt", "daT_dt",
]

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

# Illusion target force features (for scenes where we know the target Cd/Cl)
ILLUSION_TARGET_KEYS = ["target_Cd", "target_Cl"]
ILLUSION_FEAT_KEYS_V2 = CORE_FEAT_KEYS_V2 + ILLUSION_TARGET_KEYS
ILLUSION_ALL_FEAT_KEYS_V2 = ILLUSION_FEAT_KEYS_V2 + MU_FEAT_KEYS  # with mu for joint

# Physics-only feature keys: NO action history terms (no aF_lag1, no daF, etc.)
# These are the clean inputs for learning d(alpha)/dt as a function of physics state.
PHYSICS_FEAT_KEYS = [
    "u_m", "u_a", "u_c", "v_a",
    "Cd_tot", "Cd_rear", "Cl_tot", "Cl_diff",
]

# Illusion error-state: physics features + force error (not raw target forces)
ILLUSION_ERR_KEYS = PHYSICS_FEAT_KEYS + ["Cd_err", "Cl_err"]

# Observation lag (1-step delayed) and derivative (time-normalized) features
# These add temporal/phasing information to the otherwise static physics features.
LAG_FEAT_KEYS = [
    "u_m_lag1", "u_a_lag1", "u_c_lag1", "v_a_lag1",
    "Cd_tot_lag1", "Cd_rear_lag1", "Cl_tot_lag1", "Cl_diff_lag1",
]
DERIV_FEAT_KEYS = [
    "du_m_dt", "du_a_dt", "du_c_dt", "dv_a_dt",
    "dCd_tot_dt", "dCd_rear_dt", "dCl_tot_dt", "dCl_diff_dt",
]
# Action lag (for ablation level 4 — comparing against old approach)
ACTION_LAG_KEYS = ["aF_lag1", "aB_lag1", "aT_lag1"]

# Augmented levels for ablation study:
# Level 0 = x_n (PHYSICS_FEAT_KEYS, no memory at all)
# Level 1 = x_n + x_{n-1} (current + 1-step lag)
# Level 2 = x_n + dx/dt (current + derivative)
# Level 3 = x_n + x_{n-1} + dx/dt (full temporal context)
# Level 4 = Level 3 + a_{n-1} (add action history for comparison)
AUG_LEVEL_1_KEYS = PHYSICS_FEAT_KEYS + LAG_FEAT_KEYS
AUG_LEVEL_2_KEYS = PHYSICS_FEAT_KEYS + DERIV_FEAT_KEYS
AUG_LEVEL_3_KEYS = PHYSICS_FEAT_KEYS + LAG_FEAT_KEYS + DERIV_FEAT_KEYS
AUG_LEVEL_4_KEYS = AUG_LEVEL_3_KEYS + ACTION_LAG_KEYS

# Illusion equivalents with error-state
ILLUSION_LAG_KEYS = [
    "Cd_err_lag1", "Cl_err_lag1",
]
ILLUSION_DERIV_KEYS = [
    "dCd_err_dt", "dCl_err_dt",
]
ILLUSION_AUG_LEVEL_1_KEYS = ILLUSION_ERR_KEYS + LAG_FEAT_KEYS + ILLUSION_LAG_KEYS
ILLUSION_AUG_LEVEL_2_KEYS = ILLUSION_ERR_KEYS + DERIV_FEAT_KEYS + ILLUSION_DERIV_KEYS
ILLUSION_AUG_LEVEL_3_KEYS = ILLUSION_AUG_LEVEL_1_KEYS + DERIV_FEAT_KEYS + ILLUSION_DERIV_KEYS
ILLUSION_AUG_LEVEL_4_KEYS = ILLUSION_AUG_LEVEL_3_KEYS + ACTION_LAG_KEYS

# Phase-state features: low-dimensional dynamic state z = [phase_obs, phase_deriv, static_force]
# This replaces the full x_n + x_{n-1} with a compact representation.
# The idea: u_a + du_a/dt encodes oscillation phase, Cl_tot + dCl_tot/dt encodes lift dynamics,
# and Cd_tot/Cd_rear provide static force feedback.
PHASE_STATE_KEYS = [
    "u_a", "du_a_dt",        # oscillation phase (cross-stream asymmetry + rate)
    "Cl_tot", "dCl_tot_dt",   # lift dynamics
    "Cd_tot", "Cd_rear",      # static drag feedback
]

# Karman expanded: phase-state + supplementary static quantities
# For Karman, 6-dim phase-state may not capture enough information.
# Adding u_m (mean streamwise), u_c (center sensor), v_a (cross asymmetry), Cl_diff (lift distribution)
KARMAN_EXPANDED_KEYS = PHASE_STATE_KEYS + [
    "u_m", "u_c", "v_a", "Cl_diff",
]

# Illusion phase-state: error-based + dynamics
ILLUSION_PHASE_KEYS = PHASE_STATE_KEYS + [
    "Cd_err", "Cl_err", "dCd_err_dt", "dCl_err_dt",
]

ALL_FEAT_KEYS = CORE_FEAT_KEYS + MU_FEAT_KEYS
# V2 all features (for cross-Re joint fitting with mu)
ALL_FEAT_KEYS_V2 = CORE_FEAT_KEYS_V2 + 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,
    include_cos_sin: bool = True,  # if True, include sin_ua/cos_ua features
    dt_c: float = 1.0,         # control interval in T0 units (for time normalization)
    u0: float = U0,            # inlet velocity for omega->alpha conversion
    target_forces: Optional[np.ndarray] = None,  # (T, 2) raw lattice target forces [fx,fy]
    rho: float = 1.0,          # fluid density for Cd/Cl conversion
    sensors_raw: Optional[np.ndarray] = None,  # (T, 6) raw lattice sensors, for obs dynamics
    forces_raw: Optional[np.ndarray] = None,   # (T, 6) raw lattice forces, for obs dynamics
) -> 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
        include_cos_sin: include sin_ua/cos_ua phase encoding
        dt_c: control interval (in T0 = D/U0 units), for time-normalized deltas
        u0: inlet velocity (lattice), used only when alpha_mode=False

    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 = {key: np.asarray(value, dtype=np.float64).copy() for key, value in dim.items()}

    # Complete symmetric/antisymmetric coordinates. Under G, u_a/v_s/v_c,
    # Cd_rear_a, Cl_F/Cl_rear_s/Cl_tot are odd; the complementary coordinates
    # are even.
    sym["u_s"] = (u_T + u_B) / 2.0
    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_s"] = (v_T + v_B) / 2.0
    sym["v_a"] = (v_T - v_B) / 2.0
    sym["v_c"] = v_C.copy()

    sym["Cd_tot"] = Cd_F + Cd_T + Cd_B
    sym["Cd_rear"] = Cd_T + Cd_B
    sym["Cd_rear_s"] = (Cd_T + Cd_B) / 2.0
    sym["Cd_rear_a"] = (Cd_T - Cd_B) / 2.0
    sym["Cl_tot"] = Cl_F + Cl_T + Cl_B
    sym["Cl_diff"] = Cl_T - Cl_B
    sym["Cl_rear_s"] = (Cl_T + Cl_B) / 2.0
    sym["Cl_rear_a"] = (Cl_T - Cl_B) / 2.0

    # Phase (optional, may obscure linear structure)
    if include_cos_sin:
        sym["sin_ua"] = np.sin(np.pi * sym["u_a"])
        sym["cos_ua"] = np.cos(np.pi * sym["u_a"])

    # Memory (nondim alpha) -- discrete version
    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]

    # Time-normalized deltas (for cross-scene coefficient comparison)
    sym["daF_dt"] = sym["daF"] / dt_c
    sym["daB_dt"] = sym["daB"] / dt_c
    sym["daT_dt"] = sym["daT"] / dt_c

    # Target forces (for illusion scenes) — convert lattice forces to Cd/Cl
    if target_forces is not None:
        tf = np.asarray(target_forces, dtype=np.float64)
        if tf.ndim == 1:
            tf = tf.reshape(1, -1)
        sym["target_Cd"] = 2.0 * tf[:, 0] / (rho * u0**2 * D_CYL)
        sym["target_Cl"] = 2.0 * tf[:, 1] / (rho * u0**2 * D_CYL)
        # Error-state: deviation between actual and target force
        sym["Cd_err"] = sym["Cd_tot"] - sym["target_Cd"]
        sym["Cl_err"] = sym["Cl_tot"] - sym["target_Cl"]

    # 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
        sym["mu_Cl_tot"] = sym["Cl_tot"] * mu  # additional for phase-state

    # Observation dynamics: 1-step lag + time-normalized derivative
    # These add temporal/phase info that static features lack.
    if sensors_raw is not None and forces_raw is not None:
        sr = np.asarray(sensors_raw, dtype=np.float64)
        fr = np.asarray(forces_raw, dtype=np.float64)
        # Lag-1 observations: first-step history is the current observation.
        # This makes first-step derivatives zero and matches online initialization.
        s_lag1 = np.empty_like(sr)
        f_lag1 = np.empty_like(fr)
        s_lag1[0] = sr[0]
        f_lag1[0] = fr[0]
        s_lag1[1:] = sr[:-1]
        f_lag1[1:] = fr[:-1]
        # Compute dim for lag-1
        dim_lag1 = compute_dimensionless(s_lag1, f_lag1, u0=u0, d=D_CYL, rho=rho)
        
        # Compute combined features for lag-1
        uB_1, uC_1, uT_1 = dim_lag1["u_hat_B"], dim_lag1["u_hat_C"], dim_lag1["u_hat_T"]
        vB_1, vC_1, vT_1 = dim_lag1["v_hat_B"], dim_lag1["v_hat_C"], dim_lag1["v_hat_T"]
        CdF_1, CdT_1, CdB_1 = dim_lag1["Cd_F"], dim_lag1["Cd_T"], dim_lag1["Cd_B"]
        ClF_1, ClT_1, ClB_1 = dim_lag1["Cl_F"], dim_lag1["Cl_T"], dim_lag1["Cl_B"]
        
        u_m_1 = (uB_1 + uC_1 + uT_1) / 3.0
        u_a_1 = (uT_1 - uB_1) / 2.0
        u_c_1 = uC_1.copy()
        v_a_1 = (vT_1 - vB_1) / 2.0
        Cd_tot_1 = CdF_1 + CdT_1 + CdB_1
        Cd_rear_1 = CdT_1 + CdB_1
        Cl_tot_1 = ClF_1 + ClT_1 + ClB_1
        Cl_diff_1 = ClT_1 - ClB_1
        
        # Trim to T (in case raw arrays are longer — e.g. validator passes 2 rows)
        def _trim(x):
            return x[-T:] if x.shape[0] > T else x
        
        # Store lag-1 features (trimmed to T)
        sym["u_m_lag1"] = _trim(u_m_1)
        sym["u_a_lag1"] = _trim(u_a_1)
        sym["u_c_lag1"] = _trim(u_c_1)
        sym["v_a_lag1"] = _trim(v_a_1)
        sym["Cd_tot_lag1"] = _trim(Cd_tot_1)
        sym["Cd_rear_lag1"] = _trim(Cd_rear_1)
        sym["Cl_tot_lag1"] = _trim(Cl_tot_1)
        sym["Cl_diff_lag1"] = _trim(Cl_diff_1)
        
        # Time-normalized derivatives: (current - lag1) / dt_c
        eps = 1e-12
        sym["du_m_dt"] = _trim((sym["u_m"] - u_m_1) / (dt_c + eps))
        sym["du_a_dt"] = _trim((sym["u_a"] - u_a_1) / (dt_c + eps))
        sym["du_c_dt"] = _trim((sym["u_c"] - u_c_1) / (dt_c + eps))
        sym["dv_a_dt"] = _trim((sym["v_a"] - v_a_1) / (dt_c + eps))
        sym["dCd_tot_dt"] = _trim((sym["Cd_tot"] - Cd_tot_1) / (dt_c + eps))
        sym["dCd_rear_dt"] = _trim((sym["Cd_rear"] - Cd_rear_1) / (dt_c + eps))
        sym["dCl_tot_dt"] = _trim((sym["Cl_tot"] - Cl_tot_1) / (dt_c + eps))
        sym["dCl_diff_dt"] = _trim((sym["Cl_diff"] - Cl_diff_1) / (dt_c + eps))
        
        # Illusion error-state dynamics
        if target_forces is not None:
            tf = np.asarray(target_forces, dtype=np.float64)
            tf_lag1 = np.empty_like(tf)
            tf_lag1[0] = tf[0]
            tf_lag1[1:] = tf[:-1]
            tCd_1 = 2.0 * tf_lag1[:, 0] / (rho * u0**2 * D_CYL)
            tCl_1 = 2.0 * tf_lag1[:, 1] / (rho * u0**2 * D_CYL)
            sym["Cd_err_lag1"] = _trim(Cd_tot_1 - tCd_1)
            sym["Cl_err_lag1"] = _trim(Cl_tot_1 - tCl_1)
            sym["dCd_err_dt"] = _trim((sym["Cd_err"] - sym["Cd_err_lag1"]) / (dt_c + eps))
            sym["dCl_err_dt"] = _trim((sym["Cl_err"] - sym["Cl_err_lag1"]) / (dt_c + eps))

    return sym


def build_feature_matrix(
    sym: Dict[str, np.ndarray],
    feat_keys: List[str],
    add_bias: bool = True,
    strict: bool = True,
) -> np.ndarray:
    """Build a feature matrix, failing on missing features by default.

    Set ``strict=False`` to retain the historical missing-key-as-zero behavior.
    """
    if not feat_keys:
        raise ValueError("feat_keys must contain at least one feature")
    missing = [key for key in feat_keys if key not in sym]
    if strict and missing:
        raise KeyError("missing required feature(s): {}".format(", ".join(missing)))

    available = next((np.asarray(sym[key]) for key in feat_keys if key in sym), None)
    if available is None:
        available = np.asarray(sym.get("u_m", np.ones(1)))
    T = available.shape[0]
    cols = []
    if add_bias:
        cols.append(np.ones(T, dtype=np.float64))
    for key in feat_keys:
        col = np.asarray(sym[key], dtype=np.float64) if key in sym else np.zeros(T, dtype=np.float64)
        if col.ndim != 1 or col.shape[0] != T:
            raise ValueError("feature {!r} must have shape ({},), got {}".format(key, T, col.shape))
        cols.append(col)
    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
