#!/usr/bin/env python3
"""Phase 0 calibration: measure baselines and produce calibration.json for any case.

Runs on new CelerisLab solver with 2000x600 config. Produces:
  - target.npy              (FIFO_LEN, 6) sensor signals (CelerisLab native units)
  - calibration.json        FORCE_SCALE, SENS_SCALE, dtw_norm_scale, SIM_BP, SIM_VAL

Workflow:
  1. Temp sim (dist_cyl + sensors only) -> warmup -> record target -> close
  2. Training sim (all 7 objects) -> warmup
  3. Stage0: zero-action measurement (150+100 steps)
  4. Stage1: open-loop reference measurement (legacy-equiv bias)
  5. Compute calibration constants from both stages

Usage:
    conda run -n pycuda_3_10 python calibrate.py \
        --case re100 --device-id 0 \
        --config configs/config_lbm_karman_2000x600.json
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from collections import deque
from pathlib import Path

import numpy as np
import pycuda.driver as cuda; cuda.init()

_REPO = Path(__file__).resolve().parents[3]
if str(_REPO) not in sys.path:
    sys.path.insert(0, str(_REPO))

from CelerisLab import Simulation
from normalization import CALIBRATION_SCHEMA_NATIVE_V2, SENSOR_UNITS_NATIVE

# ---------------------------------------------------------------------------
# Physics / geometry constants
# ---------------------------------------------------------------------------
L0 = 20.0
U0 = 0.01
RADIUS = L0 / 2.0
NX = 2000
NY = 600
CENTER_Y = float(NY - 1) / 2.0

DIST_X = 600.0
PINBALL_FRONT_X = 1000.0
PINBALL_REAR_X = 1026.0
SENSOR_X = 1200.0

FIFO_LEN = 150
CONV_LEN = 30

N_MEASURE = 100

# Reference action: legacy-equiv bias [0,-4,4]*U0, expressed as no_bias omega
# no_bias: omega = -(action * 12 + [0,0,0]) * U0 / R
# Target omega = [0, 0.004, -0.004] (= action_norm * 12 * U0 / R with action_norm [0, -4/12, 4/12])
_REF_OMEGA = np.array([0.0, 0.004, -0.004], dtype=np.float32)

# Reward constants
K_CD = 50.0; K_CL = 100.0
K_CD_ILLUSION = 12.0; K_CL_ILLUSION = 25.0   # lower: 3-pinball forces vs 1-target cylinder, wider error range
W_CD = 0.30; W_CL = 0.30; W_SIM = 0.40
FLOOR_CD = 0.10; FLOOR_CL = 0.10; FLOOR_SIM = 0.10
FLOOR_PENALTY = 0.05

ACTION_SCALE = 12.0
ACTION_BIAS = np.array([0.0, 0.0, 0.0], dtype=np.float32)


# ---------------------------------------------------------------------------
# DTW utilities
# ---------------------------------------------------------------------------
def calc_lag(target: np.ndarray, state: np.ndarray) -> int:
    t_mean = np.mean(target); s_mean = np.mean(state)
    corr = np.correlate(target - t_mean, state - s_mean, mode="full")
    lags = np.arange(-len(target) + 1, len(target))
    return int(lags[np.argmax(corr)])


def calc_dtw_sim(target: np.ndarray, state: np.ndarray,
                 norm_scale: float = 1.0) -> float:
    n, m = len(target), len(state)
    dtw = np.full((n + 1, m + 1), np.inf); dtw[0, 0] = 0.0
    for i in range(1, n + 1):
        for j in range(1, m + 1):
            cost = abs(float(target[i - 1]) - float(state[j - 1]))
            last_min = min(dtw[i - 1, j], dtw[i, j - 1], dtw[i - 1, j - 1])
            dtw[i, j] = cost + last_min
    raw = 1.0 - dtw[n, m] / (float(n) * norm_scale)
    return float(max(0.0, raw))


def compute_similarity(target_states, fifo_states, conv_len, norm_scale):
    target = np.asarray(target_states, dtype=np.float64)
    state = np.asarray(fifo_states, dtype=np.float64)
    id_sens = 3
    target_seq = target[conv_len:2 * conv_len, id_sens]
    state_seq = state[-conv_len:, id_sens]
    lag = calc_lag(target_seq, state_seq)
    sim = 0.0
    for i in range(6):
        t_seq = np.roll(target[:, i], -lag)[conv_len:2 * conv_len]
        s_seq = state[-conv_len:, i]
        sim += calc_dtw_sim(t_seq, s_seq, norm_scale=norm_scale)
    return float(sim / 6.0)


# ---------------------------------------------------------------------------
# CtxGuard
# ---------------------------------------------------------------------------
class CtxGuard:
    def __init__(self, sim):
        self.sim = sim
    def __enter__(self):
        if self.sim is not None:
            self.sim.ctx._ctx.push()
    def __exit__(self, *exc):
        if self.sim is not None:
            self.sim.ctx._ctx.pop()
        return False


def gpu_block(sim, fn):
    with CtxGuard(sim):
        fn()


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def read_obs(sim, dist_id, sensor_ids, pinball_ids) -> np.ndarray:
    obs = []
    if dist_id is not None:
        obs = list(sim.read_force(dist_id, normalize=True))
    else:
        obs = [0.0, 0.0]  # no disturbance cylinder (illusion)
    for sid in sensor_ids:
        s = sim.read_sensor(sid, normalize=True)
        obs.extend([float(s[0]), float(s[1])])
    for pid in pinball_ids:
        obs.extend(sim.read_force(pid, normalize=True))
    return np.array(obs, dtype=np.float32)


def set_omega(sim, pinball_ids, omega):
    for pid, w in zip(pinball_ids, omega):
        sim.set_body(pid, omega=float(w))


def run_stage(sim, dist_id, sensor_ids, pinball_ids, target_states,
              omega_vec, dtw_norm_scale, si, n_measure=N_MEASURE):
    gpu_block(sim, lambda: set_omega(sim, pinball_ids, omega_vec))
    fifo = deque(maxlen=FIFO_LEN)
    for _ in range(FIFO_LEN):
        gpu_block(sim, lambda: sim.run(si, zero_obs=True))
        obs = read_obs(sim, dist_id, sensor_ids, pinball_ids)
        sl = obs[2:14].copy()
        fifo.append(sl)

    obs_slices, sims = [], []
    for _ in range(n_measure):
        gpu_block(sim, lambda: sim.run(si, zero_obs=True))
        obs = read_obs(sim, dist_id, sensor_ids, pinball_ids)
        sl = obs[2:14].copy()
        fifo.append(sl)
        obs_slices.append(sl.copy())
        sim_val = compute_similarity(target_states, np.array(list(fifo)),
                                     CONV_LEN, dtw_norm_scale)
        sims.append(sim_val)

    return {
        "obs_slices": np.array(obs_slices, dtype=np.float32),
        "sims": np.array(sims, dtype=np.float32),
    }


# ---------------------------------------------------------------------------
# Illusion calibration
# ---------------------------------------------------------------------------

# Illusion-specific geometry
_ILL_PINBALL_FRONT_X = 380.0
_ILL_PINBALL_REAR_X = 406.0
_ILL_SENSOR_X = 600.0
_ILL_TARGET_X = 400.0
_ILL_REF_OMEGA = np.array([0.0, 0.005, -0.005], dtype=np.float32)  # steady-cloak-like 5*U0 bias

def _analyze_harmonics_for_calib(states, n_harmonics=5):
    N, D = states.shape
    result = []
    for d in range(D):
        y = states[:, d]
        fft_coef = np.fft.rfft(y)
        freqs = np.fft.rfftfreq(N, d=1)
        amps = 2 * np.abs(fft_coef) / N
        phases = np.angle(fft_coef)
        idx = np.argsort(amps[1:])[::-1][:n_harmonics] + 1
        harmonics = {
            'dc': float(np.real(fft_coef[0]) / N),
            'amps': amps[idx].tolist(),
            'freqs': freqs[idx].tolist(),
            'phases': phases[idx].tolist(),
        }
        result.append(harmonics)
    return result


def _calibrate_illusion(case, config_path, device_id, si, out_dir, log, warmup, target_diam=1.0):
    target_radius = target_diam * L0
    # ---- Step 1: Record target (target cylinder + 3 sensors) ----
    log(f"Step 1: Recording illusion target (target cyl + sensors, diam={target_diam}L)...")
    sim = Simulation(lbm_config_path=config_path, device_id=device_id)
    sim._assert_object_count_contract = lambda *a, **kw: None
    sim.add_body("circle", center=(_ILL_TARGET_X, CENTER_Y, 0.0),
                 radius=target_radius)
    s0 = sim.add_body("sensor", center=(_ILL_SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0)
    s1 = sim.add_body("sensor", center=(_ILL_SENSOR_X, CENTER_Y, 0.0), radius=5.0)
    s2 = sim.add_body("sensor", center=(_ILL_SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0)
    sim.initialize()
    gpu_block(sim, lambda: sim.run(warmup, zero_obs=True))
    log(f"  Warmup done ({warmup} steps)")

    target_states = np.zeros((FIFO_LEN, 8), dtype=np.float32)
    for i in range(FIFO_LEN):
        gpu_block(sim, lambda: sim.run(si, zero_obs=True))
        target_states[i] = [
            sim.read_sensor(s0, normalize=True)[0],
            sim.read_sensor(s0, normalize=True)[1],
            sim.read_sensor(s1, normalize=True)[0],
            sim.read_sensor(s1, normalize=True)[1],
            sim.read_sensor(s2, normalize=True)[0],
            sim.read_sensor(s2, normalize=True)[1],
            sim.read_force(0, normalize=True)[0],
            sim.read_force(0, normalize=True)[1],
        ]
    sim.close()
    # Sensor-only portion for DTW
    target_sensor = target_states[:, 0:6].copy()
    np.save(str(out_dir / "target.npy"), target_sensor)

    # Harmonics
    harmonics = _analyze_harmonics_for_calib(target_states, n_harmonics=5)
    with open(out_dir / "target_harmonics.json", "w") as f:
        json.dump(harmonics, f, indent=2)

    target_std = np.std(target_sensor, axis=0)
    uy_std_avg = float(np.mean([target_std[1], target_std[3], target_std[5]]))
    dtw_norm_scale = max(uy_std_avg, np.finfo(np.float32).eps)
    log(f"  Target recorded. DTW norm_scale: {dtw_norm_scale:.4f}")

    # ---- Step 2: Training sim (3 sensors + 3 pinball) ----
    log("Step 2: Creating training sim (6 objects)...")
    sim = Simulation(lbm_config_path=config_path, device_id=device_id)
    sim._assert_object_count_contract = lambda *a, **kw: None
    sensor_ids = [
        sim.add_body("sensor", center=(_ILL_SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0),
        sim.add_body("sensor", center=(_ILL_SENSOR_X, CENTER_Y, 0.0), radius=5.0),
        sim.add_body("sensor", center=(_ILL_SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0),
    ]
    sim.add_body("circle", center=(_ILL_PINBALL_FRONT_X, CENTER_Y, 0.0), radius=RADIUS)
    sim.add_body("circle", center=(_ILL_PINBALL_REAR_X, CENTER_Y + 15.0, 0.0), radius=RADIUS)
    sim.add_body("circle", center=(_ILL_PINBALL_REAR_X, CENTER_Y - 15.0, 0.0), radius=RADIUS)
    sim.initialize()
    pinball_ids = [3, 4, 5]

    gpu_block(sim, lambda: sim.run(warmup, zero_obs=True))
    log(f"  Warmup done ({warmup} steps)")
    gpu_block(sim, lambda: sim.snapshot())

    # ---- Step 3: Stage0 (zero rotation) ----
    log("Step 3: Stage0 (zero rotation)...")
    gpu_block(sim, lambda: sim.restore())
    stage0 = run_stage(sim, None, sensor_ids, pinball_ids, target_sensor,
                       np.zeros(3, dtype=np.float32), dtw_norm_scale, si)

    # ---- Step 4: Stage1 (reference open-loop) ----
    log("Step 4: Stage1 (reference open-loop)...")
    gpu_block(sim, lambda: sim.restore())
    stage1 = run_stage(sim, None, sensor_ids, pinball_ids, target_sensor,
                       _ILL_REF_OMEGA, dtw_norm_scale, si)

    sim.close()
    log("  CFD closed.")

    # ---- Step 5: Compute calibration ----
    log("Step 5: Computing calibration constants...")

    all_forces = np.concatenate([stage0["obs_slices"][:, 6:12],
                                  stage1["obs_slices"][:, 6:12]], axis=0)
    force_scale = float(np.max(np.abs(all_forces)))
    force_scale = max(force_scale, 0.0001)

    all_sens = np.concatenate([stage0["obs_slices"][:, 0:6],
                                stage1["obs_slices"][:, 0:6]], axis=0)
    sens_scale = float(np.max(np.abs(all_sens)))
    sens_scale = max(sens_scale, np.finfo(np.float32).eps)

    s0_sim = float(np.mean(stage0["sims"]))
    s1_sim = float(np.mean(stage1["sims"]))

    # Build SIM_BP: [0, worst_measured, better_measured, ..., 1.0]
    # Normally (Karman) Stage0 < Stage1 (zero is worse, reference is better).
    # But for Illusion, Stage0 may already match the target well, and reference
    # rotation (steady-cloak-like) makes it worse.  Swap roles in that case.
    if s1_sim < s0_sim:
        # Reference action is worse than zero (typical Illusion case).
        # Stage1 (5*U0 bias) -> r_sim ~0.10, Stage0 (zero) -> r_sim ~0.35
        worst_sim = s1_sim
        better_sim = s0_sim
    else:
        # Karman case or zero is worse than bias.
        worst_sim = s0_sim
        better_sim = s1_sim

    sim_bp = [0.0, worst_sim, better_sim,
              better_sim + (1.0 - better_sim) * 0.4,
              better_sim + (1.0 - better_sim) * 0.7,
              1.0]
    # Illusion reward mapping: zero~0.35, bias~0.10, best~1.0
    sim_val = [0.0, 0.1, 0.35, 0.7, 0.85, 1.0]

    for i in range(1, len(sim_bp)):
        if sim_bp[i] <= sim_bp[i - 1]:
            sim_bp[i] = sim_bp[i - 1] + 0.01

    # Rounding
    force_scale = round(force_scale, 4)
    sens_scale = round(sens_scale, 6)
    dtw_norm_scale = round(dtw_norm_scale, 6)
    sim_bp = [round(x, 2) for x in sim_bp]
    for i in range(1, len(sim_bp)):
        if sim_bp[i] <= sim_bp[i - 1]:
            sim_bp[i] = sim_bp[i - 1] + 0.01
    sim_bp[-1] = 1.0

    log(f"  FORCE_SCALE = {force_scale:.4f}, SENS_SCALE = {sens_scale:.6f}")
    log(f"  Stage0 sim = {s0_sim:.4f}, Stage1 sim = {s1_sim:.4f}")
    log(f"  SIM_BP = {[f'{x:.2f}' for x in sim_bp]}")

    # ---- Step 6: Write calibration.json ----
    calibration = {
        "schema_version": CALIBRATION_SCHEMA_NATIVE_V2,
        "sensor_units": SENSOR_UNITS_NATIVE,
        "case": case, "scene": "illusion",
        "target_diam": target_diam,
        "grid": {"nx": NX, "ny": NY},
        "config_path": config_path,
        "SI": si,
        "FIFO_LEN": FIFO_LEN, "CONV_LEN": CONV_LEN,
        "FORCE_SCALE": force_scale, "SENS_SCALE": sens_scale,
        "dtw_norm_scale": float(dtw_norm_scale),
        "SIM_BP": [float(x) for x in sim_bp],
        "SIM_VAL": [float(x) for x in sim_val],
        "K_CD": K_CD_ILLUSION, "K_CL": K_CL_ILLUSION,
        "W_CD": W_CD, "W_CL": W_CL, "W_SIM": W_SIM,
        "FLOOR_CD": FLOOR_CD, "FLOOR_CL": FLOOR_CL, "FLOOR_SIM": FLOOR_SIM,
        "FLOOR_PENALTY": FLOOR_PENALTY,
        "ACTION_SCALE": ACTION_SCALE, "ACTION_BIAS": ACTION_BIAS.tolist(),
        "U0": U0, "RADIUS": RADIUS, "L0": L0,
    }
    with open(out_dir / "calibration.json", "w") as f:
        json.dump(calibration, f, indent=2)

    log(f"\nCalibration complete. Files in {out_dir}:")
    log(f"  calibration.json, target.npy, target_harmonics.json, calibrate.log")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main() -> int:
    parser = argparse.ArgumentParser(description="Phase 0 calibration")
    parser.add_argument("--case", type=str, required=True)
    parser.add_argument("--device-id", type=int, default=0)
    parser.add_argument("--config", type=str, required=True)
    parser.add_argument("--si", type=int, default=800)
    parser.add_argument("--out-dir", type=str, default=None)
    parser.add_argument("--scene", type=str, default="karman",
                        choices=["karman", "illusion"],
                        help="Scene type (karman or illusion)")
    parser.add_argument("--target-diam", type=float, default=1.0,
                        help="Target cylinder diameter in L0 units (illusion only, default=1.0)")
    parser.add_argument("--dist-radius", type=float, default=1.0,
                        help="Disturbance cylinder radius in L0 units (karman only, default=1.0)")
    args = parser.parse_args()

    case = args.case
    scene = args.scene
    config_path = args.config
    device_id = args.device_id
    si = args.si
    warmup = int(4.0 * NX / U0)

    if args.out_dir:
        out_dir = Path(args.out_dir).resolve()
    else:
        out_dir = Path(__file__).resolve().parent / "calibrations" / case
    out_dir.mkdir(parents=True, exist_ok=True)

    log_path = out_dir / "calibrate.log"
    def log(msg):
        line = f"[{time.strftime('%H:%M:%S')}] {msg}"
        print(line, flush=True)
        with open(log_path, "a") as f:
            f.write(line + "\n"); f.flush()

    log(f"=== Phase 0 Calibration: {case} ({scene}) ===")
    log(f"  Config: {config_path}")
    log(f"  SI: {si}")
    log(f"  Output: {out_dir}")

    if scene == "illusion":
        _calibrate_illusion(case, config_path, device_id, si, out_dir, log, warmup, target_diam=args.target_diam)
        return 0

    # ---- Step 1: Record target (Karman: dist_cyl + sensors only) ----
    log("Step 1: Recording target signal...")
    dist_radius = args.dist_radius * L0
    log(f"  Dist cylinder radius = {args.dist_radius}L = {dist_radius:.1f} lattice")
    sim = Simulation(lbm_config_path=config_path, device_id=device_id)
    sim._assert_object_count_contract = lambda *a, **kw: None
    dist_id_t = sim.add_body("circle", center=(DIST_X, CENTER_Y, 0.0), radius=dist_radius)
    s0_t = sim.add_body("sensor", center=(SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0)
    s1_t = sim.add_body("sensor", center=(SENSOR_X, CENTER_Y, 0.0), radius=5.0)
    s2_t = sim.add_body("sensor", center=(SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0)
    sim.initialize()

    gpu_block(sim, lambda: sim.run(warmup, zero_obs=True))
    log(f"  Warmup done ({warmup} steps)")

    target = np.zeros((FIFO_LEN, 6), dtype=np.float32)
    for i in range(FIFO_LEN):
        gpu_block(sim, lambda: sim.run(si, zero_obs=True))
        target[i] = [
            sim.read_sensor(s0_t, normalize=True)[0],
            sim.read_sensor(s0_t, normalize=True)[1],
            sim.read_sensor(s1_t, normalize=True)[0],
            sim.read_sensor(s1_t, normalize=True)[1],
            sim.read_sensor(s2_t, normalize=True)[0],
            sim.read_sensor(s2_t, normalize=True)[1],
        ]
    sim.close()
    np.save(str(out_dir / "target.npy"), target)
    target_std = np.std(target, axis=0)
    uy_std_avg = float(np.mean([target_std[1], target_std[3], target_std[5]]))
    dtw_norm_scale = max(uy_std_avg, np.finfo(np.float32).eps)
    log(f"  Target recorded. s1_uy std={target[:, 3].std():.4f}")
    log(f"  DTW norm_scale: {dtw_norm_scale:.4f}")

    # ---- Step 2: Create training sim ----
    log("Step 2: Creating training sim (all 7 objects)...")
    sim = Simulation(lbm_config_path=config_path, device_id=device_id)
    sim._assert_object_count_contract = lambda *a, **kw: None
    dist_id = sim.add_body("circle", center=(DIST_X, CENTER_Y, 0.0), radius=dist_radius)
    sensor_ids = [
        sim.add_body("sensor", center=(SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0),
        sim.add_body("sensor", center=(SENSOR_X, CENTER_Y, 0.0), radius=5.0),
        sim.add_body("sensor", center=(SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0),
    ]
    sim.add_body("circle", center=(PINBALL_FRONT_X, CENTER_Y, 0.0), radius=RADIUS)
    sim.add_body("circle", center=(PINBALL_REAR_X, CENTER_Y + 15.0, 0.0), radius=RADIUS)
    sim.add_body("circle", center=(PINBALL_REAR_X, CENTER_Y - 15.0, 0.0), radius=RADIUS)
    sim.initialize()
    pinball_ids = [4, 5, 6]

    gpu_block(sim, lambda: sim.run(warmup, zero_obs=True))
    log(f"  Warmup done ({warmup} steps)")
    gpu_block(sim, lambda: sim.snapshot())

    # ---- Step 3: Stage0 (zero rotation) ----
    log("Step 3: Stage0 (zero rotation)...")
    gpu_block(sim, lambda: sim.restore())
    stage0 = run_stage(sim, dist_id, sensor_ids, pinball_ids, target,
                       np.zeros(3, dtype=np.float32), dtw_norm_scale, si)

    # ---- Step 4: Stage1 (reference open-loop) ----
    log("Step 4: Stage1 (reference open-loop)...")
    gpu_block(sim, lambda: sim.restore())
    stage1 = run_stage(sim, dist_id, sensor_ids, pinball_ids, target,
                       _REF_OMEGA, dtw_norm_scale, si)

    sim.close()
    log("  CFD closed.")

    # ---- Step 5: Compute calibration ----
    log("Step 5: Computing calibration constants...")

    all_forces = np.concatenate([stage0["obs_slices"][:, 6:12],
                                  stage1["obs_slices"][:, 6:12]], axis=0)
    force_scale = float(np.max(np.abs(all_forces)))
    force_scale = max(force_scale, 0.0001)

    all_sens = np.concatenate([stage0["obs_slices"][:, 0:6],
                                stage1["obs_slices"][:, 0:6]], axis=0)
    sens_scale = float(np.max(np.abs(all_sens)))
    sens_scale = max(sens_scale, np.finfo(np.float32).eps)

    log(f"  FORCE_SCALE = {force_scale:.6f}")
    log(f"  SENS_SCALE  = {sens_scale:.4f}")

    s0_sim = float(np.mean(stage0["sims"]))
    s1_sim = float(np.mean(stage1["sims"]))
    log(f"  Stage0 sim_mean = {s0_sim:.4f}")
    log(f"  Stage1 sim_mean = {s1_sim:.4f}")

    # Build SIM_BP: always [0, worst_measured, better_measured, ..., 1.0]
    if s1_sim < s0_sim:
        worst_sim = s1_sim
        better_sim = s0_sim
    else:
        worst_sim = s0_sim
        better_sim = s1_sim

    # When stage0 and stage1 are too close (|s1-s0| < 0.05), the measured
    # SIM_BP collapses to near-zero dynamic range.  Fall back to a generic
    # mapping anchored on the well-behaved re100 baseline (s0~0.32, s1~0.82).
    # This gives the DRL a meaningful reward gradient regardless of the
    # open-loop baselines.
    spread = abs(s1_sim - s0_sim)
    if spread < 0.10 or better_sim < 0.5:
        log(f"  Spread={spread:.3f} too small or better<0.5 — using generic SIM_BP mapping")
        sim_bp = [0.0, 0.30, 0.65, 0.79, 0.89, 1.0]
        sim_val = [0.0, 0.20, 0.50, 0.80, 0.90, 1.0]
    else:
        sim_bp = [0.0, worst_sim, better_sim,
                  better_sim + (1.0 - better_sim) * 0.4,
                  better_sim + (1.0 - better_sim) * 0.7,
                  1.0]
        sim_val = [0.0, 0.2, 0.5, 0.8, 0.9, 1.0]

        for i in range(1, len(sim_bp)):
            if sim_bp[i] <= sim_bp[i - 1]:
                sim_bp[i] = sim_bp[i - 1] + 0.01

    log(f"  SIM_BP (raw) = {[f'{x:.4f}' for x in sim_bp]}")
    log(f"  SIM_VAL = {sim_val}")

    # ---- Step 5b: Round to clean values ----
    # Round measured values to reasonable precision so calibration files
    # are human-readable and robust to small measurement noise.
    force_scale = round(force_scale, 4)          # 0.002429 -> 0.0024
    sens_scale = round(sens_scale, 6)            # 0.7543 -> 0.75
    dtw_norm_scale = round(dtw_norm_scale, 6)    # 0.2043 -> 0.204
    sim_bp = [round(x, 2) for x in sim_bp]       # 0.3166 -> 0.32 etc.
    # Ensure sim_bp monotonicity preserved after rounding
    for i in range(1, len(sim_bp)):
        if sim_bp[i] <= sim_bp[i - 1]:
            sim_bp[i] = sim_bp[i - 1] + 0.01
    # Clamp last to 1.0
    sim_bp[-1] = 1.0

    log(f"  FORCE_SCALE (rounded) = {force_scale:.4f}")
    log(f"  SENS_SCALE  (rounded) = {sens_scale:.6f}")
    log(f"  dtw_norm_scale (rounded) = {dtw_norm_scale:.6f}")
    log(f"  SIM_BP (rounded) = {[f'{x:.2f}' for x in sim_bp]}")

    # ---- Step 6: Write calibration.json ----
    calibration = {
        "schema_version": CALIBRATION_SCHEMA_NATIVE_V2,
        "sensor_units": SENSOR_UNITS_NATIVE,
        "case": case,
        "dist_radius": args.dist_radius,
        "grid": {"nx": NX, "ny": NY},
        "config_path": config_path,
        "SI": si,
        "FIFO_LEN": FIFO_LEN,
        "CONV_LEN": CONV_LEN,
        "FORCE_SCALE": force_scale,
        "SENS_SCALE": sens_scale,
        "dtw_norm_scale": float(dtw_norm_scale),
        "SIM_BP": [float(x) for x in sim_bp],
        "SIM_VAL": [float(x) for x in sim_val],
        "K_CD": K_CD,
        "K_CL": K_CL,
        "W_CD": W_CD,
        "W_CL": W_CL,
        "W_SIM": W_SIM,
        "FLOOR_CD": FLOOR_CD,
        "FLOOR_CL": FLOOR_CL,
        "FLOOR_SIM": FLOOR_SIM,
        "FLOOR_PENALTY": FLOOR_PENALTY,
        "ACTION_SCALE": ACTION_SCALE,
        "ACTION_BIAS": ACTION_BIAS.tolist(),
        "U0": U0,
        "RADIUS": RADIUS,
        "L0": L0,
    }

    with open(out_dir / "calibration.json", "w") as f:
        json.dump(calibration, f, indent=2)

    log(f"\nCalibration complete. Files in {out_dir}:")
    log(f"  calibration.json, target.npy, calibrate.log")
    return 0


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