code/hybrid_redshift.py

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
from torch.optim import AdamW
from torch.utils.data import DataLoader, WeightedRandomSampler
from tqdm import tqdm

from .data import SpectraListDataset, collect_mmu_desi, compute_sample_stats, split_indices, valid_pixel_mask
from .metrics import LossConfig, masked_huber, redshift_losses, redshift_metrics
from .model import fourier_loglam
from .plots import plot_reconstruction_batch, plot_redshift_scatter


@dataclass
class RawCollatorConfig:
    target_length: int = 4096
    min_scale: float = 1e-3
    random_mask_ratio: float = 0.0
    eval_mask_ratio: float = 0.25
    mask_mode: str = "pixel"
    mask_span_min: int = 16
    mask_span_max: int = 64
    line_region_percentile: float = 90.0
    augment_ood: bool = False
    crop_prob: float = 0.0
    bad_window_prob: float = 0.0
    throughput_prob: float = 0.0
    noise_prob: float = 0.0
    resolution_prob: float = 0.0
    downsample_prob: float = 0.0
    line_dropout_prob: float = 0.0
    span_dropout_prob: float = 0.0
    redshift_shift: float = 0.0


class RawSpectraCollator:
    def __init__(self, cfg: RawCollatorConfig, train: bool = True, seed: int = 17):
        self.cfg = cfg
        self.train = train
        self.seed = seed
        self.rng = np.random.default_rng(seed)

    def __call__(self, samples: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
        items = [self._prepare_sample(s) for s in samples]
        x = np.stack([item["x"] for item in items], axis=0).astype(np.float32)
        valid = np.stack([item["valid"] for item in items], axis=0).astype(np.bool_)
        loglam = np.stack([item["loglam"] for item in items], axis=0).astype(np.float32)
        target_flux = np.stack([item["target_flux"] for item in items], axis=0).astype(np.float32)
        loss_mask = np.stack([item["loss_mask"] for item in items], axis=0).astype(np.bool_)
        line_weight = np.stack([item["line_weight"] for item in items], axis=0).astype(np.float32)
        line_region = np.stack([item["line_region"] for item in items], axis=0).astype(np.bool_)
        z = np.asarray([item["z"] for item in items], dtype=np.float32)
        y = np.asarray([item["y"] for item in items], dtype=np.float32)
        zwarn = np.asarray([item["zwarn"] for item in items], dtype=np.bool_)
        return {
            "x": torch.from_numpy(x),
            "valid": torch.from_numpy(valid),
            "loglam": torch.from_numpy(loglam),
            "target_flux": torch.from_numpy(target_flux),
            "loss_mask": torch.from_numpy(loss_mask),
            "line_weight": torch.from_numpy(line_weight),
            "line_region": torch.from_numpy(line_region),
            "z": torch.from_numpy(z),
            "y": torch.from_numpy(y),
            "zwarn": torch.from_numpy(zwarn),
        }

    def _prepare_sample(self, sample: dict[str, Any]) -> dict[str, Any]:
        rng = self.rng if self.train else self._eval_rng(sample)
        flux = np.asarray(sample["flux"], dtype=np.float32).copy()
        ivar = np.asarray(sample["ivar"], dtype=np.float32).copy()
        lam = np.asarray(sample["lambda"], dtype=np.float32)
        lsf = np.asarray(sample["lsf_sigma"], dtype=np.float32)
        bad = np.asarray(sample["bad_mask"], dtype=np.bool_).copy()

        if self.cfg.augment_ood:
            bad = self._augment_bad_windows(bad, rng)
            flux = self._augment_flux_calibration(flux, lam, rng)
            flux = self._augment_resolution(flux, rng)
            flux, ivar = self._augment_downsample_resample(flux, ivar, lam, rng)
            flux = self._augment_noise(flux, ivar, rng)

        valid = np.isfinite(flux) & np.isfinite(ivar) & np.isfinite(lam) & (ivar > 0) & (~bad)
        loglam = np.log(np.clip(lam.astype(np.float64), 1.0, None)).astype(np.float32)
        if valid.sum() < 16:
            valid = valid_pixel_mask(sample)

        grid = np.linspace(float(np.nanmin(loglam)), float(np.nanmax(loglam)), self.cfg.target_length, dtype=np.float32)
        flux_grid = self._interp_valid(loglam, flux, valid, grid, fill=0.0)
        ivar_grid = self._interp_valid(loglam, ivar, valid, grid, fill=0.0)
        lsf_grid = self._interp_valid(loglam, lsf, valid, grid, fill=0.0)
        valid_grid = np.interp(grid, loglam, valid.astype(np.float32), left=0.0, right=0.0) > 0.5

        center = float(np.nanmedian(flux_grid[valid_grid])) if valid_grid.any() else 0.0
        dev = np.abs(flux_grid[valid_grid] - center) if valid_grid.any() else np.asarray([1.0], dtype=np.float32)
        scale = float(np.nanmedian(dev) * 1.4826)
        if not math.isfinite(scale) or scale < self.cfg.min_scale:
            scale = max(float(np.nanmedian(np.abs(flux_grid[valid_grid]))) if valid_grid.any() else 1.0, self.cfg.min_scale)

        norm_flux = np.arcsinh((flux_grid - center) / scale).astype(np.float32)
        norm_ivar = np.log1p(np.maximum(ivar_grid * scale * scale, 0.0)).astype(np.float32)
        norm_ivar = np.clip(norm_ivar / 8.0, 0.0, 4.0)
        lsf_norm = np.nan_to_num(lsf_grid / 3.0, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32)
        loglam_norm = ((grid - math.log(6000.0)) / 0.45).astype(np.float32)

        grad = np.gradient(norm_flux, grid).astype(np.float32)
        good_grad = np.abs(grad[valid_grid])
        grad_scale = float(np.percentile(good_grad, 95)) if len(good_grad) else 1.0
        if not math.isfinite(grad_scale) or grad_scale <= 0:
            grad_scale = 1.0
        grad = np.clip(grad / grad_scale, -5.0, 5.0).astype(np.float32)
        abs_grad = np.abs(grad).astype(np.float32)

        target_flux = norm_flux.copy()
        line_weight = self._line_weights(abs_grad, valid_grid)
        line_region = self._line_region(abs_grad, valid_grid)
        corrupt = self._sample_input_dropout(abs_grad, valid_grid, rng)
        if corrupt.any():
            norm_flux = norm_flux.copy()
            grad = grad.copy()
            abs_grad = abs_grad.copy()
            norm_flux[corrupt] = 0.0
            grad[corrupt] = 0.0
            abs_grad[corrupt] = 0.0

        y = math.log1p(float(sample["z"]))
        if self.train and self.cfg.redshift_shift > 0:
            delta = float(self.rng.uniform(-self.cfg.redshift_shift, self.cfg.redshift_shift))
            y = max(0.0, y + delta)
            grid = (grid + delta).astype(np.float32)
            loglam_norm = ((grid - math.log(6000.0)) / 0.45).astype(np.float32)

        x = np.stack(
            [
                norm_flux,
                norm_ivar,
                valid_grid.astype(np.float32),
                lsf_norm,
                loglam_norm,
                grad,
                abs_grad,
                corrupt.astype(np.float32),
            ],
            axis=0,
        )
        return {
            "x": x,
            "valid": valid_grid,
            "loglam": grid,
            "target_flux": target_flux,
            "loss_mask": corrupt & valid_grid,
            "line_weight": line_weight,
            "line_region": line_region,
            "z": sample["z"],
            "y": np.float32(y),
            "zwarn": sample["zwarn"],
        }

    def _eval_rng(self, sample: dict[str, Any]) -> np.random.Generator:
        object_id = str(sample.get("object_id", ""))
        lam = np.asarray(sample["lambda"], dtype=np.float32)
        key = f"{self.seed}|{object_id}|{float(sample['z']):.8g}|{len(lam)}|{float(lam[0]):.4f}|{float(lam[-1]):.4f}"
        digest = hashlib.blake2b(key.encode("utf-8"), digest_size=8).digest()
        return np.random.default_rng(int.from_bytes(digest, "little", signed=False))

    def _interp_valid(self, x: np.ndarray, y: np.ndarray, valid: np.ndarray, x_new: np.ndarray, fill: float) -> np.ndarray:
        good = valid & np.isfinite(x) & np.isfinite(y)
        if good.sum() < 2:
            return np.full_like(x_new, fill, dtype=np.float32)
        return np.interp(x_new, x[good], y[good], left=fill, right=fill).astype(np.float32)

    def _augment_bad_windows(self, bad: np.ndarray, rng: np.random.Generator) -> np.ndarray:
        out = bad.copy()
        n = len(out)
        if rng.random() < self.cfg.crop_prob:
            frac = float(rng.uniform(0.62, 0.96))
            width = max(32, int(n * frac))
            start = int(rng.integers(0, max(1, n - width)))
            keep = np.zeros(n, dtype=np.bool_)
            keep[start : start + width] = True
            out |= ~keep
        if rng.random() < self.cfg.bad_window_prob:
            for _ in range(int(rng.integers(1, 5))):
                width = int(rng.integers(max(8, n // 240), max(12, n // 45)))
                start = int(rng.integers(0, max(1, n - width)))
                out[start : start + width] = True
        return out

    def _augment_flux_calibration(self, flux: np.ndarray, lam: np.ndarray, rng: np.random.Generator) -> np.ndarray:
        if rng.random() >= self.cfg.throughput_prob:
            return flux
        x = np.linspace(-1.0, 1.0, len(flux), dtype=np.float32)
        coeff = rng.normal(0.0, [0.05, 0.025, 0.015]).astype(np.float32)
        curve = 1.0 + coeff[0] * x + coeff[1] * (x * x - 0.33) + coeff[2] * np.sin(np.pi * x)
        return (flux * np.clip(curve, 0.65, 1.35)).astype(np.float32)

    def _augment_noise(self, flux: np.ndarray, ivar: np.ndarray, rng: np.random.Generator) -> np.ndarray:
        if rng.random() >= self.cfg.noise_prob:
            return flux
        sigma = np.zeros_like(flux, dtype=np.float32)
        good = np.isfinite(ivar) & (ivar > 0)
        sigma[good] = 1.0 / np.sqrt(np.maximum(ivar[good], 1e-8))
        scale = float(rng.uniform(0.15, 0.75))
        return (flux + rng.normal(0.0, sigma * scale).astype(np.float32)).astype(np.float32)

    def _augment_resolution(self, flux: np.ndarray, rng: np.random.Generator) -> np.ndarray:
        if rng.random() >= self.cfg.resolution_prob:
            return flux
        finite = np.isfinite(flux)
        fill = float(np.nanmedian(flux[finite])) if finite.any() else 0.0
        base = np.nan_to_num(flux, nan=fill, posinf=fill, neginf=fill).astype(np.float32)
        sigma = float(rng.uniform(0.6, 3.0))
        radius = max(2, int(math.ceil(4.0 * sigma)))
        x = np.arange(-radius, radius + 1, dtype=np.float32)
        kernel = np.exp(-0.5 * (x / sigma) ** 2)
        kernel = (kernel / kernel.sum()).astype(np.float32)
        padded = np.pad(base, (radius, radius), mode="edge")
        return np.convolve(padded, kernel, mode="valid").astype(np.float32)

    def _augment_downsample_resample(
        self,
        flux: np.ndarray,
        ivar: np.ndarray,
        lam: np.ndarray,
        rng: np.random.Generator,
    ) -> tuple[np.ndarray, np.ndarray]:
        if rng.random() >= self.cfg.downsample_prob:
            return flux, ivar
        n = len(flux)
        if n < 32:
            return flux, ivar
        factor = int(rng.choice(np.asarray([2, 3, 4, 6, 8], dtype=np.int64)))
        offset = int(rng.integers(0, factor))
        idx = np.arange(offset, n, factor, dtype=np.int64)
        if len(idx) < 4:
            return flux, ivar
        lam_good = np.asarray(lam[idx], dtype=np.float32)
        flux_good = np.asarray(flux[idx], dtype=np.float32)
        ivar_good = np.asarray(ivar[idx], dtype=np.float32)
        good = np.isfinite(lam_good) & np.isfinite(flux_good) & np.isfinite(ivar_good)
        if np.count_nonzero(good) < 4:
            return flux, ivar
        lam_good = lam_good[good]
        order = np.argsort(lam_good)
        lam_good = lam_good[order]
        flux_good = flux_good[good][order]
        ivar_good = ivar_good[good][order]
        flux_out = np.interp(lam, lam_good, flux_good, left=flux_good[0], right=flux_good[-1]).astype(np.float32)
        ivar_out = np.interp(lam, lam_good, ivar_good, left=0.0, right=0.0).astype(np.float32)
        ivar_out *= float(rng.uniform(0.25, 0.85))
        return flux_out, ivar_out

    def _sample_input_dropout(self, abs_grad: np.ndarray, valid: np.ndarray, rng: np.random.Generator) -> np.ndarray:
        corrupt = np.zeros_like(valid, dtype=np.bool_)
        if valid.sum() < 16:
            return corrupt
        n = len(valid)
        valid_idx = np.where(valid)[0]
        ratio = self.cfg.random_mask_ratio if self.train else self.cfg.eval_mask_ratio
        if ratio > 0:
            n_rand = max(1, int(round(len(valid_idx) * min(float(ratio), 1.0))))
            if self.cfg.mask_mode == "pixel":
                corrupt[rng.choice(valid_idx, size=min(n_rand, len(valid_idx)), replace=False)] = True
            else:
                line_bias = self.cfg.mask_mode in {"line_span", "mixed_span"}
                self._add_spans_to_mask(corrupt, valid, abs_grad, n_rand, rng, line_bias=line_bias)
        if self.train and rng.random() < self.cfg.span_dropout_prob:
            for _ in range(int(rng.integers(1, 4))):
                width = int(rng.integers(max(4, n // 220), max(8, n // 55)))
                start = int(rng.integers(0, max(1, n - width)))
                corrupt[start : start + width] |= valid[start : start + width]
        if self.train and rng.random() < self.cfg.line_dropout_prob:
            score = abs_grad.copy()
            score[~valid] = 0.0
            if np.count_nonzero(score) > 0:
                k = max(4, n // 96)
                peaks = np.argsort(score)[-k:]
                for j in peaks:
                    width = int(rng.integers(max(2, n // 900), max(4, n // 280)))
                    lo = max(0, int(j) - width)
                    hi = min(n, int(j) + width + 1)
                    corrupt[lo:hi] |= valid[lo:hi]
        return corrupt & valid

    def _add_spans_to_mask(
        self,
        corrupt: np.ndarray,
        valid: np.ndarray,
        abs_grad: np.ndarray,
        target_count: int,
        rng: np.random.Generator,
        *,
        line_bias: bool,
    ) -> None:
        valid_idx = np.where(valid)[0]
        if len(valid_idx) == 0:
            return
        lo_w = max(1, int(self.cfg.mask_span_min))
        hi_w = max(lo_w + 1, int(self.cfg.mask_span_max) + 1)
        probs = None
        if line_bias:
            score = abs_grad[valid_idx].astype(np.float64)
            positive = score[np.isfinite(score) & (score > 0)]
            scale = float(np.percentile(positive, 90)) if len(positive) else 1.0
            if not math.isfinite(scale) or scale <= 0:
                scale = 1.0
            score = np.clip(score / scale, 0.0, 5.0) + 0.05
            probs = score / score.sum()
        max_tries = max(32, target_count * 4)
        tries = 0
        while int(np.count_nonzero(corrupt & valid)) < target_count and tries < max_tries:
            tries += 1
            center = int(rng.choice(valid_idx, p=probs))
            width = int(rng.integers(lo_w, hi_w))
            lo = max(0, center - width // 2)
            hi = min(len(valid), lo + width)
            corrupt[lo:hi] |= valid[lo:hi]

    def _line_weights(self, abs_grad: np.ndarray, valid: np.ndarray) -> np.ndarray:
        weight = np.ones_like(abs_grad, dtype=np.float32)
        if valid.sum() < 16:
            return weight
        scale = float(np.percentile(abs_grad[valid], 90))
        if math.isfinite(scale) and scale > 0:
            weight += 2.0 * np.clip(abs_grad / scale, 0.0, 2.0)
        weight[~valid] = 1.0
        return np.clip(weight, 1.0, 5.0).astype(np.float32)

    def _line_region(self, abs_grad: np.ndarray, valid: np.ndarray) -> np.ndarray:
        region = np.zeros_like(valid, dtype=np.bool_)
        if valid.sum() < 16:
            return region
        pct = min(max(float(self.cfg.line_region_percentile), 0.0), 100.0)
        thresh = float(np.percentile(abs_grad[valid], pct))
        if math.isfinite(thresh) and thresh > 0:
            region = (abs_grad >= thresh) & valid
        return region.astype(np.bool_)


class ConvBlock(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 7, stride: int = 1, dropout: float = 0.0):
        super().__init__()
        padding = kernel_size // 2
        self.net = nn.Sequential(
            nn.Conv1d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=False),
            nn.BatchNorm1d(out_channels),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Conv1d(out_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding, bias=False),
            nn.BatchNorm1d(out_channels),
        )
        self.skip = (
            nn.Conv1d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False)
            if stride != 1 or in_channels != out_channels
            else nn.Identity()
        )
        self.act = nn.GELU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.act(self.net(x) + self.skip(x))


class LayerScaleEncoderLayer(nn.Module):
    def __init__(self, d_model: int, heads: int, dropout: float, layerscale_init: float):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)
        self.self_attn = nn.MultiheadAttention(d_model, heads, dropout=dropout, batch_first=True)
        self.dropout1 = nn.Dropout(dropout)
        self.norm2 = nn.LayerNorm(d_model)
        self.linear1 = nn.Linear(d_model, d_model * 4)
        self.linear2 = nn.Linear(d_model * 4, d_model)
        self.dropout = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.act = nn.GELU()
        init = float(layerscale_init)
        self.ls1 = nn.Parameter(torch.full((d_model,), init))
        self.ls2 = nn.Parameter(torch.full((d_model,), init))

    def forward(
        self,
        src: torch.Tensor,
        src_mask: torch.Tensor | None = None,
        src_key_padding_mask: torch.Tensor | None = None,
        is_causal: bool = False,
    ) -> torch.Tensor:
        q = self.norm1(src)
        attn, _ = self.self_attn(
            q,
            q,
            q,
            attn_mask=src_mask,
            key_padding_mask=src_key_padding_mask,
            need_weights=False,
            is_causal=is_causal,
        )
        src = src + self.ls1 * self.dropout1(attn)
        ff = self.linear2(self.dropout(self.act(self.linear1(self.norm2(src)))))
        return src + self.ls2 * self.dropout2(ff)


class HybridSpecZ(nn.Module):
    def __init__(
        self,
        in_channels: int = 8,
        d_model: int = 256,
        conv_width: int = 128,
        layers: int = 5,
        heads: int = 8,
        dropout: float = 0.1,
        fourier_freqs: int = 32,
        z_bins: int = 64,
        y_min: float = 0.0,
        y_max: float = math.log1p(6.0),
        prediction_mode: str = "regression",
        bin_temperature: float = 1.0,
        residual_scale: float = 0.06,
        candidate_topk: int = 5,
        stem_stride: int = 8,
        rec_hidden_mult: int = 0,
        rec_refine_width: int = 16,
        rec_refine_kernel: int = 5,
        layerscale_init: float = 0.0,
    ):
        super().__init__()
        allowed_modes = {"regression", "softbin", "hybrid", "bin_residual", "ranked_bin_residual"}
        if prediction_mode not in allowed_modes:
            raise ValueError(f"prediction_mode must be one of {sorted(allowed_modes)}, got {prediction_mode!r}")
        self.fourier_freqs = fourier_freqs
        self.z_bins = z_bins
        self.y_min = y_min
        self.y_max = y_max
        self.prediction_mode = prediction_mode
        self.bin_temperature = bin_temperature
        self.residual_scale = residual_scale
        self.candidate_topk = max(1, min(int(candidate_topk), z_bins))
        if stem_stride not in {4, 8}:
            raise ValueError(f"stem_stride must be 4 or 8, got {stem_stride}")
        self.stem_stride = int(stem_stride)
        self.rec_pixels_per_token = int(stem_stride)
        self.stride_stages = int(round(math.log2(self.stem_stride)))
        bin_width = (y_max - y_min) / z_bins
        centers = torch.linspace(y_min + 0.5 * bin_width, y_max - 0.5 * bin_width, z_bins)
        self.register_buffer("z_bin_centers", centers, persistent=False)

        if self.stem_stride == 8:
            self.stem = nn.Sequential(
                ConvBlock(in_channels, conv_width, kernel_size=9, stride=2, dropout=dropout * 0.5),
                ConvBlock(conv_width, conv_width, kernel_size=7, stride=2, dropout=dropout * 0.5),
                ConvBlock(conv_width, d_model, kernel_size=7, stride=2, dropout=dropout * 0.5),
                ConvBlock(d_model, d_model, kernel_size=5, stride=1, dropout=dropout * 0.5),
            )
        else:
            self.stem = nn.Sequential(
                ConvBlock(in_channels, conv_width, kernel_size=9, stride=2, dropout=dropout * 0.5),
                ConvBlock(conv_width, d_model, kernel_size=7, stride=2, dropout=dropout * 0.5),
                ConvBlock(d_model, d_model, kernel_size=5, stride=1, dropout=dropout * 0.5),
            )
        self.pos_proj = nn.Sequential(nn.Linear(fourier_freqs * 2, d_model), nn.LayerNorm(d_model))
        self.cls = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)
        # The model never receives true z; this learned query is the always-masked z token.
        self.z_query = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

        if layerscale_init > 0:
            enc_layer = LayerScaleEncoderLayer(d_model, heads, dropout, layerscale_init)
        else:
            enc_layer = nn.TransformerEncoderLayer(
                d_model=d_model,
                nhead=heads,
                dim_feedforward=d_model * 4,
                dropout=dropout,
                batch_first=True,
                norm_first=True,
                activation="gelu",
            )
        self.encoder = nn.TransformerEncoder(enc_layer, num_layers=layers)
        self.pool_gate = nn.Sequential(nn.LayerNorm(d_model), nn.Linear(d_model, 1))
        head_dim = d_model * 5
        self.z_head = nn.Sequential(nn.LayerNorm(head_dim), nn.Linear(head_dim, d_model), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_model, 2))
        self.z_bin_head = nn.Sequential(nn.LayerNorm(head_dim), nn.Linear(head_dim, d_model), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_model, z_bins))
        self.z_candidate_head = nn.Sequential(
            nn.LayerNorm(head_dim),
            nn.Linear(head_dim, d_model),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_model, z_bins),
        )
        if rec_hidden_mult > 0:
            rec_hidden = int(d_model * rec_hidden_mult)
            self.rec_head = nn.Sequential(
                nn.LayerNorm(d_model),
                nn.Linear(d_model, rec_hidden),
                nn.GELU(),
                nn.Dropout(dropout),
                nn.Linear(rec_hidden, self.rec_pixels_per_token),
            )
        else:
            self.rec_head = nn.Sequential(nn.LayerNorm(d_model), nn.Linear(d_model, self.rec_pixels_per_token))
        rec_pad = int(rec_refine_kernel) // 2
        self.rec_refine = nn.Sequential(
            nn.Conv1d(1, rec_refine_width, kernel_size=rec_refine_kernel, padding=rec_pad),
            nn.GELU(),
            nn.Conv1d(rec_refine_width, 1, kernel_size=rec_refine_kernel, padding=rec_pad),
        )

    def forward(self, x: torch.Tensor, valid: torch.Tensor, loglam: torch.Tensor) -> dict[str, torch.Tensor]:
        bsz = x.shape[0]
        h = self.stem(x).transpose(1, 2)
        tok_valid = valid.float().unsqueeze(1)
        tok_loglam = loglam.unsqueeze(1)
        for _ in range(self.stride_stages):
            tok_valid = F.avg_pool1d(tok_valid, kernel_size=2, stride=2, ceil_mode=True)
            tok_loglam = F.avg_pool1d(tok_loglam, kernel_size=2, stride=2, ceil_mode=True)
        tok_valid = tok_valid.squeeze(1) > 0.20
        tok_loglam = tok_loglam.squeeze(1)
        if tok_valid.shape[1] != h.shape[1]:
            tok_valid = tok_valid[:, : h.shape[1]]
            tok_loglam = tok_loglam[:, : h.shape[1]]
            h = h[:, : tok_valid.shape[1]]

        h = h + self.pos_proj(fourier_loglam(tok_loglam, self.fourier_freqs))
        cls = self.cls.expand(bsz, -1, -1)
        z_query = self.z_query.expand(bsz, -1, -1)
        src = torch.cat([cls, z_query, h], dim=1)
        special_valid = torch.ones((bsz, 2), dtype=torch.bool, device=x.device)
        src_valid = torch.cat([special_valid, tok_valid], dim=1)
        padding = ~src_valid
        memory = self.encoder(src, src_key_padding_mask=padding)

        spec = memory[:, 2:]
        spec_valid = src_valid[:, 2:]
        spec_mask = spec_valid.unsqueeze(-1)
        rec = self.rec_head(spec).reshape(bsz, -1)
        rec = rec + self.rec_refine(rec.unsqueeze(1)).squeeze(1)
        if rec.shape[1] > x.shape[-1]:
            rec = rec[:, : x.shape[-1]]
        elif rec.shape[1] < x.shape[-1]:
            rec = F.pad(rec, (0, x.shape[-1] - rec.shape[1]))
        denom = spec_valid.float().sum(dim=1).clamp_min(1.0).unsqueeze(-1)
        mean_pool = (spec * spec_mask.float()).sum(dim=1) / denom
        max_pool = spec.masked_fill(~spec_mask, -1e4).max(dim=1).values
        gate_logits = self.pool_gate(spec).squeeze(-1).masked_fill(~spec_valid, -1e4)
        gate = torch.softmax(gate_logits, dim=1)
        attn_pool = torch.einsum("bn,bnd->bd", gate, spec)
        feat = torch.cat([memory[:, 0], memory[:, 1], mean_pool, max_pool, attn_pool], dim=-1)
        z_params = self.z_head(feat)
        z_bin_logits = self.z_bin_head(feat)
        candidate_residual = self.residual_scale * torch.tanh(self.z_candidate_head(feat))
        centers = self.z_bin_centers.to(dtype=z_bin_logits.dtype, device=z_bin_logits.device)
        candidate_y = (centers.unsqueeze(0) + candidate_residual).clamp(self.y_min, self.y_max)
        topk_logits, topk_bins = torch.topk(z_bin_logits, k=self.candidate_topk, dim=-1)
        candidate_topk_y = candidate_y.gather(1, topk_bins)
        y_reg = z_params[:, 0]
        bin_prob = torch.softmax(z_bin_logits / max(self.bin_temperature, 1e-4), dim=-1)
        y_bin = (bin_prob * self.z_bin_centers.to(dtype=bin_prob.dtype, device=bin_prob.device)).sum(dim=-1)
        y_ranked = (bin_prob * candidate_y.to(dtype=bin_prob.dtype)).sum(dim=-1)
        y_legacy_bin_residual = y_bin + self.residual_scale * torch.tanh(y_reg)
        if self.prediction_mode == "regression":
            y_pred = y_reg
        elif self.prediction_mode == "softbin":
            y_pred = y_bin
        elif self.prediction_mode == "hybrid":
            y_pred = 0.35 * y_reg + 0.65 * y_bin
        elif self.prediction_mode == "ranked_bin_residual":
            y_pred = 0.5 * y_legacy_bin_residual + 0.5 * y_ranked
        else:
            y_pred = y_legacy_bin_residual
        y_pred = y_pred.clamp(self.y_min, self.y_max)
        return {
            "rec": rec,
            "y_mu": y_pred,
            "y_pred": y_pred,
            "y_reg": y_reg,
            "y_bin": y_bin,
            "y_ranked": y_ranked,
            "y_top1_candidate": candidate_topk_y[:, 0],
            "y_logvar": torch.clamp(z_params[:, 1], -8.0, 4.0),
            "z_bin_logits": z_bin_logits,
            "z_feat": feat,
            "candidate_y": candidate_y,
            "candidate_topk_y": candidate_topk_y,
            "candidate_topk_bins": topk_bins,
            "candidate_topk_logits": topk_logits,
        }

    def y_to_bin(self, y: torch.Tensor) -> torch.Tensor:
        scaled = (y - self.y_min) / max(self.y_max - self.y_min, 1e-6)
        return torch.clamp((scaled * self.z_bins).long(), 0, self.z_bins - 1)


def move_to_device(batch: dict[str, torch.Tensor], device: torch.device) -> dict[str, torch.Tensor]:
    return {k: v.to(device, non_blocking=True) if torch.is_tensor(v) else v for k, v in batch.items()}


def limit_batch_examples(batch: dict[str, torch.Tensor], max_examples: int | None, seen_examples: int) -> dict[str, torch.Tensor] | None:
    if max_examples is None or max_examples <= 0:
        return batch
    remaining = int(max_examples) - int(seen_examples)
    if remaining <= 0:
        return None
    bsz = int(batch["y"].shape[0])
    if remaining >= bsz:
        return batch
    return {k: v[:remaining] if torch.is_tensor(v) and v.shape[:1] == (bsz,) else v for k, v in batch.items()}


def load_checkpoint_into_model(model: nn.Module, state: dict[str, torch.Tensor], allow_mismatched: bool = False) -> None:
    if not allow_mismatched:
        try:
            model.load_state_dict(state, strict=True)
        except RuntimeError:
            missing, unexpected = model.load_state_dict(state, strict=False)
            print(f"RESUME_NONSTRICT missing={list(missing)} unexpected={list(unexpected)}")
        return

    target_state = model.state_dict()
    compatible = {}
    skipped = []
    for key, value in state.items():
        target = target_state.get(key)
        if target is not None and tuple(target.shape) == tuple(value.shape):
            compatible[key] = value
        else:
            skipped.append(key)
    missing, unexpected = model.load_state_dict(compatible, strict=False)
    print(
        "RESUME_FILTERED "
        f"loaded={len(compatible)} skipped={len(skipped)} "
        f"missing={list(missing)} unexpected={list(unexpected)} skipped_keys={skipped[:20]}"
    )


def redshift_total_loss(model: HybridSpecZ, out: dict[str, torch.Tensor], batch: dict[str, torch.Tensor], cfg: LossConfig) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
    parts = redshift_losses(model, out, batch["y"], batch["zwarn"], cfg)
    if "rec" in out and "target_flux" in batch and "loss_mask" in batch:
        line_weight = batch.get("line_weight")
        if line_weight is not None:
            line_weight = line_weight.pow(cfg.line_weight_power)
        rec = masked_huber(out["rec"], batch["target_flux"], batch["loss_mask"], weight=line_weight)
    else:
        rec = parts["z_huber"].sum() * 0.0
    total = (
        cfg.rec_weight * rec
        + cfg.z_weight * parts["z_huber"]
        + cfg.z_bin_weight * parts["z_bin"]
        + cfg.z_candidate_weight * parts["z_candidate"]
        + cfg.z_nll_weight * parts["z_nll"]
    )
    metrics = {"loss": total.detach(), "rec": rec.detach(), **{k: v.detach() for k, v in parts.items()}}
    return total, metrics


def plot_spectra_batch(path: str | Path, batch: dict[str, torch.Tensor], y_pred: np.ndarray, max_items: int = 4) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    x = batch["x"].detach().cpu().numpy()
    loglam = batch["loglam"].detach().cpu().numpy()
    valid = batch["valid"].detach().cpu().numpy()
    z = batch["z"].detach().cpu().numpy()
    bsz = min(max_items, x.shape[0])
    fig, axes = plt.subplots(bsz, 1, figsize=(13, 3.0 * bsz), squeeze=False)
    for i in range(bsz):
        ax = axes[i, 0]
        wave = np.exp(loglam[i])
        good = valid[i].astype(bool)
        ax.plot(wave[good], x[i, 0, good], color="black", linewidth=0.8, label="input flux")
        ax.plot(wave[good], x[i, 6, good], color="#1f77b4", linewidth=0.6, alpha=0.55, label="line score")
        masked = x[i, 7] > 0
        if masked.any():
            ax.scatter(wave[masked], np.zeros(masked.sum()), s=5, color="#d62728", alpha=0.55, label="redshift dropout")
        ax.set_title(f"z true={z[i]:.5f}  z pred={np.expm1(y_pred[i]):.5f}")
        ax.set_ylabel("normalized")
        ax.grid(alpha=0.2)
        if i == 0:
            ax.legend(loc="best", fontsize=8)
    axes[-1, 0].set_xlabel("wavelength Angstrom")
    fig.tight_layout()
    fig.savefig(path, dpi=150)
    plt.close(fig)


def add_redshift_slice_metrics(metrics: dict[str, float], prefix: str, y_true: np.ndarray, y_pred: np.ndarray) -> None:
    z_true = np.expm1(y_true)
    z_pred = np.expm1(y_pred)
    slices = {
        "z_lt_0p4": z_true < 0.4,
        "z_0p4_1p0": (z_true >= 0.4) & (z_true < 1.0),
        "z_1p0_2p0": (z_true >= 1.0) & (z_true < 2.0),
        "z_gte_2p0": z_true >= 2.0,
    }
    for name, mask in slices.items():
        count = int(np.count_nonzero(mask))
        metrics[f"{prefix}/{name}_count"] = float(count)
        if count >= 5:
            err = z_pred[mask] - z_true[mask]
            denom = 1.0 + z_true[mask]
            metrics[f"{prefix}/{name}_mae_z"] = float(np.mean(np.abs(err)))
            metrics[f"{prefix}/{name}_bias_z"] = float(np.mean(err))
            metrics[f"{prefix}/{name}_cat_0p05"] = float(np.mean(np.abs(err / denom) > 0.05))


def add_candidate_metrics(
    metrics: dict[str, float],
    prefix: str,
    y_true: np.ndarray,
    candidate_y: np.ndarray,
    candidate_bins: np.ndarray | None,
    *,
    z_bins: int,
    y_min: float,
    y_max: float,
) -> None:
    if candidate_y.size == 0:
        return
    z_true = np.expm1(y_true)
    z_candidate = np.expm1(candidate_y)
    abs_dz = np.abs(z_candidate - z_true[:, None])
    norm_dz = abs_dz / (1.0 + z_true[:, None])
    top_limits = [1, 3, 5]
    for k in top_limits:
        kk = min(k, candidate_y.shape[1])
        best_abs = np.min(abs_dz[:, :kk], axis=1)
        best_norm = np.min(norm_dz[:, :kk], axis=1)
        metrics[f"{prefix}/candidate_top{kk}_best_mae_z"] = float(np.mean(best_abs))
        metrics[f"{prefix}/candidate_top{kk}_hit_0p003"] = float(np.mean(best_norm <= 0.003))
        metrics[f"{prefix}/candidate_top{kk}_hit_0p01"] = float(np.mean(best_norm <= 0.01))
        metrics[f"{prefix}/candidate_top{kk}_hit_0p05"] = float(np.mean(best_norm <= 0.05))
    if candidate_bins is not None and candidate_bins.size:
        scaled = (y_true - y_min) / max(y_max - y_min, 1e-6)
        true_bins = np.clip((scaled * z_bins).astype(np.int64), 0, z_bins - 1)
        for k in top_limits:
            kk = min(k, candidate_bins.shape[1])
            metrics[f"{prefix}/candidate_top{kk}_bin_hit"] = float(np.mean(np.any(candidate_bins[:, :kk] == true_bins[:, None], axis=1)))


@torch.no_grad()
def evaluate(
    model: HybridSpecZ,
    loader: DataLoader,
    loss_cfg: LossConfig,
    device: torch.device,
    run_dir: Path,
    step: int,
    prefix: str = "val",
    max_batches: int | None = 50,
    max_examples: int | None = None,
) -> dict[str, float]:
    model.eval()
    losses: dict[str, list[float]] = {}
    y_true_all: list[np.ndarray] = []
    y_pred_all: list[np.ndarray] = []
    candidate_y_all: list[np.ndarray] = []
    candidate_bins_all: list[np.ndarray] = []
    y_true_clean: list[np.ndarray] = []
    y_pred_clean: list[np.ndarray] = []
    candidate_y_clean: list[np.ndarray] = []
    candidate_bins_clean: list[np.ndarray] = []
    zwarn_all: list[np.ndarray] = []
    first_batch = None
    first_pred = None
    first_rec = None
    seen_examples = 0
    for bi, batch in enumerate(loader):
        if max_batches is not None and max_batches > 0 and bi >= max_batches:
            break
        batch = limit_batch_examples(batch, max_examples, seen_examples)
        if batch is None:
            break
        seen_examples += int(batch["y"].shape[0])
        batch = move_to_device(batch, device)
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"):
            out = model(batch["x"], batch["valid"], batch["loglam"])
            _, parts = redshift_total_loss(model, out, batch, loss_cfg)
        y_pred = out.get("y_pred", out["y_mu"])
        for k, v in parts.items():
            losses.setdefault(k, []).append(float(v.detach().cpu()))
        if "rec" in out and "target_flux" in batch and "loss_mask" in batch:
            rec_err = F.smooth_l1_loss(out["rec"].float(), batch["target_flux"].float(), reduction="none", beta=0.5)
            loss_mask = batch["loss_mask"].bool()
            line_region = batch.get("line_region")
            if line_region is not None:
                line_mask = loss_mask & line_region.bool()
                cont_mask = loss_mask & (~line_region.bool())
                for name, mask in (("rec_line", line_mask), ("rec_continuum", cont_mask)):
                    denom = mask.float().sum()
                    if float(denom.detach().cpu()) > 0:
                        losses.setdefault(name, []).append(float(((rec_err * mask.float()).sum() / denom.clamp_min(1.0)).detach().cpu()))
            context_mask = batch["valid"].bool() & (~loss_mask)
            denom = context_mask.float().sum(dim=1).clamp_min(1.0)
            baseline = (batch["target_flux"].float() * context_mask.float()).sum(dim=1, keepdim=True) / denom.unsqueeze(1)
            baseline_err = F.smooth_l1_loss(baseline.expand_as(batch["target_flux"]).float(), batch["target_flux"].float(), reduction="none", beta=0.5)
            mask_denom = loss_mask.float().sum().clamp_min(1.0)
            losses.setdefault("rec_mean_baseline", []).append(float(((baseline_err * loss_mask.float()).sum() / mask_denom).detach().cpu()))
        finite = torch.isfinite(batch["y"]).detach().cpu().numpy()
        clean = ((~batch["zwarn"].bool()) & torch.isfinite(batch["y"])).detach().cpu().numpy()
        zw = batch["zwarn"].detach().cpu().numpy().astype(bool)
        if finite.any():
            y_true_all.append(batch["y"].detach().cpu().numpy()[finite])
            y_pred_all.append(y_pred.float().detach().cpu().numpy()[finite])
            zwarn_all.append(zw[finite])
            if "candidate_topk_y" in out:
                candidate_y_all.append(out["candidate_topk_y"].float().detach().cpu().numpy()[finite])
            if "candidate_topk_bins" in out:
                candidate_bins_all.append(out["candidate_topk_bins"].detach().cpu().numpy()[finite])
        if clean.any():
            y_true_clean.append(batch["y"].detach().cpu().numpy()[clean])
            y_pred_clean.append(y_pred.float().detach().cpu().numpy()[clean])
            if "candidate_topk_y" in out:
                candidate_y_clean.append(out["candidate_topk_y"].float().detach().cpu().numpy()[clean])
            if "candidate_topk_bins" in out:
                candidate_bins_clean.append(out["candidate_topk_bins"].detach().cpu().numpy()[clean])
        if first_batch is None:
            first_batch = {k: v.detach().cpu() if torch.is_tensor(v) else v for k, v in batch.items()}
            first_pred = y_pred.float().detach().cpu().numpy()
            if "rec" in out:
                first_rec = out["rec"].float().detach().cpu().numpy()

    metrics = {f"{prefix}/{k}": float(np.mean(v)) for k, v in losses.items()}
    if y_true_all:
        y_true = np.concatenate(y_true_all)
        y_pred = np.concatenate(y_pred_all)
        for k, v in redshift_metrics(y_true, y_pred).items():
            metrics[f"{prefix}/{k}"] = v
        add_redshift_slice_metrics(metrics, prefix, y_true, y_pred)
        if candidate_y_all:
            candidate_y_np = np.concatenate(candidate_y_all)
            candidate_bins_np = np.concatenate(candidate_bins_all) if candidate_bins_all else None
            add_candidate_metrics(
                metrics,
                prefix,
                y_true,
                candidate_y_np,
                candidate_bins_np,
                z_bins=model.z_bins,
                y_min=model.y_min,
                y_max=model.y_max,
            )
        metrics[f"{prefix}/z_count"] = float(len(y_true))
        metrics[f"{prefix}/zwarn_fraction"] = float(np.mean(np.concatenate(zwarn_all))) if zwarn_all else 0.0
        plot_redshift_scatter(run_dir / "plots" / f"{prefix}_redshift_step_{step:06d}.png", y_true, y_pred)
    if y_true_clean:
        clean_true = np.concatenate(y_true_clean)
        clean_pred = np.concatenate(y_pred_clean)
        if len(clean_true) >= 5:
            for k, v in redshift_metrics(clean_true, clean_pred).items():
                metrics[f"{prefix}_clean/{k}"] = v
            if candidate_y_clean:
                candidate_y_clean_np = np.concatenate(candidate_y_clean)
                candidate_bins_clean_np = np.concatenate(candidate_bins_clean) if candidate_bins_clean else None
                add_candidate_metrics(
                    metrics,
                    f"{prefix}_clean",
                    clean_true,
                    candidate_y_clean_np,
                    candidate_bins_clean_np,
                    z_bins=model.z_bins,
                    y_min=model.y_min,
                    y_max=model.y_max,
                )
            metrics[f"{prefix}_clean/z_count"] = float(len(clean_true))
    if first_batch is not None and first_pred is not None:
        if first_rec is not None and "target_flux" in first_batch and "loss_mask" in first_batch:
            plot_reconstruction_batch(
                run_dir / "plots" / f"{prefix}_reconstruction_step_{step:06d}.png",
                first_batch["loglam"].numpy(),
                first_batch["target_flux"].numpy(),
                first_rec,
                first_batch["loss_mask"].numpy(),
                first_batch["valid"].numpy(),
                first_batch["z"].numpy(),
                np.expm1(first_pred),
            )
        plot_spectra_batch(run_dir / "plots" / f"{prefix}_spectra_step_{step:06d}.png", first_batch, first_pred)
    model.train()
    return metrics


def make_loader(
    samples: list[dict[str, Any]],
    indices: np.ndarray,
    cfg: RawCollatorConfig,
    args: argparse.Namespace,
    train: bool,
    sampler: WeightedRandomSampler | None = None,
) -> DataLoader:
    return DataLoader(
        SpectraListDataset(samples, indices),
        batch_size=args.batch_size,
        shuffle=train and sampler is None,
        sampler=sampler,
        num_workers=args.num_workers,
        pin_memory=True,
        collate_fn=RawSpectraCollator(cfg, train=train, seed=args.seed + (0 if train else 1000)),
    )


def checkpoint_score(mode: str, val_metrics: dict[str, float], ood_metrics: dict[str, float] | None, z_alpha: float = 0.6) -> float:
    def score_prefix(metrics: dict[str, float], prefix: str) -> float:
        z_score = (
            metrics.get(f"{prefix}/nmad", math.inf)
            + metrics.get(f"{prefix}/cat_0p01", 1.0)
            + metrics.get(f"{prefix}/mae_log1p", 1.0)
        )
        rec_score = metrics.get(f"{prefix}/rec")
        if rec_score is None or not math.isfinite(float(rec_score)):
            return z_score
        alpha = min(max(float(z_alpha), 0.0), 1.0)
        return alpha * z_score + (1.0 - alpha) * float(rec_score)

    val_score = score_prefix(val_metrics, "val")
    if mode == "rec":
        return float(val_metrics.get("val/rec", math.inf))
    if mode == "val" or ood_metrics is None:
        return val_score
    ood_score = score_prefix(ood_metrics, "ood")
    if mode == "ood":
        return ood_score
    return 0.5 * val_score + 0.5 * ood_score


def scheduled_lr(base_lr: float, min_lr: float, step: int, total_steps: int, warmup_steps: int) -> float:
    if warmup_steps > 0 and step <= warmup_steps:
        return base_lr * float(step) / float(max(1, warmup_steps))
    if min_lr < 0 or total_steps <= warmup_steps:
        return base_lr
    progress = (step - warmup_steps) / float(max(1, total_steps - warmup_steps))
    progress = min(max(progress, 0.0), 1.0)
    return min_lr + 0.5 * (base_lr - min_lr) * (1.0 + math.cos(math.pi * progress))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset-name", default="MultimodalUniverse/desi")
    parser.add_argument("--max-samples", type=int, default=4096)
    parser.add_argument("--cache-dir", default="/workspace/native_specz_mae/cache")
    parser.add_argument("--hf-cache-dir", default=os.environ.get("HF_DATASETS_CACHE", "/workspace/hf_cache/datasets"))
    parser.add_argument("--run-dir", default="/workspace/runs/hybrid_specz")
    parser.add_argument("--resume-checkpoint", default="")
    parser.add_argument("--allow-mismatched-checkpoint", action="store_true")
    parser.add_argument("--refresh-data", action="store_true")
    parser.add_argument("--epochs", type=int, default=8)
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--num-workers", type=int, default=2)
    parser.add_argument("--target-length", type=int, default=4096)
    parser.add_argument("--d-model", type=int, default=256)
    parser.add_argument("--conv-width", type=int, default=128)
    parser.add_argument("--layers", type=int, default=5)
    parser.add_argument("--heads", type=int, default=8)
    parser.add_argument("--dropout", type=float, default=0.1)
    parser.add_argument("--z-bins", type=int, default=64)
    parser.add_argument("--stem-stride", type=int, choices=[4, 8], default=8)
    parser.add_argument("--rec-hidden-mult", type=int, default=0)
    parser.add_argument("--rec-refine-width", type=int, default=16)
    parser.add_argument("--rec-refine-kernel", type=int, default=5)
    parser.add_argument("--layerscale-init", type=float, default=0.0)
    parser.add_argument(
        "--prediction-mode",
        choices=["regression", "softbin", "hybrid", "bin_residual", "ranked_bin_residual"],
        default="regression",
    )
    parser.add_argument("--bin-temperature", type=float, default=1.0)
    parser.add_argument("--residual-scale", type=float, default=0.06)
    parser.add_argument("--candidate-topk", type=int, default=5)
    parser.add_argument("--lr", type=float, default=2e-4)
    parser.add_argument("--min-lr", type=float, default=-1.0)
    parser.add_argument("--warmup-steps", type=int, default=0)
    parser.add_argument("--weight-decay", type=float, default=0.03)
    parser.add_argument("--grad-clip", type=float, default=1.0)
    parser.add_argument("--grad-accum-steps", type=int, default=1)
    parser.add_argument("--eval-every", type=int, default=100)
    parser.add_argument("--eval-max-val", type=int, default=800)
    parser.add_argument("--eval-max-ood", type=int, default=480)
    parser.add_argument("--max-steps", type=int, default=0)
    parser.add_argument("--checkpoint-score", choices=["val", "ood", "combined", "rec"], default="combined")
    parser.add_argument("--score-z-alpha", type=float, default=0.6)
    parser.add_argument("--objective", choices=["joint", "rec_only", "z_only"], default="joint")
    parser.add_argument("--balance-redshift", action="store_true")
    parser.add_argument("--train-clean-only", action="store_true")
    parser.add_argument("--clean-sample-boost", type=float, default=1.0)
    parser.add_argument("--augment-ood", action="store_true")
    parser.add_argument("--eval-ood", action="store_true")
    parser.add_argument("--random-mask-ratio", type=float, default=0.0)
    parser.add_argument("--eval-mask-ratio", type=float, default=0.25)
    parser.add_argument("--mask-mode", choices=["pixel", "span", "line_span", "mixed_span"], default="pixel")
    parser.add_argument("--mask-span-min", type=int, default=16)
    parser.add_argument("--mask-span-max", type=int, default=64)
    parser.add_argument("--line-region-percentile", type=float, default=90.0)
    parser.add_argument("--crop-prob", type=float, default=0.0)
    parser.add_argument("--bad-window-prob", type=float, default=0.0)
    parser.add_argument("--throughput-prob", type=float, default=0.0)
    parser.add_argument("--noise-prob", type=float, default=0.0)
    parser.add_argument("--resolution-prob", type=float, default=0.0)
    parser.add_argument("--downsample-prob", type=float, default=0.0)
    parser.add_argument("--line-dropout-prob", type=float, default=0.0)
    parser.add_argument("--span-dropout-prob", type=float, default=0.0)
    parser.add_argument("--redshift-shift", type=float, default=0.0)
    parser.add_argument("--rec-weight", type=float, default=0.0)
    parser.add_argument("--z-weight", type=float, default=1.0)
    parser.add_argument("--z-bin-weight", type=float, default=0.25)
    parser.add_argument("--z-candidate-weight", type=float, default=0.0)
    parser.add_argument("--z-nll-weight", type=float, default=0.05)
    parser.add_argument("--zwarn-weight", type=float, default=0.3)
    parser.add_argument("--high-z-boost", type=float, default=1.0)
    parser.add_argument("--high-z-threshold", type=float, default=1.0)
    parser.add_argument("--clean-z-only", action="store_true")
    parser.add_argument("--seed", type=int, default=17)
    args = parser.parse_args()

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    run_dir = Path(args.run_dir) / time.strftime("%Y%m%d_%H%M%S")
    run_dir.mkdir(parents=True, exist_ok=True)
    (run_dir / "args.json").write_text(json.dumps(vars(args), indent=2), encoding="utf-8")

    samples = collect_mmu_desi(
        cache_file=Path(args.cache_dir) / f"desi_{args.max_samples}.pt",
        max_samples=args.max_samples,
        dataset_name=args.dataset_name,
        hf_cache_dir=args.hf_cache_dir,
        refresh=args.refresh_data,
    )
    stats = compute_sample_stats(samples)
    (run_dir / "data_stats.json").write_text(json.dumps(stats.__dict__, indent=2), encoding="utf-8")
    print("DATA_STATS", json.dumps(stats.__dict__, sort_keys=True))

    train_idx, val_idx = split_indices(len(samples), val_fraction=0.15, seed=args.seed)
    if args.train_clean_only:
        clean_train = np.asarray([not bool(samples[int(i)]["zwarn"]) for i in train_idx], dtype=np.bool_)
        train_idx = train_idx[clean_train]
        if len(train_idx) == 0:
            raise RuntimeError("No clean ZWARN==0 samples are available for --train-clean-only.")
        print(f"TRAIN_CLEAN_ONLY n_train={len(train_idx)}")

    sampler = None
    if args.balance_redshift or args.clean_sample_boost != 1.0:
        weights = np.ones(len(train_idx), dtype=np.float32)
        y_train = np.asarray([np.log1p(float(samples[int(i)]["z"])) for i in train_idx], dtype=np.float32)
        if args.balance_redshift:
            bins = np.linspace(float(y_train.min()), float(y_train.max()) + 1e-6, 28)
            bin_id = np.clip(np.digitize(y_train, bins) - 1, 0, len(bins) - 2)
            counts = np.bincount(bin_id, minlength=len(bins) - 1).astype(np.float32)
            weights *= 1.0 / np.maximum(counts[bin_id], 1.0)
        if args.clean_sample_boost != 1.0:
            clean = np.asarray([not bool(samples[int(i)]["zwarn"]) for i in train_idx], dtype=np.bool_)
            weights *= np.where(clean, float(args.clean_sample_boost), 1.0).astype(np.float32)
        weights = weights / weights.mean()
        sampler = WeightedRandomSampler(torch.as_tensor(weights, dtype=torch.double), num_samples=len(weights), replacement=True)

    train_cfg = RawCollatorConfig(
        target_length=args.target_length,
        random_mask_ratio=args.random_mask_ratio,
        eval_mask_ratio=args.eval_mask_ratio,
        mask_mode=args.mask_mode,
        mask_span_min=args.mask_span_min,
        mask_span_max=args.mask_span_max,
        line_region_percentile=args.line_region_percentile,
        augment_ood=args.augment_ood,
        crop_prob=args.crop_prob,
        bad_window_prob=args.bad_window_prob,
        throughput_prob=args.throughput_prob,
        noise_prob=args.noise_prob,
        resolution_prob=args.resolution_prob,
        downsample_prob=args.downsample_prob,
        line_dropout_prob=args.line_dropout_prob,
        span_dropout_prob=args.span_dropout_prob,
        redshift_shift=args.redshift_shift,
    )
    val_cfg = RawCollatorConfig(
        target_length=args.target_length,
        eval_mask_ratio=args.eval_mask_ratio,
        mask_mode=args.mask_mode,
        mask_span_min=args.mask_span_min,
        mask_span_max=args.mask_span_max,
        line_region_percentile=args.line_region_percentile,
    )
    ood_cfg = RawCollatorConfig(
        target_length=args.target_length,
        eval_mask_ratio=args.eval_mask_ratio,
        mask_mode=args.mask_mode,
        mask_span_min=args.mask_span_min,
        mask_span_max=args.mask_span_max,
        line_region_percentile=args.line_region_percentile,
        augment_ood=True,
        crop_prob=0.65,
        bad_window_prob=0.45,
        throughput_prob=0.65,
        noise_prob=0.35,
        resolution_prob=0.45,
        downsample_prob=0.35,
    )
    train_loader = make_loader(samples, train_idx, train_cfg, args, train=True, sampler=sampler)
    val_loader = make_loader(samples, val_idx, val_cfg, args, train=False)
    ood_loader = make_loader(samples, val_idx, ood_cfg, args, train=False) if args.eval_ood else None

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = HybridSpecZ(
        d_model=args.d_model,
        conv_width=args.conv_width,
        layers=args.layers,
        heads=args.heads,
        dropout=args.dropout,
        z_bins=args.z_bins,
        stem_stride=args.stem_stride,
        rec_hidden_mult=args.rec_hidden_mult,
        rec_refine_width=args.rec_refine_width,
        rec_refine_kernel=args.rec_refine_kernel,
        layerscale_init=args.layerscale_init,
        prediction_mode=args.prediction_mode,
        bin_temperature=args.bin_temperature,
        residual_scale=args.residual_scale,
        candidate_topk=args.candidate_topk,
    ).to(device)
    if args.resume_checkpoint:
        ckpt = torch.load(args.resume_checkpoint, map_location=device, weights_only=False)
        state = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt
        load_checkpoint_into_model(model, state, allow_mismatched=args.allow_mismatched_checkpoint)
        print(f"RESUME_CHECKPOINT {args.resume_checkpoint}")
    n_params = sum(p.numel() for p in model.parameters())
    print(f"MODEL_PARAMS {n_params}")

    optimizer = AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay, betas=(0.9, 0.95))
    rec_weight = args.rec_weight
    z_weight = args.z_weight
    z_bin_weight = args.z_bin_weight
    z_candidate_weight = args.z_candidate_weight
    z_nll_weight = args.z_nll_weight
    if args.objective == "rec_only":
        rec_weight = rec_weight if rec_weight > 0 else 1.0
        z_weight = 0.0
        z_bin_weight = 0.0
        z_candidate_weight = 0.0
        z_nll_weight = 0.0
    elif args.objective == "z_only":
        rec_weight = 0.0

    loss_cfg = LossConfig(
        rec_weight=rec_weight,
        z_weight=z_weight,
        z_bin_weight=z_bin_weight,
        z_candidate_weight=z_candidate_weight,
        z_nll_weight=z_nll_weight,
        zwarn_weight=args.zwarn_weight,
        clean_z_only=args.clean_z_only,
        high_z_boost=args.high_z_boost,
        high_z_threshold=math.log1p(args.high_z_threshold),
    )
    best_score = math.inf
    global_step = 0
    micro_step = 0
    grad_accum_steps = max(1, int(args.grad_accum_steps))
    total_train_steps = args.max_steps if args.max_steps else int(math.ceil(len(train_loader) / grad_accum_steps) * args.epochs)
    model.train()
    optimizer.zero_grad(set_to_none=True)
    for epoch in range(args.epochs):
        pbar = tqdm(train_loader, desc=f"hybrid epoch {epoch}")
        for batch in pbar:
            micro_step += 1
            batch = move_to_device(batch, device)
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"):
                out = model(batch["x"], batch["valid"], batch["loglam"])
                loss, parts = redshift_total_loss(model, out, batch, loss_cfg)
            (loss / grad_accum_steps).backward()
            if micro_step % grad_accum_steps != 0:
                pbar.set_postfix(
                    loss=float(parts["loss"].detach().cpu()),
                    rec=float(parts["rec"].detach().cpu()),
                    huber=float(parts["z_huber"].detach().cpu()),
                    accum=f"{micro_step % grad_accum_steps}/{grad_accum_steps}",
                )
                continue
            next_step = global_step + 1
            lr_now = scheduled_lr(args.lr, args.min_lr, next_step, total_train_steps, int(args.warmup_steps))
            for group in optimizer.param_groups:
                group["lr"] = lr_now
            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
            optimizer.step()
            optimizer.zero_grad(set_to_none=True)
            global_step = next_step
            pbar.set_postfix(
                loss=float(parts["loss"].detach().cpu()),
                rec=float(parts["rec"].detach().cpu()),
                huber=float(parts["z_huber"].detach().cpu()),
                lr=lr_now,
                grad=float(grad_norm.detach().cpu()) if torch.is_tensor(grad_norm) else float(grad_norm),
            )
            if global_step == 1 or global_step % args.eval_every == 0:
                val_metrics = evaluate(
                    model,
                    val_loader,
                    loss_cfg,
                    device,
                    run_dir,
                    global_step,
                    prefix="val",
                    max_batches=None,
                    max_examples=args.eval_max_val,
                )
                print("VAL", global_step, json.dumps(val_metrics, sort_keys=True))
                ood_metrics = None
                if ood_loader is not None:
                    ood_metrics = evaluate(
                        model,
                        ood_loader,
                        loss_cfg,
                        device,
                        run_dir,
                        global_step,
                        prefix="ood",
                        max_batches=None,
                        max_examples=args.eval_max_ood,
                    )
                    print("OOD", global_step, json.dumps(ood_metrics, sort_keys=True))
                score = checkpoint_score(args.checkpoint_score, val_metrics, ood_metrics, z_alpha=args.score_z_alpha)
                if score < best_score:
                    best_score = score
                    best_metrics = {"step": global_step, "score": best_score, **val_metrics}
                    if ood_metrics is not None:
                        best_metrics.update(ood_metrics)
                    torch.save(
                        {"model": model.state_dict(), "args": vars(args), "step": global_step, "score": best_score, "metrics": best_metrics},
                        run_dir / "best.pt",
                    )
                    (run_dir / "best_metrics.json").write_text(json.dumps(best_metrics, indent=2), encoding="utf-8")
            if args.max_steps and global_step >= args.max_steps:
                break
        if args.max_steps and global_step >= args.max_steps:
            break

    final_metrics = evaluate(
        model,
        val_loader,
        loss_cfg,
        device,
        run_dir,
        global_step,
        prefix="val",
        max_batches=None,
        max_examples=args.eval_max_val,
    )
    if ood_loader is not None:
        final_metrics.update(
            evaluate(
                model,
                ood_loader,
                loss_cfg,
                device,
                run_dir,
                global_step,
                prefix="ood",
                max_batches=None,
                max_examples=args.eval_max_ood,
            )
        )
    torch.save({"model": model.state_dict(), "args": vars(args), "step": global_step, "metrics": final_metrics}, run_dir / "last.pt")
    (run_dir / "final_metrics.json").write_text(json.dumps(final_metrics, indent=2), encoding="utf-8")
    print("FINAL", json.dumps(final_metrics, sort_keys=True))
    print("RUN_DIR", run_dir)


if __name__ == "__main__":
    main()