"""
DANTE v6 (full-adjustable, no SINDy pretraining)

Core decisions for this version:
1) Disable SINDy prior structure usage entirely.
2) Use full simplified basis pool and optimize all controller coefficients.
3) Keep minimal loop: no resume/checkpoint restore, live DB save + reward-only TB.
4) Continuous rollout objective with protective reset only on done/truncated.
"""

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

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_OBS = 2
N_ACT = 3
BIAS_SCALE = 1.0
COEF_SCALE = 2.0

# Full simplified basis pool: all terms are trainable in v6.
FULL_BASIS_TERMS = [
    "obs0", "obs1", "dobs0", "dobs1",
    "sin_obs0", "sin_obs1", "cos_obs0", "cos_obs1",
    "tanh_obs0", "tanh_obs1",
    "act0_l1", "act1_l1", "act2_l1",
]

BASIS_PROFILES = {
    # Compact profile (<20 dims) with derivative + nonlinear terms.
    "compact_deriv_nl": ["obs0", "obs1", "dobs0", "dobs1", "tanh_obs1"],
    # Constrained search top performer (<20 dims) with nonlinear + history.
    "compact_nl_hist": ["obs1", "sin_obs0", "cos_obs0", "act1_l1"],
}


# Fixed run configuration (kept in-file for reproducibility and stricter control)
V6_CONFIG = {
    "name": "d1a3o12_250421_forces02_dante_v6",
    "device_id": 1,
    "surrogate_gpu_id": 1,
    "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,
    "basis_profile": "compact_deriv_nl",
    # Use d-dependent initialization, then clamp into [min_num_initial, max_num_initial].
    "num_initial_per_dim": 10,
    "min_num_initial": 100,
    "max_num_initial": 240,
    "samples_per_acq": 18,
    "max_init_attempts_factor": 2.0,
    "surrogate_mode": "ensemble",  # mlp | cnn | ensemble
    # Match PPO script budget: 100 learn iterations * 2048 rollout steps
    "target_cfd_steps": 204800*2,
    # Keep consistent with surrogate study runs in this repo.
    "surrogate_epochs": 400,
}


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

    def close(self) -> None:
        pass


class LinearBasisController:
    def __init__(self, basis_terms: List[str]):
        self.basis_terms = list(basis_terms)
        self.num_basis = len(self.basis_terms)
        self.total_params = N_ACT * (1 + self.num_basis)

        self.params = np.zeros(self.total_params, dtype=np.float64)
        self.obs_l1 = np.zeros(2, dtype=np.float64)
        self.prev_action = np.zeros(3, dtype=np.float64)

    def reset_state(self, obs0: np.ndarray) -> None:
        obs0 = np.asarray(obs0, dtype=np.float64).reshape(-1)
        if obs0.size < 2:
            if obs0.size == 1:
                obs0 = np.array([obs0[0], obs0[0]], dtype=np.float64)
            else:
                obs0 = np.zeros(2, dtype=np.float64)
        self.obs_l1 = obs0[:2].copy()
        self.prev_action = np.zeros(3, dtype=np.float64)

    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}")
        self.params = np.clip(x, -1.0, 1.0)

    def _feature_dict(self, obs: np.ndarray) -> Dict[str, float]:
        o = np.asarray(obs, dtype=np.float64).reshape(-1)
        if o.size < 2:
            if o.size == 1:
                o = np.array([o[0], o[0]], dtype=np.float64)
            else:
                o = np.zeros(2, dtype=np.float64)

        o0, o1 = float(o[0]), float(o[1])
        o0_l1, o1_l1 = float(self.obs_l1[0]), float(self.obs_l1[1])
        a0, a1, a2 = float(self.prev_action[0]), float(self.prev_action[1]), float(self.prev_action[2])

        return {
            "obs0": o0,
            "obs1": o1,
            "dobs0": o0 - o0_l1,
            "dobs1": o1 - o1_l1,
            "sin_obs0": float(np.sin(np.pi * o0)),
            "sin_obs1": float(np.sin(np.pi * o1)),
            "cos_obs0": float(np.cos(np.pi * o0)),
            "cos_obs1": float(np.cos(np.pi * o1)),
            "tanh_obs0": float(np.tanh(o0)),
            "tanh_obs1": float(np.tanh(o1)),
            "act0_l1": a0,
            "act1_l1": a1,
            "act2_l1": a2,
        }

    def predict(self, obs: np.ndarray) -> np.ndarray:
        feat = self._feature_dict(obs)
        out = np.zeros(3, dtype=np.float64)
        stride = 1 + self.num_basis

        for ch in range(3):
            off = ch * stride
            qb = np.tanh(1.25 * self.params[off])
            y = BIAS_SCALE * qb
            for k, term in enumerate(self.basis_terms):
                qk = np.tanh(1.25 * self.params[off + 1 + k])
                y += (COEF_SCALE * qk) * feat.get(term, 0.0)
            out[ch] = y

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

        obs2 = np.asarray(obs, dtype=np.float64).reshape(-1)
        if obs2.size < 2:
            if obs2.size == 1:
                obs2 = np.array([obs2[0], obs2[0]], dtype=np.float64)
            else:
                obs2 = np.zeros(2, dtype=np.float64)

        self.obs_l1 = obs2[:2].copy()
        self.prev_action = out.copy()
        return out.astype(np.float32)


class FlowControlObjectiveV6(ObjectiveFunction):
    def __init__(
        self,
        basis_terms: List[str],
        eval_steps: int = 200,
        tail_steps: int = 100,
        startup_steps: int = 200,
        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.basis_terms = list(basis_terms)
        self.controller = LinearBasisController(self.basis_terms)
        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 = 1
        self.obs_current = None
        self.recover_count = 0

    def init_env(self, device_id: int = 1) -> 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(3, 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(self.obs_current)

    @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:
        x = self._preprocess(x)
        self.controller.set_params(x)

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

        # DANTE candidate evaluations are isolated from each other.
        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(obs)

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

            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)),
            }

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

    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 v6 full-adjustable (no SINDy pretraining)")
    p.add_argument(
        "--name",
        type=str,
        default=V6_CONFIG["name"],
        help="Optional run name override. Core budgets remain fixed in-file.",
    )
    p.add_argument(
        "--basis-profile",
        type=str,
        default=V6_CONFIG["basis_profile"],
        choices=sorted(BASIS_PROFILES.keys()),
        help="Controller basis profile to optimize.",
    )
    return p.parse_args()


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


def save_live_db(path: str, x: np.ndarray, y: np.ndarray, meta: Dict) -> 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,
) -> None:
    print(
        f"[{phase}] {idx}/{total} | reward={reward:.4f} | best={best_reward:.4f} | "
        f"recover={recover_total} | eval={global_step}/{total_expected} | elapsed={elapsed_sec:.1f}s",
        flush=True,
    )


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

    name = str(args.name)
    cfg = dict(V6_CONFIG)
    cfg["name"] = name

    basis_profile = str(args.basis_profile)
    basis_terms = list(BASIS_PROFILES[basis_profile])

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

    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)

    obj = FlowControlObjectiveV6(
        basis_terms=basis_terms,
        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"]))

    controller_dims = int(obj.dims)
    surrogate_gpu_id = int(cfg["surrogate_gpu_id"])
    num_initial_raw = int(max(1, round(float(cfg["num_initial_per_dim"]) * controller_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)))
    if max_init_attempts < num_initial:
        raise RuntimeError("max_init_attempts must be >= num_initial")
    total_expected = int(num_initial + num_acquisitions * int(cfg["samples_per_acq"]))

    config_payload = {
        "name": name,
        "use_sindy_prior": False,
        "reason": "forced_disabled_by_v6_full_adjustable",
        "basis_profile": basis_profile,
        "basis_terms": basis_terms,
        "controller_dims": controller_dims,
        "num_initial": 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"]),
        "surrogate_gpu_id": int(surrogate_gpu_id),
        "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("use_sindy_prior:", False)
    print("basis_profile:", basis_profile)
    print("basis_terms:", len(basis_terms))
    print("controller_dims:", controller_dims)
    print("num_initial:", num_initial)
    print("num_initial_raw:", int(num_initial_raw))
    print("num_initial_per_dim:", float(cfg["num_initial_per_dim"]))
    print("min_num_initial:", int(cfg["min_num_initial"]))
    print("max_num_initial:", int(cfg["max_num_initial"]))
    print("max_init_attempts:", max_init_attempts)
    print("max_init_attempts_factor:", float(cfg["max_init_attempts_factor"]))
    print("num_acquisitions:", int(num_acquisitions))
    print("samples_per_acq:", int(cfg["samples_per_acq"]))
    print("surrogate_mode:", str(cfg["surrogate_mode"]))
    print("surrogate_gpu_id:", surrogate_gpu_id)
    print("eval_steps:", int(cfg["eval_steps"]))
    print("tail_steps:", int(cfg["tail_steps"]))
    print("startup_steps:", int(cfg["startup_steps"]))
    print("target_cfd_steps:", int(cfg["target_cfd_steps"]))
    print("matched_cfd_steps:", int(total_expected * int(cfg["eval_steps"])))
    print("expected_total_evals:", total_expected)

    ckpt_path = Path(os.path.join(model_dir, f"{name}_surrogate.keras"))
    surrogate = FlowControlSurrogateV6(
        input_dims=controller_dims,
        epochs=int(cfg["surrogate_epochs"]),
        patience=30,
        check_point_path=ckpt_path,
        tf_device_id=int(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")

    if SummaryWriter is None:
        writer = NullWriter()
    else:
        writer = 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\n")

    input_x = np.empty((0, controller_dims), dtype=np.float64)
    input_y = np.empty((0,), dtype=np.float64)
    history = []
    best_reward = -1.0
    best_params = np.zeros(controller_dims, dtype=np.float64)
    invalid_skips = 0

    t0 = time.time()
    global_step = 0

    try:
        # Phase 1: random init points (collect exactly num_initial valid samples)
        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: "
                    f"valid={valid_init}/{num_initial}, attempts={init_attempts}/{max_init_attempts}, "
                    f"invalid_skips={invalid_skips}"
                )

            init_attempts += 1
            x = sample_params(controller_dims, obj.turn)
            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
                print(
                    f"[init] invalid sample skipped: attempt={init_attempts}, "
                    f"failure_code={info.get('failure_code', -1)}, "
                    f"valid={valid_init}/{num_initial}, invalid_skips={invalid_skips}",
                    flush=True,
                )

            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,
                            "basis_terms": basis_terms,
                            "controller_dims": controller_dims,
                            "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)),
                        "basis_terms": basis_terms,
                        "use_sindy_prior": False,
                        "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)}\n"
                )

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

        # Phase 2: DANTE acquisitions
        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=controller_dims,
                        epochs=int(cfg["surrogate_epochs"]),
                        patience=30,
                        check_point_path=cand_ckpt,
                        tf_device_id=int(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"device={fit_metrics.get('device', 'CPU')}",
                    flush=True,
                )
            print(
                f"[acq {acq + 1}/{int(num_acquisitions)}] surrogate fit done, rolling out {int(cfg['samples_per_acq'])} candidates...",
                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"]),
            )
            candidates = explorer.rollout(input_x, input_y, iteration=acq)

            acq_rewards = []
            for j, x in enumerate(candidates):
                info = obj.evaluate_candidate(x)
                reward = float(info["reward"])
                is_valid = int(info.get("failure_code", 0)) == 0
                if is_valid:
                    acq_rewards.append(reward)

                    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 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,
                                    "basis_terms": basis_terms,
                                    "controller_dims": controller_dims,
                                    "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)),
                            "basis_terms": basis_terms,
                            "use_sindy_prior": False,
                            "recover_count": int(obj.recover_count),
                            "acq": int(acq + 1),
                            "invalid_skips": int(invalid_skips),
                        },
                    )
                else:
                    invalid_skips += 1
                    print(
                        f"[acq {acq + 1}] invalid sample skipped: cand={j + 1}, failure_code={info.get('failure_code', -1)}",
                        flush=True,
                    )

                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)}\n"
                    )

                global_step += 1
                print_progress_line(
                    phase=f"acq{acq + 1}",
                    idx=j + 1,
                    total=len(candidates),
                    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,
                )

            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),
                    "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}, recover_total={obj.recover_count}"
            )

        with open(final_meta_path, "wb") as f:
            pickle.dump(
                {
                    "name": name,
                    "best_reward": best_reward,
                    "best_params": best_params,
                    "basis_terms": basis_terms,
                    "controller_dims": controller_dims,
                    "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()
