"""Phase 3: closed-loop for ablation modes.

Usage::

    conda run -n pycuda_3_10 python phase3_ablation_val.py \\
        --ablation-json output/analysis_crossre/sindy/ablation_results.json \\
        --mode v21 --validate-re 70 --device 2
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from collections import deque

import numpy as np

_PROJ = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
if _PROJ not in sys.path:
    sys.path.insert(0, _PROJ)
from LegacyCelerisLab import FlowField

from utils import (
    nu_from_re, load_legacy_configs, build_karman_cloak_env, add_pinball,
    build_observation, scale_action, action_to_physical,
    compute_dimensionless, compute_physical_symbols,
    save_vorticity_png, vorticity_from_ddf, compute_similarity,
)
from cfg import (
    CONFIG_DIR, OUTPUT_DIR, SAMPLE_INTERVAL, FIFO_LEN, CONV_LEN,
    S_DIM, ACTION_SCALE, ACTION_BIAS, U0,
)

DATA_TYPE = np.float32


def load_ablation_coef(ablation_path, mode, channels_to_load=("front", "bottom", "top")):
    """Load coefficients for a specific ablation mode."""
    with open(ablation_path) as f:
        data = json.load(f)
    
    mode_data = data[mode]
    coefs = {}
    for ch in mode_data["channels"]:
        name = ch["cylinder"]
        if name in channels_to_load:
            coefs[name] = {
                "coef": np.array(ch["best_coef"], dtype=np.float64),
                "has_bias": ch["has_bias"],
            }
    return coefs


def predict_ablation(obs_slice, actions_prev, actions_prev2, coefs, mu, mode, u0=0.01):
    """Predict action using ablation mode coefficients.

    obs_slice: (12,) raw lattice [sensor(6), force(6)]
    """
    sensors = obs_slice[0:6].astype(np.float64).reshape(1, 6)
    forces = obs_slice[6:12].astype(np.float64).reshape(1, 6)
    a_prev = actions_prev.astype(np.float64).reshape(1, 3)
    a_prev2 = actions_prev2.astype(np.float64).reshape(1, 3)

    is_dim = "v22" in mode or "v23" in mode or "v24" in mode or mode in ("v2_dimless",)

    if is_dim:
        dim = compute_dimensionless(sensors, forces, u0=u0, d=20.0)
        # Build dimensionless features
        sym = _build_dimensionless_features(dim, a_prev[0], a_prev2[0], mu)
        Y_scale = 1.0 / u0  # predict nondim alpha = omega / U0
    else:
        # Lattice features (v2/v21)
        a_prev_f = np.zeros((1, 3), dtype=np.float64)
        a_prev2_f = np.zeros((1, 3), dtype=np.float64)
        a_prev_f[0] = a_prev
        a_prev2_f[0] = a_prev2
        sym = _build_lattice_features(sensors, forces, a_prev_f, a_prev2_f, mu)
        Y_scale = 1.0  # predict raw omega

    # Build feature vector
    feat_keys = [k for k in sym.keys() if k not in ("mu", "mu_u_a", "mu_v_a", "mu_Cd_tot", "mu_Cl_diff")]
    feat_keys_mu = ["mu", "mu_u_a", "mu_v_a", "mu_Cd_tot", "mu_Cl_diff"]
    
    feats = []
    for k in feat_keys:
        feats.append(float(sym[k][0]) if isinstance(sym[k], np.ndarray) else float(sym[k]))
    if mu > 0:
        for k in feat_keys_mu:
            feats.append(float(sym[k][0]) if isinstance(sym[k], np.ndarray) else float(sym[k]))

    omega = np.zeros(3, dtype=np.float64)
    for ci, name in enumerate(["front", "bottom", "top"]):
        c = coefs.get(name)
        if c is None:
            continue
        coef_arr = c["coef"]
        has_bias = c["has_bias"]
        
        if has_bias:
            feat_vec = np.array([1.0] + feats) if len(coef_arr) == len(feats) + 1 else np.array(feats)
        else:
            feat_vec = np.array(feats) if len(coef_arr) == len(feats) else np.array([1.0] + feats)
        
        if len(feat_vec) != len(coef_arr):
            feat_vec = np.array(feats)  # fallback
        
        pred = float(feat_vec @ coef_arr) * Y_scale
        omega[ci] = pred
    
    return omega


def _build_lattice_features(sensors, forces, a_prev, a_prev2, mu):
    """Build v2-style lattice features. All args 2D (1, N)."""
    # Ensure 2D
    if sensors.ndim == 1:
        sensors = sensors.reshape(1, -1)
    if forces.ndim == 1:
        forces = forces.reshape(1, -1)
    if a_prev.ndim == 1:
        a_prev = a_prev.reshape(1, -1)
    if a_prev2.ndim == 1:
        a_prev2 = a_prev2.reshape(1, -1)
    sym = compute_physical_symbols(sensors, forces, a_prev, a_prev2)
    sym["mu"] = np.array([mu])
    sym["mu_u_a"] = sym["u_a"] * mu
    sym["mu_v_a"] = sym["v_a"] * mu
    sym["mu_Cd_tot"] = sym["Fx_tot"] * mu
    sym["mu_Cl_diff"] = sym["Fy_diff"] * mu
    return sym


def _build_dimensionless_features(dim, a_prev, a_prev2, mu):
    """Build dimensionless features. a_prev/a_prev2 are 1D (3,) arrays."""
    if a_prev.ndim > 1:
        a_prev = a_prev.flatten()
    if a_prev2.ndim > 1:
        a_prev2 = a_prev2.flatten()
    """Build dimensionless features."""
    T = 1
    u_B, u_C, u_T = dim["u_hat_B"][0], dim["u_hat_C"][0], dim["u_hat_T"][0]
    v_B, v_C, v_T = dim["v_hat_B"][0], dim["v_hat_C"][0], dim["v_hat_T"][0]
    
    sym = {}
    sym["u_m"] = np.array([(u_B + u_C + u_T) / 3.0])
    sym["u_a"] = np.array([(u_T - u_B) / 2.0])
    sym["u_c"] = np.array([u_C])
    sym["u_curv"] = np.array([u_B - 2.0*u_C + u_T])
    sym["v_a"] = np.array([(v_T - v_B) / 2.0])
    sym["v_curv"] = np.array([v_B - 2.0*v_C + v_T])
    sym["sin_ua"] = np.sin(np.pi * sym["u_a"])
    sym["cos_ua"] = np.cos(np.pi * sym["u_a"])
    
    sym["Cd_tot"] = np.array([dim["Cd_F"][0] + dim["Cd_T"][0] + dim["Cd_B"][0]])
    sym["Cd_rear"] = np.array([dim["Cd_T"][0] + dim["Cd_B"][0]])
    sym["Cd_diff"] = np.array([dim["Cd_T"][0] - dim["Cd_B"][0]])
    sym["Cl_tot"] = np.array([dim["Cl_F"][0] + dim["Cl_T"][0] + dim["Cl_B"][0]])
    sym["Cl_diff"] = np.array([dim["Cl_T"][0] - dim["Cl_B"][0]])
    
    # Nondim actions
    a_prev_n = a_prev / U0
    a_prev2_n = a_prev2 / U0
    sym["a0_lag1"] = np.array([a_prev_n[0]])
    sym["a1_lag1"] = np.array([a_prev_n[1]])
    sym["a2_lag1"] = np.array([a_prev_n[2]])
    sym["da0"] = np.array([a_prev_n[0] - a_prev2_n[0]])
    sym["da1"] = np.array([a_prev_n[1] - a_prev2_n[1]])
    sym["da2"] = np.array([a_prev_n[2] - a_prev2_n[2]])
    
    sym["mu"] = np.array([mu])
    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 run_closed_loop(re_code, coefs, mode, device_id, output_root, n_steps=100):
    """Run closed-loop for one ablation mode."""
    os.makedirs(output_root, exist_ok=True)
    nu = nu_from_re(re_code, u0=U0)
    mu = 2.0 / re_code
    
    cuda_cfg, field_cfg = load_legacy_configs(CONFIG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(nu))
    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
    
    target_states, _ = build_karman_cloak_env(
        ff, u0=U0, l0=20.0, sample_interval=SAMPLE_INTERVAL,
        fifo_len=FIFO_LEN, data_type=DATA_TYPE)
    norm = add_pinball(
        ff, l0=20.0, u0=U0, sample_interval=SAMPLE_INTERVAL,
        fifo_len=FIFO_LEN, data_type=DATA_TYPE, action_bias=ACTION_BIAS)
    
    np.savez(os.path.join(output_root, "target.npz"), target_states=target_states)
    
    # Controlled rollout
    ff.restore_ddf()
    ff.apply_ddf()
    fifo = deque(maxlen=FIFO_LEN)
    bias_action = scale_action(
        np.zeros(3, dtype=np.float32), scale=ACTION_SCALE,
        bias=ACTION_BIAS, u0=U0, n_total_bodies=7)
    for _ in range(FIFO_LEN):
        ff.run(SAMPLE_INTERVAL, bias_action)
        fifo.append(ff.obs.copy()[2:14])
    
    sens_sc = []
    actions_prev = action_to_physical(
        np.zeros((1,3), dtype=np.float32), scale=ACTION_SCALE,
        bias=ACTION_BIAS, u0=U0).flatten()
    actions_prev2 = actions_prev.copy()
    
    for step in range(n_steps):
        obs_slice = fifo[-1] if len(fifo) > 0 else np.zeros(12, dtype=np.float32)
        omega_pred = predict_ablation(obs_slice, actions_prev, actions_prev2, coefs, mu, mode, u0=U0)
        
        norm_action = (omega_pred / U0 - np.array(ACTION_BIAS, dtype=np.float64)) / ACTION_SCALE
        norm_action = np.clip(norm_action, -1.0, 1.0).astype(np.float32)
        action_arr = scale_action(
            norm_action, scale=ACTION_SCALE, bias=ACTION_BIAS, u0=U0, n_total_bodies=7)
        ff.run(SAMPLE_INTERVAL, action_arr)
        
        obs_slice_new = ff.obs.copy()[2:14]
        fifo.append(obs_slice_new)
        sens_sc.append(obs_slice_new[0:6])
        actions_prev2 = actions_prev.copy()
        actions_prev = omega_pred.copy()
    
    sens_arr = np.array(sens_sc, dtype=np.float32)
    sim = compute_similarity(target_states, sens_arr, CONV_LEN)
    
    omega_vort = vorticity_from_ddf(ff, u0=U0)
    save_vorticity_png(os.path.join(output_root, f"vorticity_{mode}.png"),
                       omega_vort, title=f"Re{re_code} {mode}")
    
    del ff
    result = {"re_code": re_code, "mode": mode, "similarity": sim}
    with open(os.path.join(output_root, "result.json"), "w") as f:
        json.dump(result, f, indent=2)
    print(f"  Re{re_code} {mode}: similarity={sim:.4f}")
    return result


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ablation-json", type=str, required=True)
    ap.add_argument("--mode", type=str, default="v21")
    ap.add_argument("--validate-re", type=str, default="70")
    ap.add_argument("--device", type=int, default=2)
    ap.add_argument("--steps", type=int, default=100)
    ap.add_argument("--out-dir", type=str, default=os.path.join(OUTPUT_DIR, "sindy_val"))
    args = ap.parse_args()
    
    validate_re = [int(r) for r in args.validate_re.split(",")]
    coefs = load_ablation_coef(args.ablation_json, args.mode)
    os.makedirs(args.out_dir, exist_ok=True)
    
    for rc in validate_re:
        out_sub = os.path.join(args.out_dir, f"re{rc}")
        run_closed_loop(rc, coefs, args.mode, args.device, out_sub, n_steps=args.steps)


if __name__ == "__main__":
    main()