#!/usr/bin/env python3
"""A/B test: does Karman cylinder order in add_pinball() matter?

Runs the same PPO checkpoint on two FlowFields with different cylinder addition
orders, then compares sensor signals and actions.
"""
import os, sys, json, time, numpy as np
from collections import deque

_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
sys.path.insert(0, _REPO)
_SRC = os.path.join(_REPO, "src"); sys.path.insert(0, _SRC)

from LegacyCelerisLab import FlowField
from SR_analysis.utils.cfd_interface import load_legacy_configs, scale_action, build_karman_cloak_env, add_pinball as add_pinball_original, build_observation, load_ppo_model, compute_similarity
from SR_analysis.configs import get_scene, LEGACY_CFG_DIR, FIFO_LEN, CONV_LEN

DATA_TYPE = np.float32
SI = 800

def add_pinball_corrected(ff, **kw):
    """Same as add_pinball but with bottom-first order matching training env."""
    import copy
    from SR_analysis.utils.cfd_interface import add_pinball as _orig
    # We can't easily modify add_pinball inline, so we'll do the steps manually
    l0 = kw.get("l0", 20.0)
    u0 = kw.get("u0", 0.01)
    si_local = kw.get("sample_interval", SI)
    fl = kw.get("fifo_len", FIFO_LEN)
    dt = kw.get("data_type", DATA_TYPE)
    ab = kw.get("action_bias", (0.0, -4.0, 4.0))
    pfx = kw.get("pinball_front_x", 30.0)
    prx = kw.get("pinball_rear_x", 31.3)
    oss = kw.get("obs_slice_start", 2)
    ose = kw.get("obs_slice_end", 14)

    cy = (ff.FIELD_SHAPE[1] - 1) / 2.0
    u0_f = float(u0)

    # CORRECTED: front → BOTTOM(y=-0.75) → TOP(y=+0.75)
    ff.add_cylinder((pfx * l0, cy, 0.0), l0 / 2.0)           # front
    ff.add_cylinder((prx * l0, cy - 0.75 * l0, 0.0), l0 / 2.0)  # BOTTOM first
    ff.add_cylinder((prx * l0, cy + 0.75 * l0, 0.0), l0 / 2.0)  # TOP second

    n_obj = ff.obs.size // 2
    print(f"  [corrected] bodies after pinball: {n_obj}")

    # stabilize
    ff.run(int(4 * ff.FIELD_SHAPE[0] / u0_f), np.zeros(n_obj, dtype=dt))
    ff.get_ddf(); ff.save_ddf()

    # norm (zero-action)
    fifo = deque(maxlen=fl)
    for _ in range(fl):
        ff.run(si_local, np.zeros(n_obj, dtype=dt))
        fifo.append(ff.obs.copy()[oss:ose])
    tsa = np.array(fifo, dtype=np.float32)
    fnf = 6.0 * float(np.max(np.abs(tsa[:, 6:12])))
    sd = np.mean(tsa[:, 0:6], axis=0).astype(np.float32)
    snf = np.zeros(6, dtype=np.float32)
    for i in range(6):
        snf[i] = 5.0 * float(np.max(np.abs(tsa[:, i] - sd[i])))

    # bias FIFO
    ff.apply_ddf()
    bias = np.zeros(n_obj, dtype=dt)
    # CORRECTED order: id=4=front(0), id=5=BOTTOM(-4), id=6=TOP(+4)
    bias[n_obj - 3] = float(ab[0] * u0_f)   # front
    bias[n_obj - 2] = float(ab[1] * u0_f)   # BOTTOM gets bias=-4
    bias[n_obj - 1] = float(ab[2] * u0_f)   # TOP gets bias=+4
    print(f"  [corrected] bias: front={bias[n_obj-3]:.6f}, bottom={bias[n_obj-2]:.6f}, top={bias[n_obj-1]:.6f}")
    fifo.clear()
    for _ in range(fl):
        ff.run(si_local, bias); fifo.append(ff.obs.copy()[oss:ose])
    save_states = np.array(list(fifo), dtype=dt)
    ff.get_ddf(); ff.save_ddf(); ff.apply_ddf()

    return {"force_norm_fact": fnf, "sens_deviation": sd.tolist(),
            "sens_norm_fact": snf.tolist(), "action_bias": list(ab),
            "save_states": save_states}


def run_inference(ff, norm, n_obj_total, action_scale, action_bias, u0,
                  target_states, model_path, s_dim, n_steps, label):
    """Run controlled PPO inference and return sensors, actions, rewards."""
    model = load_ppo_model(model_path, device="cuda:2", s_dim=s_dim)
    
    # bias FIFO init
    fifo = deque(maxlen=FIFO_LEN)
    ba = scale_action(np.zeros(3, dtype=np.float32),
                      scale=action_scale, bias=action_bias, u0=u0, n_total_bodies=n_obj_total)
    for _ in range(FIFO_LEN):
        ff.context.push(); ff.run(SI, ba); ff.context.pop()
        fifo.append(ff.obs.copy()[2:14])

    s_c, f_c, a_c = [], [], []
    obs = np.zeros(12, dtype=np.float32)
    ba_arr = np.array(list(ba), dtype=np.float32)

    for step in range(n_steps):
        act, _ = model.predict(obs, deterministic=True)
        act = act.astype(np.float32).flatten(); a_c.append(act.copy())
        ff.context.push()
        ff.run(SI, scale_action(act, scale=action_scale, bias=action_bias, u0=u0, n_total_bodies=n_obj_total))
        ff.context.pop()
        osl = ff.obs.copy()[2:14]; fifo.append(osl)
        s_c.append(osl[0:6]); f_c.append(osl[6:12])
        forces_n = osl[6:12] / norm["force_norm_fact"]
        sens_n = (osl[0:6] - norm["sens_deviation"]) / norm["sens_norm_fact"]
        obs = np.clip(np.hstack([forces_n, sens_n]), -1, 1).astype(np.float32)

    sim = compute_similarity(target_states, np.array(s_c, dtype=np.float32), CONV_LEN)
    return np.array(s_c), np.array(a_c), sim, fifo


def main():
    cfg = get_scene("karman_re100")
    u0, nu, l0 = cfg["u0"], cfg["nu"], 20.0
    si, action_scale = cfg["sample_interval"], cfg["action_scale"]
    action_bias = tuple(cfg["action_bias"])
    n_obj_total = cfg["n_objects_env"]
    s_dim = cfg["s_dim"]
    model_path = "models/old/d1a3o12_re100.zip"
    n_steps = 100  # shorter for quick test

    cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(nu))

    print("=" * 60)
    print("A/B TEST: Karman Cylinder Order")
    print("=" * 60)

    # ── Run A: ORIGINAL order (front → TOP → BOTTOM) ──
    print("\n--- A: ORIGINAL add_pinball (front→TOP→BOTTOM) ---")
    ff = FlowField(field_cfg, cuda_cfg, device_id=2)
    target_states, _ = build_karman_cloak_env(ff, u0=u0, l0=l0, sample_interval=si, fifo_len=FIFO_LEN, data_type=DATA_TYPE)
    norm_a = add_pinball_original(ff, l0=l0, u0=u0, sample_interval=si, fifo_len=FIFO_LEN, data_type=DATA_TYPE,
                                   action_bias=action_bias, pinball_front_x=cfg["pinball_front_x"], pinball_rear_x=cfg["pinball_rear_x"],
                                   obs_slice_start=2, obs_slice_end=14)
    ff.restore_ddf(); ff.apply_ddf()
    # Show which cylinder gets which ID
    print(f"  Cylinder IDs: {list(ff.objects.keys())}")
    for k, v in ff.objects.items():
        print(f"    id={k}: center=({v['center'][0]:.0f}, {v['center'][1]:.1f}, {v['center'][2]:.0f})")

    sensors_a, actions_a, sim_a, _ = run_inference(
        ff, norm_a, n_obj_total, action_scale, action_bias, u0,
        target_states, model_path, s_dim, n_steps, "A")
    del ff
    print(f"  A similarity: {sim_a:.4f}")

    # ── Run B: CORRECTED order (front → BOTTOM → TOP) ──
    print("\n--- B: CORRECTED order (front→BOTTOM→TOP) ---")
    ff = FlowField(field_cfg, cuda_cfg, device_id=2)
    target_states2, _ = build_karman_cloak_env(ff, u0=u0, l0=l0, sample_interval=si, fifo_len=FIFO_LEN, data_type=DATA_TYPE)
    norm_b = add_pinball_corrected(ff, l0=l0, u0=u0, sample_interval=si, fifo_len=FIFO_LEN, data_type=DATA_TYPE,
                                    action_bias=action_bias, pinball_front_x=cfg["pinball_front_x"], pinball_rear_x=cfg["pinball_rear_x"],
                                    obs_slice_start=2, obs_slice_end=14)
    ff.restore_ddf(); ff.apply_ddf()
    print(f"  Cylinder IDs: {list(ff.objects.keys())}")
    for k, v in ff.objects.items():
        print(f"    id={k}: center=({v['center'][0]:.0f}, {v['center'][1]:.1f}, {v['center'][2]:.0f})")

    sensors_b, actions_b, sim_b, _ = run_inference(
        ff, norm_b, n_obj_total, action_scale, action_bias, u0,
        target_states2, model_path, s_dim, n_steps, "B")
    del ff
    print(f"  B similarity: {sim_b:.4f}")

    # ── Compare ──
    print("\n" + "=" * 60)
    print("COMPARISON")
    print("=" * 60)
    print(f"  Similarity: A={sim_a:.6f}, B={sim_b:.6f}, Δ={abs(sim_a - sim_b):.6f}")
    
    sensor_diff = np.max(np.abs(sensors_a - sensors_b))
    print(f"  Max sensor diff (A-B): {sensor_diff:.8f}")
    
    # Channel-wise action comparison
    for i, name in enumerate(["Front", "Bottom_channel", "Top_channel"]):
        corr = np.corrcoef(actions_a[:, i], actions_b[:, i])[0, 1]
        print(f"  Action {name}: corr={corr:.6f}")
    
    # Physical action (omega) comparison
    omega_a = (actions_a * action_scale + np.array(action_bias)) * u0
    omega_b = (actions_b * action_scale + np.array(action_bias)) * u0
    
    print(f"\n  Physical omega stats:")
    for i, name in enumerate(["Front", "Bottom(phys ch1)", "Top(phys ch2)"]):
        print(f"    A {name}: μ={omega_a[:,i].mean():.6f}, σ={omega_a[:,i].std():.6f}")
        print(f"    B {name}: μ={omega_b[:,i].mean():.6f}, σ={omega_b[:,i].std():.6f}")
    
    # Check: is the "bottom channel" action in B actually controlling the bottom cylinder?
    # If B correctly maps bias[-2]=bottom=cy-0.75 with bias=-4:
    #   bottom (cy-0.75): bias=-4, action from channel 1
    #   top (cy+0.75): bias=+4, action from channel 2
    # If A incorrectly maps:
    #   top (cy+0.75): bias=-4, action from channel 1
    #   bottom (cy-0.75): bias=+4, action from channel 2
    
    print(f"\n  VERDICT:")
    if abs(sim_a - sim_b) < 0.01 and sensor_diff < 1e-4:
        print(f"  → NO DIFFERENCE. Cylinder order does not affect PPO output.")
        print(f"  → Either PPO's rear channels are identical, or LegacyCelerisLab remaps actions.")
    elif sim_b > sim_a + 0.01:
        print(f"  → B is BETTER (corrected order improves similarity).")
        print(f"  → add_pinball order was indeed swapped and matters.")
    else:
        print(f"  → Difference exists but not conclusive.")

    return 0

if __name__ == "__main__":
    raise SystemExit(main())
