# analysis_crossre/scripts/validate_v23.py
"""Validate v23: front no-bias + rear shared-head.

Front: v2 coeffs with bias=0.
Top: v2 coeffs unchanged.
Bottom: -top(Gx), using G-transformed raw observations.

Usage:
    conda run -n pycuda_3_10 python validate_v23.py --re 70 --device 2
"""
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, action_to_physical, scale_action, build_karman_cloak_env,
    add_pinball, build_observation, compute_physical_symbols,
    save_vorticity_png, vorticity_from_ddf, compute_similarity,
    load_legacy_configs,
)
from cfg import (
    OUTPUT_DIR, SAMPLE_INTERVAL, FIFO_LEN, CONV_LEN, S_DIM,
    ACTION_SCALE, ACTION_BIAS, U0, CONFIG_DIR,
)

DATA_TYPE = np.float32

# v2 feature keys matching sindy_results_v2.json
V2_FEAT_KEYS = [
    "u_m", "u_a", "u_c", "v_a",
    "Fx_tot", "Fx_rear", "Fy_tot", "Fy_diff",
    "sin_ua", "cos_ua",
    "a0_lag1", "a1_lag1", "a2_lag1",
    "da0", "da1", "da2",
]
V2_MU_KEYS = ["mu", "mu_u_a", "mu_v_a", "mu_Fx_tot", "mu_Fy_diff", "mu_Fy_tot"]


def apply_G_raw(obs_slice, a_prev, a_prev2):
    """Apply mirror operator G to raw observations and actions.

    obs_slice: (12,) [s0_ux,s0_uy, s1_ux,s1_uy, s2_ux,s2_uy, front_fx,front_fy, bot_fx,bot_fy, top_fx,top_fy]
    a_prev: (3,) [aF, aB, aT]
    a_prev2: (3,) [aF_prev2, aB_prev2, aT_prev2]

    Returns (G_obs, G_a_prev, G_a_prev2)
    """
    # Sensors: top<->bottom swap, cross components negate
    G_obs = np.zeros(12, dtype=np.float64)
    G_obs[0] = obs_slice[4]    # s0_ux <- s2_ux (streamwise: no sign)
    G_obs[1] = -obs_slice[5]   # s0_uy <- -s2_uy (cross: negate)
    G_obs[2] = obs_slice[2]    # s1_ux unchanged
    G_obs[3] = -obs_slice[3]   # s1_uy negate
    G_obs[4] = obs_slice[0]    # s2_ux <- s0_ux
    G_obs[5] = -obs_slice[1]   # s2_uy <- -s0_uy

    # Forces: front unchanged (but lift sign flips), bottom<->top with sign flips
    G_obs[6] = obs_slice[6]      # front_fx unchanged
    G_obs[7] = -obs_slice[7]     # front_fy negate
    G_obs[8] = obs_slice[10]     # bot_fx <- top_fx
    G_obs[9] = -obs_slice[11]    # bot_fy <- -top_fy
    G_obs[10] = obs_slice[8]     # top_fx <- bot_fx
    G_obs[11] = -obs_slice[9]    # top_fy <- -bot_fy

    # Actions: all negate, B<->T swap
    G_a_prev = np.array([-a_prev[0], -a_prev[2], -a_prev[1]], dtype=np.float64)
    G_a_prev2 = np.array([-a_prev2[0], -a_prev2[2], -a_prev2[1]], dtype=np.float64)

    return G_obs, G_a_prev, G_a_prev2


def build_v2_feature_vec(obs_slice, actions_prev, actions_prev2, mu, add_bias):
    """Build feature vector matching v2 layout."""
    sensors = obs_slice[0:6].astype(np.float64).reshape(1, 6)
    forces = obs_slice[6:12].astype(np.float64).reshape(1, 6)
    ap = actions_prev.astype(np.float64).reshape(1, 3)
    ap2 = actions_prev2.astype(np.float64).reshape(1, 3)

    sym = compute_physical_symbols(sensors, forces, ap, ap2)
    sym["mu"] = np.array([mu])
    sym["mu_u_a"] = sym["u_a"] * mu
    sym["mu_v_a"] = sym["v_a"] * mu
    sym["mu_Fx_tot"] = sym["Fx_tot"] * mu
    sym["mu_Fy_diff"] = sym["Fy_diff"] * mu
    sym["mu_Fy_tot"] = sym["Fy_tot"] * mu

    vals = []
    if add_bias:
        vals.append(1.0)
    for k in V2_FEAT_KEYS:
        vals.append(float(sym[k][0]))
    for k in V2_MU_KEYS:
        vals.append(float(sym[k][0]))
    return np.array(vals, dtype=np.float64)


def load_v2_coefs(v2_path):
    """Load v2 coefficients. Zero front bias. Return front + top only."""
    with open(v2_path) as f:
        data = json.load(f)
    cross = data["cross_re"]
    coefs_list = cross["channels"]

    # Zero front bias
    coefs_list[0]["best_coef"][0] = 0.0

    return {
        "front": {
            "coef": np.array(coefs_list[0]["best_coef"], dtype=np.float64),
            "has_bias": True,
        },
        "top": {
            "coef": np.array(coefs_list[2]["best_coef"], dtype=np.float64),
            "has_bias": True,
        },
    }


def predict_v23(obs_slice, a_prev, a_prev2, coefs, mu):
    """Predict physical omega using v23 shared-head.

    Front: v2 coefs (bias zeroed).
    Top: v2 coefs unchanged.
    Bottom: -top(Gx).
    """
    # Front prediction
    feat = build_v2_feature_vec(obs_slice, a_prev, a_prev2, mu, add_bias=coefs["front"]["has_bias"])
    front = float(feat @ coefs["front"]["coef"])

    # Top prediction (from original state)
    feat_top = build_v2_feature_vec(obs_slice, a_prev, a_prev2, mu, add_bias=coefs["top"]["has_bias"])
    top = float(feat_top @ coefs["top"]["coef"])

    # Bottom = -top(Gx)
    G_obs, G_a_prev, G_a_prev2 = apply_G_raw(obs_slice, a_prev, a_prev2)
    feat_G = build_v2_feature_vec(G_obs, G_a_prev, G_a_prev2, mu, add_bias=coefs["top"]["has_bias"])
    top_at_Gx = float(feat_G @ coefs["top"]["coef"])
    bottom = -top_at_Gx

    return np.array([front, bottom, top], dtype=np.float64)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--re", type=int, 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"))
    ap.add_argument("--v2-results", type=str,
                    default=os.path.join(OUTPUT_DIR, "sindy", "sindy_results_v2.json"))
    args = ap.parse_args()

    re_code = args.re
    mu = 2.0 / re_code
    output_root = os.path.join(args.out_dir, f"re{re_code}")
    os.makedirs(output_root, exist_ok=True)

    print(f"\n=== v23 validation: Re{re_code} (mu={mu:.6f}) ===")

    coefs = load_v2_coefs(args.v2_results)
    for name in ["front", "top"]:
        print(f"  {name}: {len(coefs[name]['coef'])} coefs")

    # Build environment
    cuda_cfg, field_cfg = load_legacy_configs(CONFIG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(nu_from_re(re_code, u0=U0)))
    ff = FlowField(field_cfg, cuda_cfg, device_id=args.device)

    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)

    # 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 = []
    a_prev = action_to_physical(np.zeros((1,3), dtype=np.float32),
                                 scale=ACTION_SCALE, bias=ACTION_BIAS, u0=U0).flatten()
    a_prev2 = a_prev.copy()

    for step in range(args.steps):
        obs_slice = fifo[-1] if len(fifo) > 0 else np.zeros(12, dtype=np.float32)
        omega = predict_v23(obs_slice, a_prev, a_prev2, coefs, mu)

        # Apply action
        norm_action = (omega / 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])
        a_prev2 = a_prev.copy()
        a_prev = omega.copy()

    sens_arr = np.array(sens_sc, dtype=np.float32)
    sim = compute_similarity(target_states, sens_arr, CONV_LEN)
    print(f"  v23 similarity: {sim:.4f}")

    # Vorticity
    omega_vort = vorticity_from_ddf(ff, u0=U0)
    save_vorticity_png(os.path.join(output_root, "vorticity_v23.png"),
                       omega_vort, title=f"Re{re_code} v23 (shared-head)")

    result = {"re_code": re_code, "mode": "v23", "similarity": sim}
    with open(os.path.join(output_root, "result_v23.json"), "w") as f:
        json.dump(result, f, indent=2)

    del ff
    print(f"  Done -> {output_root}")
    return 0


if __name__ == "__main__":
    sys.exit(main())
