"""Trainium distillation training loop for the shallow Voxtral draft.

Design:
  - Data: precomputed (audio_embeds, teacher_token_ids) tuples on CPU disk.
    Load per batch, splice audio_embeds into a padded [B, max_seq_len, hidden]
    inputs_embeds tensor. Attention mask + labels mask handle variable lengths.
  - Model: draft = fresh copy of Voxtral's LLM decoder with N of 30 layers
    kept, all other weights (embed_tokens, norm, lm_head) copied from target.
  - Loss: cross-entropy on teacher_token_ids at the position immediately
    after the audio+prompt prefix, teacher-forcing the rest.
  - Optimizer: AdamW, LR 1e-4, warmup linear + cosine decay.
  - XLA: mark_step after each optimizer.step(). Static shape via padding.

Runs on trn2 (torch_xla). Single-device (LNC=2 -> 4 logical cores, TP=1 here).

Usage:
    python train_distill.py \
        --target mistralai/Voxtral-Mini-3B-2507 \
        --data-dir /mnt/data/pseudo_labels \
        --keep-layers 0,10,20,29 \
        --output-dir /mnt/drafts/voxtral-mini-3B-draft-trained-v1 \
        --max-seq-len 512 \
        --batch-size 2 \
        --lr 1e-4 \
        --warmup-steps 20 \
        --total-steps 200 \
        --save-every 50
"""

from __future__ import annotations

import argparse
import copy
import gc
import json
import math
import random
import time
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import VoxtralForConditionalGeneration, AutoProcessor

import torch_xla
import torch_xla.core.xla_model as xm


# ---- Draft construction --------------------------------------------------------

def prune_decoder_layers(model, keep_indices: list[int]):
    """Prune the LLM decoder to keep only the specified layer indices.

    Modifies model in place. Also updates config.text_config.num_hidden_layers.
    """
    decoder = model.model.language_model
    layers = decoder.layers
    n_full = len(layers)
    assert max(keep_indices) < n_full, f"keep={keep_indices} exceeds n_full={n_full}"

    kept = nn.ModuleList([layers[i] for i in keep_indices])
    for new_idx, layer in enumerate(kept):
        if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "layer_idx"):
            layer.self_attn.layer_idx = new_idx
    decoder.layers = kept
    model.config.text_config.num_hidden_layers = len(keep_indices)
    if hasattr(model.config, "num_hidden_layers"):
        model.config.num_hidden_layers = len(keep_indices)
    return model


# ---- Dataset ------------------------------------------------------------------

class PseudoLabelDataset(torch.utils.data.Dataset):
    """Loads .pt files produced by precompute_pseudo_labels.py."""

    def __init__(self, data_dir: Path):
        self.data_dir = Path(data_dir)
        self.files = sorted([p for p in self.data_dir.glob("*.pt") if p.name != "_metadata.json"])
        assert self.files, f"No .pt files in {data_dir}"
        # Also load metadata
        meta_path = self.data_dir / "_metadata.json"
        with open(meta_path) as f:
            self.metadata = json.load(f)

    def __len__(self):
        return len(self.files)

    def __getitem__(self, i):
        return torch.load(self.files[i], map_location="cpu", weights_only=False)


def collate_batch(batch, max_seq_len: int, hidden_size: int, embed_tokens_module: nn.Module, audio_token_id: int, pad_token_id: int, dtype: torch.dtype):
    """Build padded [B, max_seq_len, hidden_size] inputs_embeds + labels for CE.

    For each sample:
      - prefix = [audio + prompt tokens], length = prefix_len (383 for standard Voxtral audio+transcribe prompt)
      - target = teacher_token_ids, length = variable
      - We construct: inputs_embeds[0:prefix_len] = <audio-spliced text embeds>
                      inputs_embeds[prefix_len:prefix_len+n_target] = <embed(teacher_token_ids[:-1] prepended by BOS-like)>
      - Actually easier: use teacher-forcing on prefix + full target. Model sees
        prefix + target[:-1], predicts target at positions prefix_len-1 to prefix_len+n_target-2.
        Label positions match. Positions before prefix_len-1 are label=-100 (ignored).

    Actually cleanest: form (prefix + target) as one sequence. Predict target
    tokens at positions prefix_len ... prefix_len + n_target - 1. The model's
    logits at position (prefix_len - 1) predicts target[0], and generally the
    logits at position (prefix_len + k - 1) predicts target[k].

    So:
      - inputs_embeds = [prefix_embeds, embed(target_tokens[:-1])] (length prefix_len + n_target - 1)
      - labels        = [-100 for prefix, target_tokens[0:n_target]] shifted appropriately

    But since we're using teacher forcing and the model outputs logits at every
    position (including prefix), a simpler way is:
      - full_seq = [prefix_tokens, target_tokens] concat, length = prefix_len + n_target
      - inputs_embeds[t] = embed(full_seq[t]) with audio-splice at prefix positions
      - labels[t] = full_seq[t+1] for t in [prefix_len-1, ..., prefix_len+n_target-2]; -100 elsewhere
      - We predict at position t using logits[t]; loss uses shift-by-one convention.

    Simpler: use HF's -100 ignore_index. labels[i] = full_seq[i+1] for target
    positions, -100 otherwise. Model outputs logits, we take logits[:-1] and
    labels[1:] and compute CE. transformers Llama-style models do this
    internally when labels are passed. Here we do it manually to be explicit.

    Padding: pad all sequences to `max_seq_len` with pad_token_id in inputs
    and -100 in labels.
    """
    B = len(batch)
    device = "cpu"

    # Build per-sample sequences
    inputs_embeds_list = []
    labels_list = []
    attention_mask_list = []

    for sample in batch:
        input_ids = sample["input_ids"]           # [1, prefix_len]
        audio_embeds = sample["audio_embeds"]     # [n_audio, hidden]
        teacher_ids = sample["teacher_token_ids"]  # [n_target]

        prefix_len = input_ids.shape[1]
        n_target = teacher_ids.shape[0]
        total_len = prefix_len + n_target

        if total_len > max_seq_len:
            # Truncate target
            n_target = max_seq_len - prefix_len
            teacher_ids = teacher_ids[:n_target]
            total_len = max_seq_len

        # 1. Build text embeddings for prefix
        with torch.no_grad():
            prefix_text_embeds = embed_tokens_module(input_ids).to(dtype)  # [1, prefix_len, hidden]
        # 2. Splice audio embeds
        audio_mask = (input_ids == audio_token_id).unsqueeze(-1)  # [1, prefix_len, 1]
        prefix_text_embeds = prefix_text_embeds.masked_scatter(audio_mask, audio_embeds)
        prefix_embeds = prefix_text_embeds.squeeze(0)  # [prefix_len, hidden]

        # 3. Build target embeddings
        with torch.no_grad():
            target_embeds = embed_tokens_module(teacher_ids.unsqueeze(0)).to(dtype).squeeze(0)  # [n_target, hidden]

        # 4. Concat: [prefix, target]. inputs_embeds is [prefix + target]
        full_embeds = torch.cat([prefix_embeds, target_embeds], dim=0)  # [total_len, hidden]

        # 5. Build labels for CE. labels[t] = the token the model should predict AT position t.
        # For an autoregressive model, logits at position t predict the token at position t+1.
        # So labels[t] = full_seq_token[t+1]. Convert: we want the logits at positions
        # (prefix_len - 1 ... prefix_len + n_target - 2) to predict tokens
        # (teacher_ids[0] ... teacher_ids[n_target - 1]).
        # Layout labels aligned to sequence positions [0 ... total_len - 1] with -100 for ignore.
        # labels[t] should be teacher_ids[t - (prefix_len - 1)] when t in [prefix_len-1, prefix_len+n_target-2].
        labels = torch.full((total_len,), -100, dtype=torch.long)
        for k in range(n_target):
            pos = prefix_len - 1 + k
            if 0 <= pos < total_len:
                labels[pos] = teacher_ids[k]

        # Attention mask (all-ones since no padding at this stage; padding added below)
        attn = torch.ones(total_len, dtype=torch.long)

        # 6. Pad to max_seq_len on the right
        pad_len = max_seq_len - total_len
        if pad_len > 0:
            pad_embeds = torch.zeros(pad_len, prefix_embeds.shape[-1], dtype=dtype)
            full_embeds = torch.cat([full_embeds, pad_embeds], dim=0)
            labels = torch.cat([labels, torch.full((pad_len,), -100, dtype=torch.long)], dim=0)
            attn = torch.cat([attn, torch.zeros(pad_len, dtype=torch.long)], dim=0)

        inputs_embeds_list.append(full_embeds)
        labels_list.append(labels)
        attention_mask_list.append(attn)

    inputs_embeds = torch.stack(inputs_embeds_list, dim=0)  # [B, max_seq_len, hidden]
    labels = torch.stack(labels_list, dim=0)  # [B, max_seq_len]
    attention_mask = torch.stack(attention_mask_list, dim=0)  # [B, max_seq_len]
    return {
        "inputs_embeds": inputs_embeds,
        "labels": labels,
        "attention_mask": attention_mask,
    }


# ---- Training loop ------------------------------------------------------------

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--target", default="mistralai/Voxtral-Mini-3B-2507")
    p.add_argument("--data-dir", required=True, type=Path)
    p.add_argument("--output-dir", required=True, type=Path)
    p.add_argument("--keep-layers", default="0,10,20,29",
                   help="Comma-separated 0-indexed decoder layers to keep in the draft")
    p.add_argument("--max-seq-len", type=int, default=512)
    p.add_argument("--batch-size", type=int, default=2)
    p.add_argument("--lr", type=float, default=1e-4)
    p.add_argument("--warmup-steps", type=int, default=20)
    p.add_argument("--total-steps", type=int, default=200)
    p.add_argument("--save-every", type=int, default=50)
    p.add_argument("--log-every", type=int, default=1)
    p.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float32"])
    p.add_argument("--seed", type=int, default=42)
    return p.parse_args()


def build_lr_schedule(warmup_steps: int, total_steps: int):
    def fn(step):
        if step < warmup_steps:
            return step / max(1, warmup_steps)
        # cosine decay to 0.1
        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
        progress = min(1.0, progress)
        return 0.1 + 0.9 * 0.5 * (1.0 + math.cos(math.pi * progress))
    return fn


def main() -> int:
    args = parse_args()
    torch.manual_seed(args.seed)
    random.seed(args.seed)

    dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float32
    keep_layers = [int(x) for x in args.keep_layers.split(",")]

    print(f"[train] Loading target on CPU: {args.target}")
    target_full = VoxtralForConditionalGeneration.from_pretrained(
        args.target, torch_dtype=dtype, low_cpu_mem_usage=True,
    ).eval()
    processor = AutoProcessor.from_pretrained(args.target)
    hidden_size = target_full.config.text_config.hidden_size
    audio_token_id = target_full.config.audio_token_id
    pad_token_id = (processor.tokenizer.pad_token_id
                    if processor.tokenizer.pad_token_id is not None
                    else processor.tokenizer.eos_token_id)
    print(f"[train] hidden_size={hidden_size}, audio_token_id={audio_token_id}, pad_token_id={pad_token_id}")

    # Keep only the embed_tokens module on CPU for collate; we don't need the full target during training
    embed_tokens_module_cpu = target_full.model.language_model.embed_tokens
    embed_tokens_module_cpu.eval()
    for p in embed_tokens_module_cpu.parameters():
        p.requires_grad_(False)

    # ---- Build the draft (a copy of the target with pruned layers)
    print(f"[train] Building draft with keep_layers={keep_layers}")
    draft = VoxtralForConditionalGeneration.from_pretrained(
        args.target, torch_dtype=dtype, low_cpu_mem_usage=True,
    )
    prune_decoder_layers(draft, keep_layers)
    print(f"[train] Draft has {draft.config.text_config.num_hidden_layers} decoder layers "
          f"(from 30 in target)")

    # ---- Free the target (we don't need it beyond embed_tokens which we already extracted)
    # But we need embed_tokens for the collate function. We can also use the DRAFT's embed_tokens
    # since it was copied from target. Simpler: keep a lightweight reference to the DRAFT's
    # embed_tokens (still on CPU before we move draft to XLA).
    del target_full
    gc.collect()

    # ---- Move draft to XLA device
    xla = xm.xla_device()
    print(f"[train] XLA device: {xla}")

    # We train the shallow decoder ONLY, but the embed_tokens + norm + lm_head
    # travel with the model. Freeze them so we only optimize the decoder layers.
    for name, p in draft.named_parameters():
        # Trainable: everything under model.language_model.layers.
        # Frozen: embed_tokens, norm, lm_head, and audio encoder + projector (though the
        # audio path won't run on XLA since we splice pre-computed audio_embeds).
        # Freeze audio path since it never runs in training.
        if "language_model.layers." in name:
            p.requires_grad_(True)
        else:
            p.requires_grad_(False)

    # Only need the LLM decoder + lm_head on XLA (skip audio encoder entirely to save HBM)
    lang = draft.model.language_model.to(xla)
    lm_head = draft.lm_head.to(xla)
    xm.mark_step()
    print(f"[train] LLM on XLA")

    trainable_params = [p for p in lang.parameters() if p.requires_grad]
    # embed_tokens + norm are inside lang -- keep them frozen. Only .layers get gradients.
    for n, p in lang.named_parameters():
        if not n.startswith("layers."):
            p.requires_grad_(False)
    trainable_params = [p for p in lang.parameters() if p.requires_grad]
    n_trainable = sum(p.numel() for p in trainable_params)
    print(f"[train] Trainable params: {n_trainable:,} "
          f"({n_trainable / 1e6:.1f}M)")

    # ---- Dataset + dataloader
    dataset = PseudoLabelDataset(args.data_dir)
    print(f"[train] Dataset: {len(dataset)} clips from {args.data_dir}")

    def _collate(batch):
        return collate_batch(
            batch, args.max_seq_len, hidden_size,
            embed_tokens_module_cpu, audio_token_id, pad_token_id, dtype,
        )

    # For small datasets, use a simple round-robin (no shuffle for reproducibility of prototype)
    loader = torch.utils.data.DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=0,   # inline; small dataset
        collate_fn=_collate,
        drop_last=True,
    )

    # ---- Optimizer + LR schedule
    optim = torch.optim.AdamW(trainable_params, lr=args.lr, betas=(0.9, 0.95), weight_decay=0.01)
    lr_fn = build_lr_schedule(args.warmup_steps, args.total_steps)

    # ---- Training loop
    print(f"[train] Starting training. total_steps={args.total_steps}")
    step = 0
    loader_iter = iter(loader)
    losses = []
    t_start = time.perf_counter()

    while step < args.total_steps:
        try:
            batch = next(loader_iter)
        except StopIteration:
            loader_iter = iter(loader)
            batch = next(loader_iter)

        # LR update
        lr_scale = lr_fn(step)
        for pg in optim.param_groups:
            pg["lr"] = args.lr * lr_scale

        # Move batch to XLA
        inputs_embeds = batch["inputs_embeds"].to(xla)
        labels = batch["labels"].to(xla)
        attn = batch["attention_mask"].to(xla)

        # Forward
        out = lang(
            inputs_embeds=inputs_embeds,
            attention_mask=attn,
            use_cache=False,
        )
        hidden = out.last_hidden_state  # [B, L, hidden]
        logits = lm_head(hidden)  # [B, L, vocab]

        # CE loss
        B, L, V = logits.shape
        loss = F.cross_entropy(
            logits.view(-1, V).float(),
            labels.view(-1),
            ignore_index=-100,
        )

        # Backward
        optim.zero_grad()
        loss.backward()
        # Gradient clipping (optional but stabilizes)
        torch.nn.utils.clip_grad_norm_(trainable_params, max_norm=1.0)
        optim.step()

        # XLA sync
        xm.mark_step()

        loss_val = loss.item()
        losses.append(loss_val)

        if step % args.log_every == 0:
            elapsed = time.perf_counter() - t_start
            print(f"[train] step {step:4d} | loss {loss_val:.4f} | lr {args.lr*lr_scale:.2e} | "
                  f"elapsed {elapsed:.1f}s | avg step time {elapsed/max(1,step+1)*1000:.0f}ms")

        if (step + 1) % args.save_every == 0 or step + 1 == args.total_steps:
            save_dir = args.output_dir / f"ckpt_step_{step+1:06d}"
            save_dir.mkdir(parents=True, exist_ok=True)
            # Move draft weights back to CPU for saving
            print(f"[train] Saving checkpoint to {save_dir}")
            # Sync XLA weights to CPU
            xm.mark_step()
            cpu_state = {k: v.detach().to("cpu") for k, v in draft.state_dict().items()}
            # Reload to a full VoxtralForConditionalGeneration and save
            draft_cpu = VoxtralForConditionalGeneration.from_pretrained(
                args.target, torch_dtype=dtype, low_cpu_mem_usage=True,
            )
            prune_decoder_layers(draft_cpu, keep_layers)
            # Overwrite trained params
            draft_cpu.load_state_dict(cpu_state, strict=False)
            draft_cpu.save_pretrained(save_dir, safe_serialization=True)
            processor.save_pretrained(save_dir)
            # Save training state
            torch.save({
                "step": step + 1,
                "losses": losses,
                "args": vars(args),
                "keep_layers": keep_layers,
            }, save_dir / "training_state.pt")
            del draft_cpu, cpu_state
            gc.collect()

        step += 1

    print(f"\n[train] Training complete. {step} steps in {time.perf_counter()-t_start:.1f}s")
    print(f"[train] Final loss: {losses[-1]:.4f}, initial loss: {losses[0]:.4f}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())