"""Custom training loop for rho-guided SFT — MLX backend.

MLX equivalent of trainer.py. Runs on Apple Silicon unified memory,
avoiding the MPS NaN gradient bugs that force the PyTorch version to CPU.

Which backend should I use?
  ┌─────────────────────────┬──────────────────────────────────────┐
  │ Use MLX (this file)     │ Use PyTorch (trainer.py)             │
  ├─────────────────────────┼──────────────────────────────────────┤
  │ Apple Silicon Mac       │ Linux / Windows / CUDA GPUs          │
  │ M1/M2/M3/M4 chips      │ Non-Apple hardware                   │
  │ ~10× faster than CPU PT │ Supports MPS, CUDA, CPU              │
  │ Unified memory (no OOM) │ Full audit() integration (PyTorch)   │
  │ pip install mlx mlx-lm  │ pip install peft transformers        │
  └─────────────────────────┴──────────────────────────────────────┘

  Both backends produce the same training signal (CE + contrastive loss)
  with the same LoRA targets (Q, K, O projections). Results should be
  statistically equivalent within normal training variance.

Key differences from PyTorch version:
  - Uses mlx-lm's LoRA infrastructure instead of peft
  - Gradients via nn.value_and_grad() instead of .backward()
  - No device management (MLX uses unified memory)
  - Manual gradient accumulation via tree_map
  - SFT data passed as raw text strings, tokenized on-the-fly

Usage:
    from rho_eval.alignment.mlx_trainer import mlx_rho_guided_sft
    from rho_eval.alignment.dataset import BehavioralContrastDataset

    # model, tokenizer = mlx_lm.load("Qwen/Qwen2.5-7B-Instruct")
    sft_texts = ["### Instruction:\\n...\\n### Response:\\n...", ...]
    contrast_data = BehavioralContrastDataset()
    result = mlx_rho_guided_sft(model, tokenizer, sft_texts, contrast_data)
"""

from __future__ import annotations

import gc
import math
import random
import time
from typing import Callable, Optional

import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as opt
from mlx.utils import tree_flatten, tree_map, tree_unflatten

from .mlx_losses import mlx_rho_auxiliary_loss


# ── SVD Compression (MLX) ──────────────────────────────────────────────

def mlx_svd_compress(model, keep_ratio: float = 0.7) -> int:
    """SVD compress Q/K/O attention projections in an MLX model.

    Converts each weight to float32 numpy for SVD (MLX bfloat16 can't
    convert directly to numpy), then writes back the low-rank approximation.

    Args:
        model: MLX causal LM model (from mlx_lm.load()).
        keep_ratio: Fraction of singular values to retain (0.7 = 70%).

    Returns:
        Number of matrices compressed.
    """
    import numpy as np

    compressed = 0
    safe_projections = ["q_proj", "k_proj", "o_proj"]

    for i, layer in enumerate(model.model.layers):
        attn = layer.self_attn
        for proj_name in safe_projections:
            if not hasattr(attn, proj_name):
                continue

            proj = getattr(attn, proj_name)
            W_mx = proj.weight
            # Convert to float32 numpy for SVD (MLX bfloat16 can't go to numpy)
            W = np.array(W_mx.astype(mx.float32))

            if len(W.shape) != 2 or min(W.shape) <= 10:
                continue

            rank = max(1, int(min(W.shape) * keep_ratio))

            try:
                U, S, Vh = np.linalg.svd(W, full_matrices=False)
                W_approx = (U[:, :rank] * S[:rank]) @ Vh[:rank, :]
                proj.weight = mx.array(W_approx).astype(W_mx.dtype)
                compressed += 1
            except Exception as e:
                print(f"    SVD failed for layer {i} {proj_name}: {e}",
                      flush=True)
                continue

    mx.eval(model.parameters())
    return compressed


# ── Batching ────────────────────────────────────────────────────────────

def _tokenize_and_pad(
    texts: list[str],
    tokenizer,
    max_length: int = 256,
) -> list[list[int]]:
    """Tokenize texts, truncating to max_length."""
    encoded = []
    for text in texts:
        tokens = tokenizer.encode(text)
        if len(tokens) > max_length:
            tokens = tokens[:max_length]
        if len(tokens) >= 2:  # need at least 2 tokens for CE
            encoded.append(tokens)
    return encoded


def _iterate_batches(
    encoded: list[list[int]],
    batch_size: int,
    pad_token_id: int,
    shuffle: bool = True,
    rng: random.Random | None = None,
):
    """Yield (input_ids, lengths) batches from pre-tokenized texts.

    Pads each batch to the longest sequence in that batch.

    Args:
        encoded: List of token ID lists.
        batch_size: Batch size.
        pad_token_id: Token ID for padding.
        shuffle: Whether to shuffle.
        rng: Random instance for reproducibility.

    Yields:
        Tuple of (batch_ids: mx.array, lengths: mx.array)
        batch_ids shape: (batch_size, max_seq_len_in_batch)
        lengths shape: (batch_size,)
    """
    indices = list(range(len(encoded)))
    if shuffle and rng:
        rng.shuffle(indices)
    elif shuffle:
        random.shuffle(indices)

    for start in range(0, len(indices), batch_size):
        batch_indices = indices[start:start + batch_size]
        batch_tokens = [encoded[i] for i in batch_indices]
        lengths = [len(t) for t in batch_tokens]
        max_len = max(lengths)

        # Pad to max length in this batch
        padded = []
        for tokens in batch_tokens:
            padded.append(tokens + [pad_token_id] * (max_len - len(tokens)))

        batch_ids = mx.array(padded)       # (B, max_len)
        batch_lengths = mx.array(lengths)   # (B,)

        yield batch_ids, batch_lengths


# ── LoRA Utilities ──────────────────────────────────────────────────────

def _apply_lora(model, lora_rank: int, lora_alpha: int, dropout: float = 0.05):
    """Apply LoRA adapters to Q, K, O projections (not V — CF90 safety).

    Uses mlx-lm's built-in linear_to_lora_layers infrastructure.

    Returns:
        Tuple of (trainable_params, total_params).
    """
    from mlx_lm.tuner.utils import linear_to_lora_layers

    # Target Q, K, O attention projections (not V — per CF90 safety rules)
    lora_config = {
        "rank": lora_rank,
        "scale": lora_alpha / lora_rank,  # mlx-lm uses scale = alpha/rank
        "dropout": dropout,
        "keys": ["self_attn.q_proj", "self_attn.k_proj", "self_attn.o_proj"],
    }

    # Apply to all transformer layers
    num_layers = len(model.model.layers) if hasattr(model, 'model') else len(model.layers)
    linear_to_lora_layers(model, num_layers, lora_config)

    # Freeze everything, then unfreeze only LoRA params
    model.freeze()

    # Unfreeze LoRA parameters (lora_a, lora_b)
    from mlx_lm.tuner.lora import LoRALinear
    for name, module in model.named_modules():
        if isinstance(module, LoRALinear):
            module.unfreeze(keys=["lora_a", "lora_b"])

    # Count parameters
    total_params = sum(p.size for _, p in tree_flatten(model.parameters()))
    trainable_params = sum(
        p.size for _, p in tree_flatten(model.trainable_parameters())
    )

    return trainable_params, total_params


def _fuse_lora(model):
    """Fuse LoRA weights into base model (merge and unload equivalent).

    After fusing, LoRALinear layers become regular nn.Linear layers.
    """
    from mlx_lm.tuner.lora import LoRALinear

    fused_layers = []
    for name, module in model.named_modules():
        if isinstance(module, LoRALinear):
            fused_layers.append((name, module.fuse()))

    if fused_layers:
        model.update_modules(tree_unflatten(fused_layers))


def _flatten_grad(grad_tree) -> mx.array:
    """Flatten a gradient tree (nested dicts) into a single 1-D vector.

    Used for computing gradient norms and cosine similarities between
    gradient components (CE vs rho vs gamma).
    """
    leaves = [v.reshape(-1) for _, v in tree_flatten(grad_tree)]
    if not leaves:
        return mx.zeros((1,))
    return mx.concatenate(leaves)


def _save_lora_checkpoint(model, step: int, checkpoint_dir: str):
    """Save LoRA adapter weights only (not the full model).

    Saves ~16 MB per checkpoint for rank-8 on 7B, vs ~15 GB for full model.
    Only saves lora_a and lora_b parameters from LoRALinear modules.
    """
    import numpy as np
    from pathlib import Path
    from mlx_lm.tuner.lora import LoRALinear

    save_dir = Path(checkpoint_dir)
    save_dir.mkdir(parents=True, exist_ok=True)

    lora_weights = {}
    for name, module in model.named_modules():
        if isinstance(module, LoRALinear):
            lora_weights[f"{name}.lora_a"] = np.array(module.lora_a)
            lora_weights[f"{name}.lora_b"] = np.array(module.lora_b)

    path = save_dir / f"lora_step_{step:04d}.npz"
    np.savez(path, **lora_weights)
    n_params = sum(v.size for v in lora_weights.values())
    size_mb = sum(v.nbytes for v in lora_weights.values()) / 1e6
    print(f"  [checkpoint] step {step}: saved {len(lora_weights)} arrays, "
          f"{n_params:,} params, {size_mb:.1f} MB → {path}", flush=True)


# ── Training Loop ───────────────────────────────────────────────────────

def mlx_rho_guided_sft(
    model,
    tokenizer,
    sft_texts: list[str],
    contrast_dataset,
    rho_weight: float = 0.2,
    gamma_weight: float = 0.0,
    protection_dataset=None,
    epochs: int = 1,
    lr: float = 2e-4,
    batch_size: int = 2,
    gradient_accumulation_steps: int = 4,
    margin: float = 0.1,
    contrast_pairs_per_step: int = 4,
    protection_pairs_per_step: int = 4,
    lora_rank: int = 8,
    lora_alpha: int = 16,
    logging_steps: int = 50,
    max_steps: Optional[int] = None,
    warmup_ratio: float = 0.1,
    weight_decay: float = 0.01,
    max_length: int = 256,
    save_path: Optional[str] = None,
    checkpoint_steps: Optional[list[int]] = None,
    checkpoint_dir: Optional[str] = None,
    gradient_log_path: Optional[str] = None,
    gamma_schedule: Optional[Callable[[int], float]] = None,
    verbose: bool = True,
) -> dict:
    """Run rho-guided SFT with combined CE + contrastive + protection loss on MLX.

    Training loop:
      For each step:
        1. Get SFT batch → ce_loss (teacher-forced CE)
        2. Sample contrast pairs → rho_loss (contrastive confidence loss)
        3. Sample protection pairs → gamma_loss (protection confidence loss)
        4. total_loss = ce_loss + rho_weight * rho_loss + gamma_weight * gamma_loss
        5. Gradients via value_and_grad + accumulate
        6. Every N steps: optimizer.update(), mx.eval()

    The γ (gamma) protection loss prevents collateral damage on protected
    behaviors (e.g., bias) during target behavior SFT (e.g., sycophancy).
    It uses the same contrastive structure as the ρ loss but with different
    contrast data — bias pairs instead of sycophancy pairs.

    When rho_weight=0, this is equivalent to standard SFT (CE only).
    When gamma_weight=0, this is equivalent to the original rho-guided SFT.

    Args:
        model: mlx-lm model (loaded via mlx_lm.load()).
        tokenizer: mlx-lm TokenizerWrapper.
        sft_texts: List of raw text strings for SFT training.
        contrast_dataset: BehavioralContrastDataset for target behavior (sycophancy).
        rho_weight: Weight of the auxiliary rho loss (0 = CE only).
        gamma_weight: Weight of the γ protection loss (0 = no protection).
        protection_dataset: BehavioralContrastDataset for protected behaviors (bias).
            Required when gamma_weight > 0.
        epochs: Number of training epochs.
        lr: Learning rate (default: 2e-4, typical for LoRA).
        batch_size: Per-step batch size for SFT data.
        gradient_accumulation_steps: Accumulate gradients over N steps.
        margin: Contrastive margin (in CE loss units).
        contrast_pairs_per_step: Number of contrast pairs per step (for rho).
        protection_pairs_per_step: Number of protection pairs per step (for gamma).
        lora_rank: LoRA rank.
        lora_alpha: LoRA alpha scaling.
        logging_steps: Print progress every N steps.
        max_steps: Stop after this many steps (overrides epochs).
        warmup_ratio: Fraction of steps for LR warmup.
        weight_decay: AdamW weight decay.
        max_length: Max token length for encoding.
        gamma_schedule: Optional callable mapping step → γ value. When provided,
            overrides gamma_weight with a dynamic per-step value. Enables
            γ-annealing experiments (e.g., lambda s: 0.10 if s < 110 else 0.0).
        verbose: Print progress messages.

    Returns:
        Dict with training stats and the fused model:
          {ce_loss, rho_loss, gamma_loss, total_loss, steps, time,
           trainable_params, trainable_pct, lora_rank,
           merged_model, method}
    """
    t0 = time.time()

    # ── LoRA Setup ────────────────────────────────────────────────────
    if verbose:
        print(f"  [mlx-rho-sft] Applying LoRA (rank={lora_rank}, "
              f"alpha={lora_alpha})...")

    trainable_params, total_params = _apply_lora(
        model, lora_rank, lora_alpha,
    )

    if verbose:
        print(f"  [mlx-rho-sft] Trainable: {trainable_params/1e6:.2f}M / "
              f"{total_params/1e6:.1f}M ({trainable_params/total_params:.4%})")
        gamma_desc = f"gamma_schedule" if gamma_schedule else f"gamma_weight={gamma_weight}"
        print(f"  [mlx-rho-sft] rho_weight={rho_weight}, {gamma_desc}, "
              f"margin={margin}, contrast_pairs={contrast_pairs_per_step}")
        if (gamma_weight > 0 or gamma_schedule) and protection_dataset is not None:
            print(f"  [mlx-rho-sft] γ protection: {len(protection_dataset)} pairs, "
                  f"{protection_pairs_per_step}/step")

    # ── Tokenize SFT data ─────────────────────────────────────────────
    if verbose:
        print(f"  [mlx-rho-sft] Tokenizing {len(sft_texts)} SFT texts...")

    encoded = _tokenize_and_pad(sft_texts, tokenizer, max_length)

    if verbose:
        print(f"  [mlx-rho-sft] {len(encoded)} texts after filtering")

    # Determine pad token
    pad_token_id = getattr(tokenizer, 'pad_token_id', None)
    if pad_token_id is None:
        pad_token_id = getattr(tokenizer, 'eos_token_id', 0)

    # ── Optimizer + Scheduler ─────────────────────────────────────────
    n_batches = math.ceil(len(encoded) / batch_size)
    steps_per_epoch = math.ceil(n_batches / gradient_accumulation_steps)
    total_steps = max_steps or (steps_per_epoch * epochs)
    warmup_steps = int(total_steps * warmup_ratio)

    # Build schedule: linear warmup then linear decay
    if warmup_steps > 0:
        warmup_fn = opt.linear_schedule(
            init=1e-8, end=lr, steps=warmup_steps,
        )
        decay_fn = opt.linear_schedule(
            init=lr, end=0.0,
            steps=max(total_steps - warmup_steps, 1),
        )
        schedule = opt.join_schedules(
            [warmup_fn, decay_fn],
            [warmup_steps],
        )
    else:
        schedule = opt.linear_schedule(init=lr, end=0.0, steps=total_steps)

    optimizer = opt.AdamW(
        learning_rate=schedule,
        weight_decay=weight_decay,
    )

    if verbose:
        print(f"  [mlx-rho-sft] epochs={epochs}, lr={lr}, batch={batch_size}, "
              f"grad_accum={gradient_accumulation_steps}")
        print(f"  [mlx-rho-sft] total_steps={total_steps}, warmup={warmup_steps}")

    # ── Loss function for value_and_grad ──────────────────────────────
    def sft_loss_fn(model, batch_ids, batch_lengths):
        """CE loss on a padded SFT batch, with length masking."""
        inputs = batch_ids[:, :-1]    # (B, L-1)
        targets = batch_ids[:, 1:]    # (B, L-1)

        logits = model(inputs)        # (B, L-1, V)

        # Per-token CE
        ce_per_token = nn.losses.cross_entropy(logits, targets)  # (B, L-1)

        # Mask padding tokens: only count positions < length-1
        seq_positions = mx.arange(targets.shape[1])[None, :]  # (1, L-1)
        mask = seq_positions < (batch_lengths[:, None] - 1)    # (B, L-1)

        # Masked mean
        ce_sum = (ce_per_token * mask).sum()
        n_tokens = mask.sum()
        ce_loss = ce_sum / mx.maximum(n_tokens, mx.array(1.0))

        return ce_loss, n_tokens

    loss_and_grad_fn = nn.value_and_grad(model, sft_loss_fn)

    # ── Gradient Telemetry Setup ──────────────────────────────────────
    grad_log_file = None
    if gradient_log_path:
        from pathlib import Path
        Path(gradient_log_path).parent.mkdir(parents=True, exist_ok=True)
        grad_log_file = open(gradient_log_path, "w")
        grad_log_file.write(
            "step,gamma_value,ce_loss,rho_loss,gamma_loss,"
            "ce_norm,rho_norm,gamma_norm,"
            "cos_rho_gamma,cos_ce_rho,cos_ce_gamma\n"
        )
        if verbose:
            print(f"  [mlx-rho-sft] Gradient telemetry → {gradient_log_path}")

    # Validate checkpoint config
    checkpoint_steps_set = set(checkpoint_steps) if checkpoint_steps else set()
    if checkpoint_steps and not checkpoint_dir:
        raise ValueError("checkpoint_dir required when checkpoint_steps is set")

    # ── Save step-0 LoRA checkpoint (random init, before training) ──
    if 0 in checkpoint_steps_set:
        _save_lora_checkpoint(model, 0, checkpoint_dir)

    # ── Training Loop ─────────────────────────────────────────────────
    model.train()
    rng = random.Random(42)

    running_ce = 0.0
    running_rho = 0.0
    running_gamma = 0.0
    running_total = 0.0
    global_step = 0
    log_count = 0
    accumulated_grad = None
    micro_step = 0
    done = False

    for epoch in range(epochs):
        for batch_ids, batch_lengths in _iterate_batches(
            encoded, batch_size, pad_token_id, shuffle=True, rng=rng,
        ):
            # ── CE loss on SFT batch ───────────────────────────────
            (ce_loss_val, n_tokens), grad = loss_and_grad_fn(
                model, batch_ids, batch_lengths,
            )

            # Scale for gradient accumulation
            scaled_grad = tree_map(
                lambda g: g / gradient_accumulation_steps, grad,
            )

            # ── Rho auxiliary loss on contrast pairs ──────────────
            if rho_weight > 0 and len(contrast_dataset) > 0:
                pairs = contrast_dataset.sample(contrast_pairs_per_step, rng)

                # Compute rho loss and its gradient separately
                def rho_loss_fn(model):
                    return mlx_rho_auxiliary_loss(
                        model, tokenizer, pairs,
                        margin=margin, max_length=max_length,
                    )

                rho_grad_fn = nn.value_and_grad(model, rho_loss_fn)
                rho_loss_val, rho_grad = rho_grad_fn(model)

                # Scale rho gradient by weight and accumulation
                rho_scaled = tree_map(
                    lambda g: g * (rho_weight / gradient_accumulation_steps),
                    rho_grad,
                )

                # Add rho gradient to CE gradient
                scaled_grad = tree_map(
                    lambda a, b: a + b, scaled_grad, rho_scaled,
                )
            else:
                rho_loss_val = mx.array(0.0)

            # ── Gamma protection loss on protected pairs ──────────
            current_gamma = gamma_schedule(global_step) if gamma_schedule else gamma_weight
            if current_gamma > 0 and protection_dataset is not None and len(protection_dataset) > 0:
                prot_pairs = protection_dataset.sample(protection_pairs_per_step, rng)

                # Compute gamma loss and its gradient separately
                def gamma_loss_fn(model):
                    return mlx_rho_auxiliary_loss(
                        model, tokenizer, prot_pairs,
                        margin=margin, max_length=max_length,
                    )

                gamma_grad_fn = nn.value_and_grad(model, gamma_loss_fn)
                gamma_loss_val, gamma_grad = gamma_grad_fn(model)

                # Scale gamma gradient by weight and accumulation
                gamma_scaled = tree_map(
                    lambda g: g * (current_gamma / gradient_accumulation_steps),
                    gamma_grad,
                )

                # Add gamma gradient to combined gradient
                scaled_grad = tree_map(
                    lambda a, b: a + b, scaled_grad, gamma_scaled,
                )
            else:
                gamma_loss_val = mx.array(0.0)

            # ── Gradient Telemetry ──────────────────────────────────
            if grad_log_file is not None:
                # Flatten raw (pre-accumulation) gradient components
                ce_flat = _flatten_grad(grad)  # raw CE gradient (unscaled)
                ce_n = mx.sqrt((ce_flat * ce_flat).sum()).item()

                if rho_weight > 0 and 'rho_grad' in dir():
                    rho_flat = _flatten_grad(rho_grad)  # raw rho gradient
                    rho_n = mx.sqrt((rho_flat * rho_flat).sum()).item()
                else:
                    rho_flat = mx.zeros_like(ce_flat)
                    rho_n = 0.0

                if current_gamma > 0 and 'gamma_grad' in dir():
                    gamma_flat = _flatten_grad(gamma_grad)  # raw gamma gradient
                    gamma_n = mx.sqrt((gamma_flat * gamma_flat).sum()).item()
                else:
                    gamma_flat = mx.zeros_like(ce_flat)
                    gamma_n = 0.0

                # Cosine similarities between gradient components
                eps = 1e-12
                cos_rg = float((rho_flat * gamma_flat).sum().item()) / (rho_n * gamma_n + eps) if (rho_n > eps and gamma_n > eps) else 0.0
                cos_cr = float((ce_flat * rho_flat).sum().item()) / (ce_n * rho_n + eps) if (ce_n > eps and rho_n > eps) else 0.0
                cos_cg = float((ce_flat * gamma_flat).sum().item()) / (ce_n * gamma_n + eps) if (ce_n > eps and gamma_n > eps) else 0.0

                grad_log_file.write(
                    f"{global_step},{current_gamma:.6f},{ce_loss_val.item():.6f},{rho_loss_val.item():.6f},"
                    f"{gamma_loss_val.item():.6f},{ce_n:.6f},{rho_n:.6f},{gamma_n:.6f},"
                    f"{cos_rg:.6f},{cos_cr:.6f},{cos_cg:.6f}\n"
                )
                grad_log_file.flush()

            # ── Accumulate ────────────────────────────────────────
            if accumulated_grad is None:
                accumulated_grad = scaled_grad
            else:
                accumulated_grad = tree_map(
                    lambda a, b: a + b, accumulated_grad, scaled_grad,
                )

            micro_step += 1

            # Track losses (evaluate lazily)
            ce_val = ce_loss_val.item()
            rho_val = rho_loss_val.item()
            gamma_val = gamma_loss_val.item()
            running_ce += ce_val
            running_rho += rho_val
            running_gamma += gamma_val
            running_total += ce_val + rho_weight * rho_val + gamma_weight * gamma_val

            # ── Optimizer step every N micro-steps ────────────────
            if micro_step % gradient_accumulation_steps == 0:
                # Clip gradients
                accumulated_grad, _ = opt.clip_grad_norm(
                    accumulated_grad, max_norm=1.0,
                )

                # Update model
                optimizer.update(model, accumulated_grad)

                # Evaluate (materialize lazy computation)
                mx.eval(model.parameters(), optimizer.state)

                accumulated_grad = None
                global_step += 1
                log_count += 1

                # ── LoRA Checkpoint ────────────────────────────────
                if global_step in checkpoint_steps_set:
                    _save_lora_checkpoint(model, global_step, checkpoint_dir)

                if verbose and global_step % logging_steps == 0:
                    avg_ce = running_ce / log_count
                    avg_rho = running_rho / log_count
                    avg_gamma = running_gamma / log_count
                    avg_total = running_total / log_count
                    elapsed = time.time() - t0
                    gamma_str = f", gamma={avg_gamma:.4f}" if gamma_weight > 0 else ""
                    print(f"  [mlx-rho-sft] step {global_step}/{total_steps}: "
                          f"ce={avg_ce:.4f}, rho={avg_rho:.4f}{gamma_str}, "
                          f"total={avg_total:.4f}, "
                          f"elapsed={elapsed:.0f}s", flush=True)
                    running_ce = 0.0
                    running_rho = 0.0
                    running_gamma = 0.0
                    running_total = 0.0
                    log_count = 0

                if max_steps and global_step >= max_steps:
                    done = True
                    break

        if done:
            break

    # Final stats
    elapsed = time.time() - t0

    # Close gradient telemetry log
    if grad_log_file is not None:
        grad_log_file.close()
        if verbose:
            print(f"  [mlx-rho-sft] Gradient telemetry: {global_step} rows written")

    if verbose:
        print(f"  [mlx-rho-sft] Training done: {global_step} steps, "
              f"{elapsed:.1f}s")

    # ── Fuse LoRA + Cleanup ───────────────────────────────────────────
    if verbose:
        print(f"  [mlx-rho-sft] Fusing LoRA weights into base model...")

    _fuse_lora(model)

    model.eval()

    # ── Save merged model ────────────────────────────────────────────
    if save_path is not None:
        from pathlib import Path
        from mlx_lm.utils import save_model

        save_dir = Path(save_path)
        save_dir.mkdir(parents=True, exist_ok=True)

        if verbose:
            print(f"  [mlx-rho-sft] Saving merged model to {save_dir}...")

        # Save sharded model weights (mlx_lm format, compatible with mlx_lm.load)
        save_model(save_dir, model)

        # Save tokenizer
        tokenizer.save_pretrained(str(save_dir))

        # Copy config from original model if available
        if hasattr(model, "config"):
            import json as _json
            config_path = save_dir / "config.json"
            if not config_path.exists():
                if hasattr(model.config, "to_dict"):
                    config_path.write_text(_json.dumps(model.config.to_dict(), indent=2))
                elif isinstance(model.config, dict):
                    config_path.write_text(_json.dumps(model.config, indent=2))

        if verbose:
            n_files = len(list(save_dir.iterdir()))
            total_size = sum(f.stat().st_size for f in save_dir.iterdir() if f.is_file())
            print(f"  [mlx-rho-sft] Saved: {n_files} files, "
                  f"{total_size / 1e9:.1f} GB")

    if verbose:
        print(f"  [mlx-rho-sft] Done: {global_step} steps, {elapsed:.1f}s")

    return {
        "ce_loss": running_ce / max(log_count, 1) if log_count > 0 else 0.0,
        "rho_loss": running_rho / max(log_count, 1) if log_count > 0 else 0.0,
        "gamma_loss": running_gamma / max(log_count, 1) if log_count > 0 else 0.0,
        "total_loss": running_total / max(log_count, 1) if log_count > 0 else 0.0,
        "steps": global_step,
        "time": elapsed,
        "trainable_params": trainable_params,
        "trainable_pct": trainable_params / total_params if total_params > 0 else 0.0,
        "lora_rank": lora_rank,
        "rho_weight": rho_weight,
        "gamma_weight": gamma_weight,
        "margin": margin,
        "method": "mlx_rho_guided_sft",
        "merged_model": model,
    }
