"""
DANTE v7 (open-loop periodic control)

This version switches from closed-loop basis feedback to a periodic open-loop controller:
1) Optimize control points of one cycle directly.
2) Use continuous phase advance to support non-integer cycle lengths.
3) Keep v6-compatible evaluation protocol (300-step rollout, tail100 score).
"""

import argparse
import json
import os
import pickle
import shutil
import sys
import time
from pathlib import Path
from typing import Dict, List, Optional

import numpy as np

os.environ["MKL_THREADING_LAYER"] = "GNU"
os.environ["OMP_NUM_THREADS"] = "16"
os.environ["MKL_NUM_THREADS"] = "16"

try:
    from torch.utils.tensorboard import SummaryWriter
except Exception:
    SummaryWriter = None


CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.abspath(os.path.join(CURRENT_DIR, os.pardir))
os.chdir(CURRENT_DIR)
sys.path.insert(0, os.path.join(ROOT, "DANTE"))

from dante.obj_functions import ObjectiveFunction
from dante.tree_exploration import TreeExploration

from dante_v6_surrogate_torch import FlowControlSurrogateV6
from dante_pinball.env.gym_env_dante_total_force import CustomEnv


N_ACT = 3

V7_CONFIG = {
    "name": "d1a3o12_250421_forces02_dante_v7_openloop",
    "device_id": 0,
    "surrogate_gpu_id": 0,
    "eval_steps": 300,
    "tail_steps": 100,
    "startup_steps": 0,
    "max_recover_resets": 1,
    "reset_each_candidate": True,
    "obs_fail_bound": 2.0,
    "obs_clip_bound": 3.0,
    "control_points_per_channel": 8,
    "phase_interp": "linear",
    "period_mode": "auto",  # auto | fixed | optimize
    "fixed_period": 40.0,
    "period_min": 15.0,
    "period_max": 80.0,
    "period_estimate_files": [
        "output/report_dante_v2_v5_v6/raw_oldenv_seed_11.npz",
        "output/report_dante_v2_v5_v6/raw_oldenv_seed_29.npz",
        "output/report_dante_v2_v5_v6/raw_oldenv_seed_47.npz",
    ],
    "period_estimate_key": "ppo_actions",
    "num_initial_per_dim": 10,
    "min_num_initial": 100,
    "max_num_initial": 320,
    "samples_per_acq": 24,
    "max_init_attempts_factor": 2.0,
    "surrogate_mode": "ensemble",  # mlp | cnn | ensemble
    "target_cfd_steps": 204800*2,
    "surrogate_epochs": 400,
}


class NullWriter:
    def add_scalar(self, *_args, **_kwargs) -> None:
        pass

    def close(self) -> None:
        pass


def map_unit_to_range(u: float, low: float, high: float) -> float:
    u_clip = float(np.clip(u, -1.0, 1.0))
    alpha = 0.5 * (u_clip + 1.0)
    return float(low + alpha * (high - low))


def map_range_to_unit(x: float, low: float, high: float) -> float:
    if high <= low:
        return 0.0
    alpha = (float(x) - low) / (high - low)
    return float(np.clip(2.0 * alpha - 1.0, -1.0, 1.0))


def dominant_period_fft(x: np.ndarray, min_period: float, max_period: float) -> Optional[float]:
    x = np.asarray(x, dtype=np.float64).reshape(-1)
    n = int(x.size)
    if n < 16:
        return None

    xc = x - np.mean(x)
    std = float(np.std(xc))
    if std < 1e-8:
        return None

    yf = np.fft.rfft(xc)
    power = (np.abs(yf) ** 2).reshape(-1)
    freq = np.fft.rfftfreq(n, d=1.0)

    with np.errstate(divide="ignore", invalid="ignore"):
        period = np.where(freq > 0.0, 1.0 / freq, np.inf)

    valid = (freq > 0.0) & (period >= float(min_period)) & (period <= float(max_period))
    if not np.any(valid):
        return None

    idx = np.argmax(power * valid.astype(np.float64))
    f_peak = float(freq[idx])
    if f_peak <= 0.0:
        return None
    return float(1.0 / f_peak)


def load_action_array_from_npz(npz_path: str, key_hint: str = "ppo_actions") -> np.ndarray:
    with np.load(npz_path) as z:
        keys = list(z.keys())
        if key_hint in z:
            arr = z[key_hint]
            return np.asarray(arr, dtype=np.float64)

        for k in keys:
            kl = k.lower()
            if "ppo" in kl and "action" in kl:
                arr = z[k]
                return np.asarray(arr, dtype=np.float64)

        for k in keys:
            arr = np.asarray(z[k])
            if arr.ndim == 2 and arr.shape[1] == N_ACT:
                return np.asarray(arr, dtype=np.float64)

    raise KeyError(f"no action array found in {npz_path}")


def estimate_period_from_npz_list(
    rel_paths: List[str],
    key_hint: str,
    min_period: float,
    max_period: float,
) -> Dict[str, object]:
    periods: List[float] = []
    details: List[Dict[str, object]] = []

    for rel in rel_paths:
        abs_path = os.path.join(ROOT, rel)
        if not os.path.exists(abs_path):
            details.append({"file": rel, "used": False, "reason": "missing"})
            continue

        try:
            a = load_action_array_from_npz(abs_path, key_hint=key_hint)
            if a.ndim != 2 or a.shape[1] != N_ACT:
                details.append({
                    "file": rel,
                    "used": False,
                    "reason": f"invalid_shape_{tuple(a.shape)}",
                })
                continue

            file_periods = []
            for ch in range(N_ACT):
                p = dominant_period_fft(
                    a[:, ch],
                    min_period=float(min_period),
                    max_period=float(max_period),
                )
                if p is not None and np.isfinite(p):
                    file_periods.append(float(p))
                    periods.append(float(p))

            details.append({
                "file": rel,
                "used": len(file_periods) > 0,
                "num_channels": int(len(file_periods)),
                "channel_periods": [float(v) for v in file_periods],
            })
        except Exception as e:
            details.append({"file": rel, "used": False, "reason": str(e)})

    if len(periods) == 0:
        return {
            "period": None,
            "source": "none",
            "details": details,
            "count": 0,
        }

    period_med = float(np.median(np.asarray(periods, dtype=np.float64)))
    period_med = float(np.clip(period_med, min_period, max_period))

    return {
        "period": period_med,
        "source": "fft_median",
        "details": details,
        "count": int(len(periods)),
        "all_periods": [float(v) for v in periods],
    }


class PeriodicOpenLoopController:
    def __init__(
        self,
        control_points_per_channel: int,
        period_mode: str,
        base_period_steps: float,
        period_min: float,
        period_max: float,
        interp: str = "linear",
    ):
        self.k = int(control_points_per_channel)
        if self.k < 3:
            raise ValueError("control_points_per_channel must be >= 3")

        self.period_mode = str(period_mode).lower()
        if self.period_mode not in {"auto", "fixed", "optimize"}:
            raise ValueError("period_mode must be one of auto|fixed|optimize")

        self.base_period_steps = float(base_period_steps)
        self.period_min = float(period_min)
        self.period_max = float(period_max)
        self.interp = str(interp).lower()
        if self.interp not in {"linear"}:
            raise ValueError("only linear interpolation is currently supported")

        self.optimize_period = self.period_mode == "optimize"

        self.ctrl_points = np.zeros((N_ACT, self.k), dtype=np.float64)
        self.phase = 0.0
        self.current_period_steps = float(np.clip(self.base_period_steps, self.period_min, self.period_max))

        self.total_params = int(N_ACT * self.k + (1 if self.optimize_period else 0))

    def reset_state(self) -> None:
        self.phase = 0.0

    def _eval_channel(self, points: np.ndarray, phase: float) -> float:
        p = np.asarray(points, dtype=np.float64).reshape(-1)
        z = float(np.mod(phase, 1.0)) * self.k
        i0 = int(np.floor(z)) % self.k
        frac = float(z - np.floor(z))
        i1 = (i0 + 1) % self.k
        return float((1.0 - frac) * p[i0] + frac * p[i1])

    def set_params(self, x: np.ndarray) -> None:
        x = np.asarray(x, dtype=np.float64).reshape(-1)
        if x.size != self.total_params:
            raise ValueError(f"controller params mismatch: {x.size} != {self.total_params}")

        core = np.clip(x[: N_ACT * self.k], -1.0, 1.0)
        self.ctrl_points = core.reshape(N_ACT, self.k)

        if self.optimize_period:
            p_unit = float(x[-1])
            self.current_period_steps = map_unit_to_range(p_unit, self.period_min, self.period_max)
        else:
            self.current_period_steps = float(np.clip(self.base_period_steps, self.period_min, self.period_max))

    def predict(self, _obs: np.ndarray) -> np.ndarray:
        action = np.zeros(N_ACT, dtype=np.float64)
        for ch in range(N_ACT):
            action[ch] = self._eval_channel(self.ctrl_points[ch], self.phase)

        action = np.clip(action, -1.0, 1.0)

        step_phase = 1.0 / max(1e-6, float(self.current_period_steps))
        self.phase = float((self.phase + step_phase) % 1.0)

        return action.astype(np.float32)


class FlowControlObjectiveV7(ObjectiveFunction):
    def __init__(
        self,
        control_points_per_channel: int,
        period_mode: str,
        period_steps: float,
        period_min: float,
        period_max: float,
        phase_interp: str = "linear",
        eval_steps: int = 300,
        tail_steps: int = 100,
        startup_steps: int = 0,
        max_recover_resets: int = 1,
        turn: float = 0.05,
        reset_each_candidate: bool = True,
        obs_fail_bound: float = 2.0,
        obs_clip_bound: float = 3.0,
    ):
        self.eval_steps = int(eval_steps)
        self.tail_steps = int(tail_steps)
        self.startup_steps = int(startup_steps)
        self.max_recover_resets = int(max_recover_resets)
        self.turn = float(turn)
        self.reset_each_candidate = bool(reset_each_candidate)
        self.obs_fail_bound = float(obs_fail_bound)
        self.obs_clip_bound = float(obs_clip_bound)

        self.control_points_per_channel = int(control_points_per_channel)
        self.period_mode = str(period_mode)
        self.period_steps = float(period_steps)
        self.period_min = float(period_min)
        self.period_max = float(period_max)
        self.phase_interp = str(phase_interp)

        self.controller = PeriodicOpenLoopController(
            control_points_per_channel=self.control_points_per_channel,
            period_mode=self.period_mode,
            base_period_steps=self.period_steps,
            period_min=self.period_min,
            period_max=self.period_max,
            interp=self.phase_interp,
        )

        self.dims = int(self.controller.total_params)
        self.lb = -1.0 * np.ones(self.dims)
        self.ub = 1.0 * np.ones(self.dims)

        self.env = None
        self._device_id = 0
        self.obs_current = None
        self.recover_count = 0

    def init_env(self, device_id: int = 0) -> None:
        self._device_id = int(device_id)
        if self.env is not None:
            self.env.close()
        self.env = CustomEnv(
            device_id=int(device_id),
            obs_fail_bound=float(self.obs_fail_bound),
            obs_clip_bound=float(self.obs_clip_bound),
        )
        self._hard_reset_and_stabilize()

    def _hard_reset_and_stabilize(self) -> None:
        obs, _ = self.env.reset()
        obs = np.asarray(obs, dtype=np.float32)
        action = np.zeros(N_ACT, dtype=np.float32)

        for _ in range(int(self.startup_steps)):
            obs, _, done, trunc, _ = self.env.step(action)
            if done or trunc:
                obs, _ = self.env.reset()

        self.obs_current = np.asarray(obs, dtype=np.float32)
        self.controller.reset_state()

    @staticmethod
    def _tail_mean(rewards: List[float], tail_steps: int) -> float:
        rr = np.asarray(rewards, dtype=np.float64)
        if rr.size == 0:
            return 0.0
        tail = rr[-tail_steps:] if rr.size >= tail_steps else rr
        return float(np.mean(tail))

    def evaluate_candidate(self, x: np.ndarray) -> Dict[str, object]:
        x = self._preprocess(x)
        self.controller.set_params(x)

        if self.env is None:
            self.init_env(self._device_id)

        if self.reset_each_candidate:
            self._hard_reset_and_stabilize()

        for attempt in range(int(self.max_recover_resets) + 1):
            obs = np.asarray(self.obs_current, dtype=np.float32)
            self.controller.reset_state()

            rewards: List[float] = []
            steps = 0
            done = False
            trunc = False
            last_info: Dict[str, object] = {}

            for _ in range(int(self.eval_steps)):
                action = self.controller.predict(obs)
                obs, reward, done, trunc, step_info = self.env.step(action)
                if isinstance(step_info, dict):
                    last_info = step_info
                rewards.append(float(reward))
                steps += 1
                if done or trunc:
                    break

            if done or trunc and attempt < int(self.max_recover_resets):
                self.recover_count += 1
                self._hard_reset_and_stabilize()
                continue

            self.obs_current = np.asarray(obs, dtype=np.float32)
            y = self._tail_mean(rewards, tail_steps=int(self.tail_steps))
            return {
                "reward": float(y),
                "scaled": float(y * 100.0),
                "steps": int(steps),
                "done": bool(done),
                "truncated": bool(trunc),
                "recoveries_used": int(attempt),
                "failure_code": int(last_info.get("failure_code", 0)),
                "period_steps": float(self.controller.current_period_steps),
            }

        return {
            "reward": 0.0,
            "scaled": 0.0,
            "steps": 0,
            "done": True,
            "truncated": True,
            "recoveries_used": int(self.max_recover_resets),
            "failure_code": 1,
            "period_steps": float(self.controller.current_period_steps),
        }

    def scaled(self, y: float) -> float:
        return float(y * 100.0)

    def __call__(self, x: np.ndarray, apply_scaling: bool = False, track: bool = True) -> float:
        info = self.evaluate_candidate(x)
        y = float(info["reward"])
        return self.scaled(y) if apply_scaling else y


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="DANTE v7 open-loop periodic controller")
    p.add_argument("--name", type=str, default=V7_CONFIG["name"], help="Optional run name override")
    p.add_argument(
        "--control-points",
        type=int,
        default=V7_CONFIG["control_points_per_channel"],
        help="Control points per channel for one cycle",
    )
    p.add_argument(
        "--period-mode",
        type=str,
        default=V7_CONFIG["period_mode"],
        choices=["auto", "fixed", "optimize"],
        help="Cycle period mode",
    )
    p.add_argument("--fixed-period", type=float, default=V7_CONFIG["fixed_period"], help="Fixed period in steps")
    p.add_argument("--period-min", type=float, default=V7_CONFIG["period_min"], help="Min period for clipping")
    p.add_argument("--period-max", type=float, default=V7_CONFIG["period_max"], help="Max period for clipping")
    p.add_argument(
        "--phase-interp",
        type=str,
        default=V7_CONFIG["phase_interp"],
        choices=["linear"],
        help="Interpolation mode for phase to control point",
    )
    return p.parse_args()


def sample_params(dims: int, turn: float, period_mode: str, period_guess: float, pmin: float, pmax: float) -> np.ndarray:
    x = np.random.uniform(-1.0, 1.0, size=dims)
    x = np.round(x / turn) * turn
    x = np.clip(x, -1.0, 1.0)

    if str(period_mode).lower() == "optimize":
        x[-1] = map_range_to_unit(period_guess, pmin, pmax)
    return x


def save_live_db(path: str, x: np.ndarray, y: np.ndarray, meta: Dict[str, object]) -> None:
    np.savez(path, input_x=x, input_y=y, meta_json=np.array([json.dumps(meta, ensure_ascii=False)]))


def print_progress_line(
    phase: str,
    idx: int,
    total: int,
    reward: float,
    best_reward: float,
    recover_total: int,
    elapsed_sec: float,
    global_step: int,
    total_expected: int,
    period_steps: float,
) -> None:
    print(
        f"[{phase}] {idx}/{total} | reward={reward:.4f} | best={best_reward:.4f} | "
        f"period={period_steps:.3f} | recover={recover_total} | eval={global_step}/{total_expected} | elapsed={elapsed_sec:.1f}s",
        flush=True,
    )


def main() -> None:
    args = parse_args()

    cfg = dict(V7_CONFIG)
    cfg["name"] = str(args.name)
    cfg["control_points_per_channel"] = int(args.control_points)
    cfg["period_mode"] = str(args.period_mode)
    cfg["fixed_period"] = float(args.fixed_period)
    cfg["period_min"] = float(args.period_min)
    cfg["period_max"] = float(args.period_max)
    cfg["phase_interp"] = str(args.phase_interp)

    period_info = estimate_period_from_npz_list(
        rel_paths=list(cfg["period_estimate_files"]),
        key_hint=str(cfg["period_estimate_key"]),
        min_period=float(cfg["period_min"]),
        max_period=float(cfg["period_max"]),
    )

    if str(cfg["period_mode"]).lower() == "fixed":
        period_steps = float(np.clip(cfg["fixed_period"], cfg["period_min"], cfg["period_max"]))
        period_source = "fixed"
    else:
        p_est = period_info.get("period", None)
        if p_est is None:
            period_steps = float(np.clip(cfg["fixed_period"], cfg["period_min"], cfg["period_max"]))
            period_source = "fallback_fixed_no_estimate"
        else:
            period_steps = float(np.clip(float(p_est), cfg["period_min"], cfg["period_max"]))
            period_source = "auto_from_ppo_fft"

    target_evals = int(cfg["target_cfd_steps"] // cfg["eval_steps"])

    obj = FlowControlObjectiveV7(
        control_points_per_channel=int(cfg["control_points_per_channel"]),
        period_mode=str(cfg["period_mode"]),
        period_steps=float(period_steps),
        period_min=float(cfg["period_min"]),
        period_max=float(cfg["period_max"]),
        phase_interp=str(cfg["phase_interp"]),
        eval_steps=int(cfg["eval_steps"]),
        tail_steps=int(cfg["tail_steps"]),
        startup_steps=int(cfg["startup_steps"]),
        max_recover_resets=int(cfg["max_recover_resets"]),
        reset_each_candidate=bool(cfg["reset_each_candidate"]),
        obs_fail_bound=float(cfg["obs_fail_bound"]),
        obs_clip_bound=float(cfg["obs_clip_bound"]),
    )
    obj.init_env(device_id=int(cfg["device_id"]))

    dims = int(obj.dims)
    num_initial_raw = int(max(1, round(float(cfg["num_initial_per_dim"]) * dims)))
    num_initial = int(min(int(cfg["max_num_initial"]), max(int(cfg["min_num_initial"]), num_initial_raw)))
    num_acquisitions = int((target_evals - int(num_initial)) // int(cfg["samples_per_acq"]))
    if num_acquisitions <= 0:
        raise RuntimeError(
            f"computed num_acquisitions <= 0 with target_evals={target_evals}, "
            f"num_initial={num_initial}, samples_per_acq={int(cfg['samples_per_acq'])}"
        )

    max_init_attempts = int(max(num_initial, round(float(cfg["max_init_attempts_factor"]) * num_initial)))
    total_expected = int(num_initial + num_acquisitions * int(cfg["samples_per_acq"]))

    name = str(cfg["name"])
    model_dir = os.path.join(ROOT, "models", "250421")
    out_dir = os.path.join(ROOT, "output")
    tb_dir = os.path.join(ROOT, "tensorboard", name)
    os.makedirs(model_dir, exist_ok=True)
    os.makedirs(out_dir, exist_ok=True)

    config_payload = {
        "name": name,
        "control_mode": "open_loop_periodic",
        "device_id": int(cfg["device_id"]),
        "surrogate_gpu_id": int(cfg["surrogate_gpu_id"]),
        "control_points_per_channel": int(cfg["control_points_per_channel"]),
        "controller_dims": int(dims),
        "period_mode": str(cfg["period_mode"]),
        "period_steps": float(period_steps),
        "period_source": period_source,
        "period_min": float(cfg["period_min"]),
        "period_max": float(cfg["period_max"]),
        "period_estimate": period_info,
        "phase_interp": str(cfg["phase_interp"]),
        "num_initial": int(num_initial),
        "num_initial_raw": int(num_initial_raw),
        "num_initial_per_dim": float(cfg["num_initial_per_dim"]),
        "min_num_initial": int(cfg["min_num_initial"]),
        "max_num_initial": int(cfg["max_num_initial"]),
        "num_acquisitions": int(num_acquisitions),
        "samples_per_acq": int(cfg["samples_per_acq"]),
        "surrogate_mode": str(cfg["surrogate_mode"]),
        "eval_steps": int(cfg["eval_steps"]),
        "tail_steps": int(cfg["tail_steps"]),
        "startup_steps": int(cfg["startup_steps"]),
        "max_recover_resets": int(cfg["max_recover_resets"]),
        "obs_fail_bound": float(cfg["obs_fail_bound"]),
        "obs_clip_bound": float(cfg["obs_clip_bound"]),
        "reset_each_candidate": bool(cfg["reset_each_candidate"]),
        "max_init_attempts": int(max_init_attempts),
        "max_init_attempts_factor": float(cfg["max_init_attempts_factor"]),
        "target_cfd_steps": int(cfg["target_cfd_steps"]),
        "target_total_evals": int(target_evals),
        "matched_total_evals": int(total_expected),
        "matched_cfd_steps": int(total_expected * int(cfg["eval_steps"])),
    }

    with open(os.path.join(out_dir, f"{name}_structure_decision.json"), "w", encoding="utf-8") as f:
        json.dump(config_payload, f, indent=2, ensure_ascii=False)

    print("control_mode:", "open_loop_periodic")
    print("device_id:", int(cfg["device_id"]))
    print("surrogate_gpu_id:", int(cfg["surrogate_gpu_id"]))
    print("control_points_per_channel:", int(cfg["control_points_per_channel"]))
    print("controller_dims:", int(dims))
    print("period_mode:", str(cfg["period_mode"]))
    print("period_steps:", float(period_steps))
    print("period_source:", period_source)
    print("num_initial:", int(num_initial))
    print("num_initial_raw:", int(num_initial_raw))
    print("num_acquisitions:", int(num_acquisitions))
    print("samples_per_acq:", int(cfg["samples_per_acq"]))
    print("target_cfd_steps:", int(cfg["target_cfd_steps"]))
    print("matched_cfd_steps:", int(total_expected * int(cfg["eval_steps"])))

    ckpt_path = Path(os.path.join(model_dir, f"{name}_surrogate.keras"))
    surrogate = FlowControlSurrogateV6(
        input_dims=dims,
        epochs=int(cfg["surrogate_epochs"]),
        patience=30,
        check_point_path=ckpt_path,
        tf_device_id=int(cfg["surrogate_gpu_id"]),
    )

    live_db_path = os.path.join(out_dir, f"{name}_database_live.npz")
    best_params_path = os.path.join(model_dir, f"{name}_best.npy")
    best_meta_path = os.path.join(out_dir, f"{name}_best_meta.pkl")
    final_meta_path = os.path.join(out_dir, f"{name}_final_meta.pkl")
    dante_log_path = os.path.join(out_dir, f"{name}_dante_log.csv")

    writer = NullWriter() if SummaryWriter is None else SummaryWriter(log_dir=tb_dir)

    with open(dante_log_path, "w", encoding="utf-8") as f:
        f.write("timestamp,phase,acq,candidate,reward,best_reward,dataset_size,recoveries_used,failure_code,period_steps\n")

    input_x = np.empty((0, dims), dtype=np.float64)
    input_y = np.empty((0,), dtype=np.float64)
    history: List[Dict[str, object]] = []
    best_reward = -1.0
    best_params = np.zeros(dims, dtype=np.float64)
    invalid_skips = 0

    t0 = time.time()
    global_step = 0

    try:
        valid_init = 0
        init_attempts = 0
        while valid_init < num_initial:
            if init_attempts >= max_init_attempts:
                raise RuntimeError(
                    f"init failed to collect enough valid samples: valid={valid_init}/{num_initial}, "
                    f"attempts={init_attempts}/{max_init_attempts}, invalid_skips={invalid_skips}"
                )

            init_attempts += 1
            x = sample_params(
                dims=dims,
                turn=obj.turn,
                period_mode=str(cfg["period_mode"]),
                period_guess=float(period_steps),
                pmin=float(cfg["period_min"]),
                pmax=float(cfg["period_max"]),
            )

            info = obj.evaluate_candidate(x)
            reward = float(info["reward"])
            is_valid = int(info.get("failure_code", 0)) == 0

            if is_valid:
                valid_init += 1
                input_x = np.vstack((input_x, x.reshape(1, -1)))
                input_y = np.append(input_y, float(info["scaled"]))
                writer.add_scalar("Reward", reward, global_step)
            else:
                invalid_skips += 1

            if is_valid and reward > best_reward:
                best_reward = reward
                best_params = x.copy()
                np.save(best_params_path, best_params)
                with open(best_meta_path, "wb") as f:
                    pickle.dump(
                        {
                            "best_reward": best_reward,
                            "best_params": best_params,
                            "config": config_payload,
                            "timestamp": time.time(),
                        },
                        f,
                    )

            if is_valid:
                save_live_db(
                    live_db_path,
                    input_x,
                    input_y,
                    {
                        "name": name,
                        "best_reward": best_reward,
                        "dataset_size": int(len(input_y)),
                        "control_mode": "open_loop_periodic",
                        "recover_count": int(obj.recover_count),
                        "invalid_skips": int(invalid_skips),
                    },
                )

            with open(dante_log_path, "a", encoding="utf-8") as f:
                f.write(
                    f"{time.time():.3f},init,0,{init_attempts},{reward:.8f},{best_reward:.8f},{len(input_y)},{info['recoveries_used']},{info.get('failure_code', 0)},{float(info.get('period_steps', period_steps)):.6f}\n"
                )

            global_step += 1
            print(
                f"[init-attempt {init_attempts}/{max_init_attempts}] valid={valid_init}/{num_initial} | "
                f"reward={reward:.4f} | best={best_reward:.4f} | period={float(info.get('period_steps', period_steps)):.3f} | "
                f"invalid_skips={invalid_skips} | recover={int(obj.recover_count)} | "
                f"eval={global_step}/{total_expected}+ | elapsed={float(time.time() - t0):.1f}s",
                flush=True,
            )

        for acq in range(int(num_acquisitions)):
            if len(input_y) < 2:
                print(f"[acq {acq + 1}] skipped: insufficient valid samples ({len(input_y)})", flush=True)
                continue

            print(f"[acq {acq + 1}/{int(num_acquisitions)}] fitting surrogate on {len(input_y)} samples...", flush=True)

            surrogate_mode = str(cfg.get("surrogate_mode", "mlp")).lower()
            if surrogate_mode == "ensemble":
                candidates = []
                for arch in ["mlp", "cnn"]:
                    cand_ckpt = Path(str(ckpt_path).replace(".keras", f"_{arch}.keras"))
                    surr = FlowControlSurrogateV6(
                        input_dims=dims,
                        epochs=int(cfg["surrogate_epochs"]),
                        patience=30,
                        check_point_path=cand_ckpt,
                        tf_device_id=int(cfg["surrogate_gpu_id"]),
                        architecture=arch,
                    )
                    try:
                        m = surr(input_x, input_y, verbose=0)
                        fm = dict(getattr(surr, "last_fit_metrics", {}) or {})
                        candidates.append((arch, m, fm, cand_ckpt))
                    except Exception as e:
                        print(f"[acq {acq + 1}] surrogate {arch} failed: {e}", flush=True)

                if not candidates:
                    raise RuntimeError("all surrogate candidates failed in ensemble mode")

                candidates.sort(key=lambda z: float(z[2].get("val_r2", -1e9)), reverse=True)
                best_arch, model, fit_metrics, best_ckpt = candidates[0]
                if best_ckpt.exists():
                    shutil.copy2(best_ckpt, ckpt_path)
                fit_metrics["selected_architecture"] = best_arch
            else:
                surrogate.architecture = "cnn" if surrogate_mode == "cnn" else "mlp"
                model = surrogate(input_x, input_y, verbose=0)
                fit_metrics = dict(getattr(surrogate, "last_fit_metrics", {}) or {})
                fit_metrics["selected_architecture"] = str(surrogate.architecture)

            if fit_metrics:
                writer.add_scalar("Surrogate/val_r2", float(fit_metrics.get("val_r2", 0.0)), acq)
                writer.add_scalar("Surrogate/val_mae", float(fit_metrics.get("val_mae", 0.0)), acq)
                print(
                    f"[acq {acq + 1}] surrogate metrics: "
                    f"val_r2={fit_metrics.get('val_r2', float('nan')):.4f}, "
                    f"val_mae={fit_metrics.get('val_mae', float('nan')):.4f}, "
                    f"selected={fit_metrics.get('selected_architecture', 'na')}, "
                    f"device={fit_metrics.get('device', 'CPU')}",
                    flush=True,
                )

            if ckpt_path.exists():
                shutil.copy2(ckpt_path, os.path.join(model_dir, f"{name}_surrogate_live.keras"))

            explorer = TreeExploration(
                func=obj,
                model=model,
                num_samples_per_acquisition=int(cfg["samples_per_acq"]),
            )
            candidate_x = explorer.rollout(input_x, input_y, iteration=acq)

            acq_rewards: List[float] = []
            acq_periods: List[float] = []
            for j, x in enumerate(candidate_x):
                info = obj.evaluate_candidate(x)
                reward = float(info["reward"])
                period_used = float(info.get("period_steps", period_steps))
                is_valid = int(info.get("failure_code", 0)) == 0

                if is_valid:
                    acq_rewards.append(reward)
                    acq_periods.append(period_used)

                    input_x = np.vstack((input_x, np.asarray(x, dtype=np.float64).reshape(1, -1)))
                    input_y = np.append(input_y, float(info["scaled"]))
                    writer.add_scalar("Reward", reward, global_step)

                    if str(cfg["period_mode"]).lower() == "optimize":
                        writer.add_scalar("Controller/period_steps", period_used, global_step)

                    if reward > best_reward:
                        best_reward = reward
                        best_params = np.asarray(x, dtype=np.float64).copy()
                        np.save(best_params_path, best_params)
                        with open(best_meta_path, "wb") as f:
                            pickle.dump(
                                {
                                    "best_reward": best_reward,
                                    "best_params": best_params,
                                    "config": config_payload,
                                    "timestamp": time.time(),
                                },
                                f,
                            )

                    save_live_db(
                        live_db_path,
                        input_x,
                        input_y,
                        {
                            "name": name,
                            "best_reward": best_reward,
                            "dataset_size": int(len(input_y)),
                            "control_mode": "open_loop_periodic",
                            "recover_count": int(obj.recover_count),
                            "acq": int(acq + 1),
                            "invalid_skips": int(invalid_skips),
                        },
                    )
                else:
                    invalid_skips += 1

                with open(dante_log_path, "a", encoding="utf-8") as f:
                    f.write(
                        f"{time.time():.3f},acq,{acq + 1},{j + 1},{reward:.8f},{best_reward:.8f},{len(input_y)},{info['recoveries_used']},{info.get('failure_code', 0)},{period_used:.6f}\n"
                    )

                global_step += 1
                print_progress_line(
                    phase=f"acq{acq + 1}",
                    idx=j + 1,
                    total=len(candidate_x),
                    reward=reward,
                    best_reward=best_reward,
                    recover_total=int(obj.recover_count),
                    elapsed_sec=float(time.time() - t0),
                    global_step=global_step,
                    total_expected=total_expected,
                    period_steps=period_used,
                )

            history.append(
                {
                    "iteration": int(acq + 1),
                    "new_mean": float(np.mean(acq_rewards) if acq_rewards else 0.0),
                    "new_max": float(np.max(acq_rewards) if acq_rewards else 0.0),
                    "period_mean": float(np.mean(acq_periods) if acq_periods else period_steps),
                    "period_std": float(np.std(acq_periods) if acq_periods else 0.0),
                    "best_reward": float(best_reward),
                    "dataset_size": int(len(input_y)),
                    "recover_total": int(obj.recover_count),
                }
            )

            print(
                f"acq {acq + 1}/{num_acquisitions}: mean={history[-1]['new_mean']:.4f}, "
                f"max={history[-1]['new_max']:.4f}, best={best_reward:.4f}, "
                f"period_mean={history[-1]['period_mean']:.3f}, recover_total={obj.recover_count}",
                flush=True,
            )

        with open(final_meta_path, "wb") as f:
            pickle.dump(
                {
                    "name": name,
                    "best_reward": best_reward,
                    "best_params": best_params,
                    "history": history,
                    "elapsed_sec": float(time.time() - t0),
                    "config": config_payload,
                    "recover_total": int(obj.recover_count),
                },
                f,
            )

        print("Training done")
        print("best_reward:", f"{best_reward:.4f}")
        print("recover_total:", obj.recover_count)
        print("invalid_skips:", invalid_skips)

    finally:
        writer.close()
        if obj.env is not None:
            obj.env.close()


if __name__ == "__main__":
    main()
