"""Online fine-tuning: OnlineFineTuner, run_finetune_cycle, and the driver loop."""

from __future__ import annotations

from collections import deque
import gc
import os

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim

from ndvi_forecast.losses import horizon_weights, make_time_mask, masked_mse, masked_slope_mse
from ndvi_forecast._utils import to_numpy, safe_r2
from ndvi_forecast.offset import pre_fine_tuning_offsetter, post_fine_tuning_offsetter
from ndvi_forecast.evaluate import save_dap_results


# ---------------------------------------------------------------------------
# OnlineFineTuner
# ---------------------------------------------------------------------------

class OnlineFineTuner:
    """Lean online fine-tuner with trust-region, trajectory/flow priors, and EMA acceptance."""

    def __init__(
        self,
        model,
        lr: float = 1e-4,
        weight_decay: float = 0.0,
        grad_clip: float = 1.0,
        alpha_ar: float = 0.35,
        horizon_gamma: float = 0.6,
        lambda_slope: float = 0.25,
        lambda_trust: float = 1e-4,
        train_subset: str = "all",
        ema_decay: float = 0.98,
        lambda_collapse: float = 1e-2,
        collapse_margin: float = 0.025,
        learn_output_affine: bool = True,
        lambda_out_affine: float = 1e-4,
        out_gain_bounds: tuple = (0.85, 1.20),
        out_bias_bounds: tuple = (-0.10, 0.10),
        lambda_traj: float = 0.10,
        traj_k: int = 20,
        traj_recent_m: int = 7,
        traj_tau: float = 8.0,
        traj_slope_clip: float = 0.05,
        lambda_flow: float = 0.15,
        flow_k: int = 20,
        flow_recent_m: int = 8,
        flow_tau: float = 10.0,
        flow_curv_clip: float = 0.02,
    ):
        self.model = model
        self.device = next(model.parameters()).device
        self.grad_clip = grad_clip
        self.alpha_ar = float(alpha_ar)
        self.horizon_gamma = float(horizon_gamma)
        self.lambda_slope = float(lambda_slope)
        self.lambda_trust = float(lambda_trust)
        self.ema_decay = float(ema_decay)
        self.lambda_collapse = float(lambda_collapse)
        self.collapse_margin = float(collapse_margin)
        self.learn_output_affine = bool(learn_output_affine)
        self.lambda_out_affine = float(lambda_out_affine)
        self._s_lo, self._s_hi = float(out_gain_bounds[0]), float(out_gain_bounds[1])
        self._c_lo, self._c_hi = float(out_bias_bounds[0]), float(out_bias_bounds[1])
        self.lambda_traj = float(max(0.0, lambda_traj))
        self.traj_k = int(max(1, traj_k))
        self.traj_recent_m = int(max(2, traj_recent_m))
        self.traj_tau = float(max(1e-4, traj_tau))
        self.traj_slope_clip = float(max(0.0, traj_slope_clip))
        self.lambda_flow = float(max(0.0, lambda_flow))
        self.flow_k = int(max(2, flow_k))
        self.flow_recent_m = int(max(3, flow_recent_m))
        self.flow_tau = float(max(1e-4, flow_tau))
        self.flow_curv_clip = float(max(0.0, flow_curv_clip))

        # Freeze encoder
        if hasattr(model, "encoder"):
            for p in model.encoder.parameters():
                p.requires_grad = False

        # Collect trainable parameters
        train_params, self._names = [], []
        def _named(scope):
            return list(getattr(model, scope).named_parameters()) if hasattr(model, scope) else []

        dec_named = _named("decoder")
        head_named = _named("delta_head")
        is_bias = lambda n: n.endswith(".bias") or n == "bias"

        # Scopes that are always included when present
        _extra_scopes = ("affine_head_gain", "affine_head_bias", "season_head", "stitch_gate")

        if train_subset == "all":
            for n, p in dec_named:
                train_params.append(p); self._names.append(("decoder", n))
            for n, p in head_named:
                train_params.append(p); self._names.append(("delta_head", n))
            for scope in _extra_scopes:
                if hasattr(model, scope):
                    for n, p in getattr(model, scope).named_parameters():
                        train_params.append(p); self._names.append((scope, n))
        elif train_subset == "bias_head":
            for n, p in dec_named:
                if is_bias(n):
                    train_params.append(p); self._names.append(("decoder", n))
            for n, p in head_named:
                if is_bias(n):
                    train_params.append(p); self._names.append(("delta_head", n))
            for scope in _extra_scopes:
                if hasattr(model, scope):
                    for n, p in getattr(model, scope).named_parameters():
                        if is_bias(n):
                            train_params.append(p); self._names.append((scope, n))
        elif train_subset == "head_only":
            for n, p in head_named:
                train_params.append(p); self._names.append(("delta_head", n))
            for scope in _extra_scopes:
                if hasattr(model, scope):
                    for n, p in getattr(model, scope).named_parameters():
                        train_params.append(p); self._names.append((scope, n))
        else:
            raise ValueError(f"train_subset must be 'all', 'bias_head', or 'head_only', got '{train_subset}'")

        # Output-affine scalars
        if self.learn_output_affine:
            self.out_gain = nn.Parameter(torch.tensor(1.0, device=self.device))
            self.out_bias = nn.Parameter(torch.tensor(0.0, device=self.device))
            train_params += [self.out_gain, self.out_bias]
        else:
            self.out_gain = nn.Parameter(torch.tensor(1.0, device=self.device), requires_grad=False)
            self.out_bias = nn.Parameter(torch.tensor(0.0, device=self.device), requires_grad=False)

        self.trainable_params = train_params
        self.optimizer = optim.Adam(self.trainable_params, lr=lr, weight_decay=weight_decay)

        # Snapshot for trust-region + EMA shadow
        self._theta0, self._ema = {}, {}
        with torch.no_grad():
            for scope, name in self._names:
                p = getattr(model, scope).get_parameter(name).detach().clone()
                self._theta0[(scope, name)] = p.clone()
                self._ema[(scope, name)] = p.clone()
            if self.learn_output_affine:
                self._ema[("aux", "out_gain")] = self.out_gain.detach().clone()
                self._ema[("aux", "out_bias")] = self.out_bias.detach().clone()

    def _apply_out_affine(self, y):
        if not self.learn_output_affine:
            return y
        return self.out_gain * y + self.out_bias

    def _ema_update(self):
        d = self.ema_decay
        with torch.no_grad():
            for scope, name in self._names:
                p = getattr(self.model, scope).get_parameter(name)
                self._ema[(scope, name)].mul_(d).add_(p.detach(), alpha=1.0 - d)
            if self.learn_output_affine:
                if ("aux", "out_gain") in self._ema:
                    self._ema[("aux", "out_gain")].mul_(d).add_(self.out_gain.detach(), alpha=1.0 - d)
                if ("aux", "out_bias") in self._ema:
                    self._ema[("aux", "out_bias")].mul_(d).add_(self.out_bias.detach(), alpha=1.0 - d)

    def _trust_region_penalty(self):
        """L2 distance to initial weights."""
        if self.lambda_trust <= 0.0:
            return torch.zeros(1, device=self.device)
        loss = torch.zeros(1, device=self.device, requires_grad=True)
        for scope, name in self._names:
            p = getattr(self.model, scope).get_parameter(name)
            loss = loss + (p - self._theta0[(scope, name)]).square().sum()
        if self.learn_output_affine:
            loss = loss + (self.out_gain - 1.0).square() + self.out_bias.square()
        return loss * self.lambda_trust

    def _collapse_penalty(self, preds):
        """Penalize low diversity across batch."""
        if self.lambda_collapse <= 0.0:
            return torch.zeros(1, device=self.device)
        B, T = preds.shape
        if B < 2:
            return torch.zeros(1, device=self.device)
        diffs = preds.unsqueeze(1) - preds.unsqueeze(0)
        d = torch.sqrt(torch.mean(diffs ** 2, dim=2) + 1e-12)
        pen = torch.relu(self.collapse_margin - d)
        return torch.sigmoid(pen * 100.0).mean() * self.lambda_collapse

    def _trajectory_persistence_loss(self, preds, input_seq, input_lengths):
        """Encourage initial predicted slope to match recent history slope."""
        if self.lambda_traj <= 0.0 or self.traj_k <= 0:
            return torch.zeros(1, device=self.device)

        B, Lin, _ = input_seq.shape
        k = min(self.traj_k, preds.size(1))
        m = self.traj_recent_m

        ema_slopes = []
        for b in range(B):
            L = int(input_lengths[b].item())
            if L <= m:
                ema_slopes.append(torch.tensor(0.0, device=self.device))
                continue
            hist = input_seq[b, L - m : L, 0]
            slopes = hist[1:] - hist[:-1]
            ema = slopes[0]
            alpha = 2.0 / (m + 1)
            for s in slopes[1:]:
                ema = ema * (1.0 - alpha) + s * alpha
            ema_slopes.append(ema.clamp(-self.traj_slope_clip, self.traj_slope_clip))

        target_slope = torch.stack(ema_slopes).unsqueeze(1)
        pred_slopes = preds[:, 1:k] - preds[:, : k - 1]
        return torch.mean((pred_slopes - target_slope) ** 2) * self.lambda_traj

    def _flow_continuation_loss(self, preds, input_seq, input_lengths):
        """Encourage initial predicted curvature to match recent history curvature."""
        if self.lambda_flow <= 0.0 or self.flow_k <= 1:
            return torch.zeros(1, device=self.device)

        B, Lin, _ = input_seq.shape
        k = min(self.flow_k, preds.size(1))
        m = self.flow_recent_m

        ema_curvs = []
        for b in range(B):
            L = int(input_lengths[b].item())
            if L <= m:
                ema_curvs.append(torch.tensor(0.0, device=self.device))
                continue
            hist = input_seq[b, L - m : L, 0]
            curvs = hist[2:] - 2 * hist[1:-1] + hist[:-2]
            if curvs.numel() == 0:
                ema_curvs.append(torch.tensor(0.0, device=self.device))
                continue
            ema = curvs[0]
            alpha = 2.0 / (m + 1)
            for c in curvs[1:]:
                ema = ema * (1.0 - alpha) + c * alpha
            ema_curvs.append(ema.clamp(-self.flow_curv_clip, self.flow_curv_clip))

        target_curv = torch.stack(ema_curvs).unsqueeze(1)
        if k < 3:
            return torch.zeros(1, device=self.device)
        pred_curvs = preds[:, 2:k] - 2 * preds[:, 1 : k - 1] + preds[:, : k - 2]
        return torch.mean((pred_curvs - target_curv) ** 2) * self.lambda_flow

    def fine_tune(
        self,
        input_seq: torch.Tensor,
        input_lengths: torch.Tensor,
        group_stats: torch.Tensor,
        target_seq: torch.Tensor,
        epochs: int = 1,
        teacher_forcing_ratio: float = 1.0,
    ):
        """Run short-window fine-tuning with TF + AR paths."""
        from torch.cuda.amp import autocast, GradScaler

        self.model.train()
        device = self.device

        input_seq = input_seq.detach().to(device, dtype=torch.float32).contiguous()
        input_lengths = input_lengths.detach().to(device, dtype=torch.long).view(-1)
        group_stats = group_stats.detach().to(device, dtype=torch.float32).contiguous()
        target_seq = target_seq.detach().to(device, dtype=torch.float32).contiguous()

        B, L_out = target_seq.shape
        w = horizon_weights(L_out, device=device, gamma=self.horizon_gamma)
        ws = horizon_weights(max(L_out - 1, 1), device=device, gamma=self.horizon_gamma)

        if not hasattr(self, "_scaler"):
            self._scaler = GradScaler(enabled=torch.cuda.is_available())

        last_loss = None
        for _ in range(max(1, int(epochs))):
            self.optimizer.zero_grad(set_to_none=True)

            with autocast(dtype=torch.float16, enabled=torch.cuda.is_available()):
                preds_tf = self.model(
                    input_seq, input_lengths, group_stats,
                    target_seq=target_seq, teacher_forcing_ratio=teacher_forcing_ratio,
                )
                if preds_tf.dim() == 3:
                    preds_tf = preds_tf.squeeze(-1)
                preds_tf = self._apply_out_affine(preds_tf).contiguous()

                tgt_lens = target_seq.new_full((B,), L_out)
                loss_tf = masked_mse(preds_tf, target_seq, tgt_lens, w=w)
                loss_slope = masked_slope_mse(preds_tf, target_seq, tgt_lens, w=ws)

                with torch.no_grad():
                    preds_ar = self.model(
                        input_seq, input_lengths, group_stats,
                        target_seq=target_seq, teacher_forcing_ratio=0.0,
                    )
                    if preds_ar.dim() == 3:
                        preds_ar = preds_ar.squeeze(-1)
                    preds_ar = self._apply_out_affine(preds_ar).contiguous()
                loss_ar = masked_mse(preds_ar, target_seq, tgt_lens, w=w)

                loss = (1.0 - self.alpha_ar) * loss_tf + self.alpha_ar * loss_ar
                loss = loss + self.lambda_slope * loss_slope
                loss = loss + self._trust_region_penalty()
                loss = loss + self._collapse_penalty(preds_tf)
                loss = loss + self._trajectory_persistence_loss(preds_tf, input_seq, input_lengths)
                loss = loss + self._flow_continuation_loss(preds_tf, input_seq, input_lengths)

            self._scaler.scale(loss).backward()
            self._scaler.unscale_(self.optimizer)
            torch.nn.utils.clip_grad_norm_(self.trainable_params, self.grad_clip)
            self._scaler.step(self.optimizer)
            self._scaler.update()
            self._ema_update()
            last_loss = float(loss.item())

        return last_loss

    def aux_state(self):
        d = {}
        if self.learn_output_affine:
            d["out_gain"] = float(self.out_gain.detach().cpu())
            d["out_bias"] = float(self.out_bias.detach().cpu())
        return d


# ---------------------------------------------------------------------------
# run_finetune_cycle
# ---------------------------------------------------------------------------

def run_finetune_cycle(
    model,
    input_seq,
    input_lengths,
    group_stats,
    target_seq,
    *,
    device="cpu",
    lr=1e-4,
    epochs=1,
    teacher_forcing_ratio=1.0,
    alpha_ar=0.35,
    grad_clip=1.0,
    train_subset="head_only",
    lambda_trust=1e-4,
    ema_decay=0.98,
    horizon_gamma=0.6,
    lambda_slope=0.25,
    lambda_collapse=1e-2,
    collapse_margin=0.025,
    learn_output_affine=True,
    lambda_out_affine=1e-4,
    out_gain_bounds=(0.85, 1.20),
    out_bias_bounds=(-0.10, 0.10),
    lambda_traj=0.10,
    traj_k=20,
    traj_recent_m=7,
    traj_tau=8.0,
    traj_slope_clip=0.05,
    lambda_flow=0.15,
    flow_k=20,
    flow_recent_m=8,
    flow_tau=10.0,
    flow_curv_clip=0.02,
    rollout_gate=False,
    seed_input_seq=None,
    seed_input_lengths=None,
    eval_gt=None,
    group_stats_seed=None,
    margin_rmse_frac=0.02,
    attempts=2,
    ema_vs_prop=True,
    lr_backoff=0.5,
    **_ignored,
):
    """Run one fine-tuning cycle and return (loss_or_rmse, calibration_dict)."""
    if any(x is None for x in (input_seq, input_lengths, group_stats, target_seq)):
        return None, {}

    device = torch.device(device)
    model.to(device)

    input_seq = input_seq.to(device, dtype=torch.float32).contiguous()
    input_lengths = input_lengths.to(device, dtype=torch.long).view(-1)
    group_stats = group_stats.to(device, dtype=torch.float32).contiguous()
    target_seq = target_seq.to(device, dtype=torch.float32).contiguous()

    finetuner = OnlineFineTuner(
        model, lr=lr, grad_clip=grad_clip, alpha_ar=alpha_ar,
        horizon_gamma=horizon_gamma, lambda_slope=lambda_slope,
        lambda_trust=lambda_trust, train_subset=train_subset,
        ema_decay=ema_decay, lambda_collapse=lambda_collapse,
        collapse_margin=collapse_margin, learn_output_affine=learn_output_affine,
        lambda_out_affine=lambda_out_affine, out_gain_bounds=out_gain_bounds,
        out_bias_bounds=out_bias_bounds, lambda_traj=lambda_traj,
        traj_k=traj_k, traj_recent_m=traj_recent_m, traj_tau=traj_tau,
        traj_slope_clip=traj_slope_clip, lambda_flow=lambda_flow,
        flow_k=flow_k, flow_recent_m=flow_recent_m, flow_tau=flow_tau,
        flow_curv_clip=flow_curv_clip,
    )

    loss = finetuner.fine_tune(
        input_seq=input_seq, input_lengths=input_lengths,
        group_stats=group_stats, target_seq=target_seq,
        epochs=max(1, int(epochs)), teacher_forcing_ratio=teacher_forcing_ratio,
    )

    calib = finetuner.aux_state()
    torch.cuda.empty_cache()
    return float(loss) if loss is not None else None, calib


# ---------------------------------------------------------------------------
# Batch-building helpers
# ---------------------------------------------------------------------------

def prepare_online_finetune_batch(
    ndvi_array: np.ndarray,
    group_stats_array: np.ndarray,
    current_dap: int,
    step: int = 5,
    max_dap: int = 127,
    *,
    min_input_len: int = 20,
    min_hist_std: float | None = None,
    drop_all_zero_target: bool = False,
    return_indices: bool = False,
    verbose: bool = False,
):
    """Prepare fine-tuning batch from NDVI array at a given DAP.

    Returns (input_seq, input_lengths, group_stats, target_seq) or None.
    If return_indices=True, returns a 5th element with kept sample indices.
    """
    ndvi = np.clip(np.asarray(ndvi_array, dtype=np.float32), 0.0, 1.0)
    gs = np.asarray(group_stats_array, dtype=np.float32)

    if ndvi.ndim != 2:
        raise ValueError("ndvi_array must be (N, T)")
    N, T = ndvi.shape
    gs = gs[:, :2].astype(np.float32, copy=False)

    current_dap = int(current_dap)
    step = int(max(1, step))
    cap = int(min(max_dap, T))

    L_in = int(min(max(0, current_dap), T))
    if L_in <= 0 or current_dap >= cap:
        return None

    L_out = int(min(step, cap - current_dap))
    if L_out <= 0:
        return None

    seqs, tgts, kept_gs, kept_idx = [], [], [], []
    for i in range(N):
        hist = ndvi[i, :L_in]
        tgt = ndvi[i, current_dap : current_dap + L_out]

        if hist.size == 0 or tgt.size == 0:
            continue
        if L_in < min_input_len:
            continue
        if min_hist_std is not None and np.std(hist) < float(min_hist_std):
            continue
        if drop_all_zero_target and np.all(tgt == 0.0):
            continue

        daps = np.arange(L_in, dtype=np.float32)
        seqs.append(np.stack([hist, daps], axis=1))
        tgts.append(tgt)
        kept_gs.append(gs[i])
        kept_idx.append(i)

    if not seqs:
        return None

    input_seq = torch.from_numpy(np.stack(seqs).astype(np.float32))
    input_len = torch.full((len(seqs),), L_in, dtype=torch.long)
    group_stats = torch.from_numpy(np.stack(kept_gs).astype(np.float32))
    target_seq = torch.from_numpy(np.stack(tgts).astype(np.float32))

    if return_indices:
        return input_seq, input_len, group_stats, target_seq, np.array(kept_idx, dtype=np.int64)
    return input_seq, input_len, group_stats, target_seq


# ---------------------------------------------------------------------------
# Driver loop
# ---------------------------------------------------------------------------

def _cuda_gc():
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()


def _pad_and_cat_replay(replay_batches):
    """Concatenate replay batches on CPU with padding."""
    max_lin = max(b[0].size(1) for b in replay_batches)
    inps, lens, gss, tgts = [], [], [], []
    for inp, ilen, gs, tgt in replay_batches:
        B, Lin, C = inp.size()
        if Lin < max_lin:
            pad = torch.zeros((B, max_lin - Lin, C), dtype=inp.dtype)
            inp = torch.cat([inp, pad], dim=1)
        inps.append(inp)
        lens.append(ilen.view(-1))
        gss.append(gs)
        tgts.append(tgt)
    return torch.cat(inps), torch.cat(lens), torch.cat(gss), torch.cat(tgts)


def _build_seed_batch(ndvi, gs, idxs, seed_dap, horizon_end):
    """Build seed batch for rollout-gated evaluation."""
    ndvi = np.clip(np.asarray(ndvi, dtype=np.float32), 0.0, 1.0)
    gs = np.asarray(gs, dtype=np.float32)
    N, T = ndvi.shape
    L_seed = int(max(1, min(seed_dap, T)))
    horizon_end = int(max(L_seed + 1, min(horizon_end, T)))
    L_eval = horizon_end - L_seed
    if L_eval <= 1:
        return None

    seeds, seed_lens, gsb, gte = [], [], [], []
    for i in idxs:
        hist = ndvi[i, :L_seed]
        fut = ndvi[i, L_seed:horizon_end]
        daps = np.arange(L_seed, dtype=np.float32)
        seeds.append(np.stack([hist, daps], axis=1))
        seed_lens.append(L_seed)
        gi = gs[i].reshape(-1)[:2]
        if gi.shape[0] < 2:
            gi = np.pad(gi, (0, 2 - gi.shape[0]))
        gsb.append(gi.astype(np.float32))
        gte.append(fut.astype(np.float32))

    return (
        torch.from_numpy(np.stack(seeds)).contiguous(),
        torch.as_tensor(seed_lens, dtype=torch.long),
        torch.from_numpy(np.stack(gsb)).contiguous(),
        torch.from_numpy(np.stack(gte)).contiguous(),
    )


def online_finetune_driver(
    model,
    ndvi_array: np.ndarray,
    group_stats_array: np.ndarray,
    save_dir: str,
    *,
    dap_start: int = 5,
    dap_end: int = 120,
    dap_step: int = 5,
    forecast_step: int = 30,
    max_dap: int = 127,
    device=None,
    base_lr: float = 1e-4,
    alpha_ar: float = 0.35,
    lambda_trust: float = 1e-4,
    lambda_slope: float = 0.25,
    horizon_gamma: float = 0.6,
    lambda_collapse: float = 1e-2,
    collapse_margin: float = 0.025,
    learn_output_affine: bool = True,
    lambda_out_affine: float = 1e-4,
    out_gain_bounds: tuple = (0.85, 1.20),
    out_bias_bounds: tuple = (-0.10, 0.10),
    lambda_traj: float = 0.10,
    traj_k: int = 20,
    traj_recent_m: int = 7,
    traj_tau: float = 8.0,
    traj_slope_clip: float = 0.05,
    lambda_flow: float = 0.15,
    flow_k: int = 20,
    flow_recent_m: int = 8,
    flow_tau: float = 10.0,
    flow_curv_clip: float = 0.02,
    replay_maxlen: int = 4,
    finetune_step: int = 10,
    ema_decay: float = 0.98,
    seed_dap: int = 60,
    gate_margin_rmse_frac: float = 0.02,
    gate_attempts: int = 2,
    gate_epochs: int = 1,
    gate_lr_backoff: float = 0.5,
    train_start_mean: float = 0.225,
    verbose: bool = True,
):
    """OOM-safe online fine-tuning driver with replay buffer and rollout-gated acceptance."""
    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
    device = torch.device(device)
    model.to(device)

    os.makedirs(save_dir, exist_ok=True)
    csv_dir = os.path.join(save_dir, "csv")
    os.makedirs(csv_dir, exist_ok=True)
    ckpt_prev = os.path.join(save_dir, "_prev_cycle_weights.pt")

    replay = deque(maxlen=int(replay_maxlen))

    if verbose:
        print(f"[CFG] DAP start={dap_start} end={dap_end} step={dap_step} "
              f"max_dap={max_dap} finetune_step={finetune_step} replay={replay_maxlen}")

    for dap in range(int(dap_start), int(dap_end) + 1, int(dap_step)):
        # Stop when a full fine-tuning horizon is no longer available.
        if dap >= max_dap - 1 or (dap + int(finetune_step) > int(max_dap)):
            break

        # Vertical offset
        offset_ndvi, offset_gs, offset_value = pre_fine_tuning_offsetter(
            ndvi_array, group_stats_array, train_start_mean,
            early_window=65, current_dap=dap, verbose=False,
        )

        # Build FT batch
        ft_pack = prepare_online_finetune_batch(
            offset_ndvi, offset_gs, current_dap=dap, step=finetune_step,
            max_dap=max_dap, min_input_len=20, min_hist_std=1e-3,
            drop_all_zero_target=True, return_indices=True,
        )
        if ft_pack is None:
            if verbose:
                print(f"[SKIP] DAP {dap}: no valid FT batch.")
            continue

        inp_cpu, ilen_cpu, gs_cpu, tgt_cpu, idxs = ft_pack
        replay.append((inp_cpu, ilen_cpu, gs_cpu, tgt_cpu))

        r_inp, r_len, r_gs, r_tgt = _pad_and_cat_replay(replay)
        r_inp, r_len, r_gs, r_tgt = r_inp.to(device), r_len.to(device), r_gs.to(device), r_tgt.to(device)

        # Seed batch for rollout gate
        horizon_end = min(max_dap, dap + finetune_step)
        seed_pack = _build_seed_batch(offset_ndvi, offset_gs, idxs, seed_dap, horizon_end)
        if seed_pack is None:
            xb_seed, ln_seed, gs_seed, gt_eval = r_inp, r_len, r_gs, r_tgt
        else:
            xb_seed, ln_seed, gs_seed, gt_eval = [t.to(device) for t in seed_pack]

        # Pre-FT score
        with torch.no_grad():
            B_eval, L_eval = gt_eval.shape
            dummy = torch.zeros((B_eval, L_eval), device=device)
            pre_preds = model(xb_seed, ln_seed, gs_seed, target_seq=dummy, teacher_forcing_ratio=0.0)
            if pre_preds.dim() == 3:
                pre_preds = pre_preds.squeeze(-1)
            pre_preds = pre_preds[:, :L_eval]
        pre_preds_unshifted = post_fine_tuning_offsetter(pre_preds, offset_value)
        pre_tgts = post_fine_tuning_offsetter(gt_eval, offset_value)
        r2_pre = safe_r2(pre_tgts, pre_preds_unshifted)

        if verbose:
            print(f"\n=== DAP {dap} ===")
            print(f"[PRE]  R2 (seed {seed_dap} -> {horizon_end}) = {np.nan_to_num(r2_pre, nan=-1.0):.4f}")

        # Backup + fine-tune
        torch.save(model.state_dict(), ckpt_prev, _use_new_zipfile_serialization=False)

        best_rmse, calib = run_finetune_cycle(
            model, r_inp, r_len, r_gs, r_tgt,
            device=device, lr=base_lr, epochs=gate_epochs, teacher_forcing_ratio=1.0,
            alpha_ar=alpha_ar, grad_clip=1.0, train_subset="all",
            lambda_trust=lambda_trust, ema_decay=ema_decay,
            horizon_gamma=horizon_gamma, lambda_slope=lambda_slope,
            lambda_collapse=lambda_collapse, collapse_margin=collapse_margin,
            learn_output_affine=learn_output_affine, lambda_out_affine=lambda_out_affine,
            out_gain_bounds=out_gain_bounds, out_bias_bounds=out_bias_bounds,
            lambda_traj=lambda_traj, traj_k=traj_k, traj_recent_m=traj_recent_m,
            traj_tau=traj_tau, traj_slope_clip=traj_slope_clip,
            lambda_flow=lambda_flow, flow_k=flow_k, flow_recent_m=flow_recent_m,
            flow_tau=flow_tau, flow_curv_clip=flow_curv_clip,
            rollout_gate=True,
            seed_input_seq=xb_seed, seed_input_lengths=ln_seed,
            eval_gt=gt_eval, group_stats_seed=gs_seed,
            margin_rmse_frac=gate_margin_rmse_frac, attempts=gate_attempts,
            ema_vs_prop=True, lr_backoff=gate_lr_backoff,
        )

        if verbose and isinstance(calib, dict) and calib:
            print("[CALIB] " + " ".join(f"{k}={v:.4f}" for k, v in calib.items()))

        # Post-FT score
        with torch.no_grad():
            post_preds = model(xb_seed, ln_seed, gs_seed, target_seq=dummy, teacher_forcing_ratio=0.0)
            if post_preds.dim() == 3:
                post_preds = post_preds.squeeze(-1)
            post_preds = post_preds[:, :L_eval]

        need_rollback = not np.isfinite(to_numpy(post_preds)).all()
        if need_rollback:
            print("[ROLLBACK] Non-finite predictions; reverting weights.")
            model.load_state_dict(torch.load(ckpt_prev, map_location=device))
            post_preds = pre_preds.clone() if isinstance(pre_preds, torch.Tensor) else pre_preds

        post_preds = post_fine_tuning_offsetter(post_preds, offset_value)
        post_tgts = post_fine_tuning_offsetter(gt_eval, offset_value)
        r2_post = safe_r2(post_tgts, post_preds)
        if verbose:
            print(f"[POST] R2 (seed {seed_dap} -> {horizon_end}) = {np.nan_to_num(r2_post, nan=-1.0):.4f}")

        # Full-horizon AR from current DAP
        with torch.no_grad():
            L_full = max(1, max_dap - dap)
            B_ft = inp_cpu.size(0)
            dummy_full = torch.zeros((B_ft, L_full), device=device)
            full_preds = model(
                inp_cpu.to(device), ilen_cpu.to(device), gs_cpu.to(device),
                target_seq=dummy_full, teacher_forcing_ratio=0.0,
            )
            if full_preds.dim() == 3:
                full_preds = full_preds.squeeze(-1)
            full_preds = full_preds[:, :L_full]

        full_preds = post_fine_tuning_offsetter(full_preds, offset_value)
        full_tgts = post_fine_tuning_offsetter(
            torch.from_numpy(np.stack([offset_ndvi[idx, dap:max_dap] for idx in idxs])).to(full_preds.device),
            offset_value,
        )
        r2_full = safe_r2(full_tgts, full_preds)
        if verbose:
            print(f"[FULL] R2 (dap {dap} -> {max_dap}) = {np.nan_to_num(r2_full, nan=-1.0):.4f}")

        # Save CSVs
        save_dap_results(
            base_dir=csv_dir, dap=dap,
            ndvi_inputs=ndvi_array.astype(np.float32),
            preds=full_preds, targets=full_tgts,
        )

        # Cleanup
        del r_inp, r_len, r_gs, r_tgt, xb_seed, ln_seed, gs_seed, gt_eval, pre_preds, post_preds
        _cuda_gc()

    if verbose:
        print("\n[DONE] Online fine-tuning loop complete.")
