"""
Text Diffusion Model for EN→DE Machine Translation
Self-contained training script.
Architecture: Masked Discrete Diffusion with DiT backbone
Inspired by MDLM (kuleshov-group) + LLaDA conditional generation
Dataset: WMT14 EN-DE

Usage:
    pip install torch transformers datasets trackio sacrebleu sacremoses sentencepiece protobuf
    python train.py
"""

import os
import math
import typing
import time
import json
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from dataclasses import dataclass
from datasets import load_dataset
from transformers import AutoTokenizer, get_cosine_schedule_with_warmup
import trackio


# ═══════════════════════════════════════════════════════════════
# MODEL ARCHITECTURE
# ═══════════════════════════════════════════════════════════════

@dataclass
class DiffusionTranslatorConfig:
    vocab_size: int = 32128
    max_src_len: int = 128
    max_tgt_len: int = 128
    hidden_dim: int = 512
    n_heads: int = 8
    n_blocks: int = 8
    dropout: float = 0.1
    cond_dim: int = 128
    mask_token_id: int = 32100
    pad_token_id: int = 0


class Rotary(nn.Module):
    def __init__(self, dim, base=10_000):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)
        self.seq_len_cached = None
        self.cos_cached = None
        self.sin_cached = None

    def forward(self, x, seq_dim=1):
        seq_len = x.shape[seq_dim]
        if seq_len != self.seq_len_cached:
            self.seq_len_cached = seq_len
            t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
            self.cos_cached = emb.cos()
            self.sin_cached = emb.sin()
        return self.cos_cached, self.sin_cached


def rotate_half(x):
    x1, x2 = x[..., :x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(q, k, cos, sin):
    cos = cos[:q.shape[1], :]
    sin = sin[:q.shape[1], :]
    cos = cos.unsqueeze(0).unsqueeze(2)
    sin = sin.unsqueeze(0).unsqueeze(2)
    q = (q * cos) + (rotate_half(q) * sin)
    k = (k * cos) + (rotate_half(k) * sin)
    return q, k


class TimestepEmbedder(nn.Module):
    def __init__(self, hidden_size, frequency_embedding_size=256):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size, bias=True),
        )
        self.frequency_embedding_size = frequency_embedding_size

    @staticmethod
    def timestep_embedding(t, dim, max_period=10000):
        half = dim // 2
        freqs = torch.exp(
            -math.log(max_period) * torch.arange(0, half, dtype=torch.float32, device=t.device) / half
        )
        args = t[:, None].float() * freqs[None]
        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
        if dim % 2:
            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
        return embedding

    def forward(self, t):
        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
        return self.mlp(t_freq)


class LayerNorm(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.weight = nn.Parameter(torch.ones([dim]))
        self.dim = dim

    def forward(self, x):
        with torch.amp.autocast('cuda', enabled=False):
            x = F.layer_norm(x.float(), [self.dim])
        return x * self.weight[None, None, :]


class DiTBlock(nn.Module):
    """Diffusion Transformer block with adaptive layer norm (adaLN)."""

    def __init__(self, dim, n_heads, cond_dim, mlp_ratio=4, dropout=0.1):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = dim // n_heads

        self.norm1 = LayerNorm(dim)
        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.k_proj = nn.Linear(dim, dim, bias=False)
        self.v_proj = nn.Linear(dim, dim, bias=False)
        self.attn_out = nn.Linear(dim, dim, bias=False)
        self.dropout1 = nn.Dropout(dropout)

        self.norm2 = LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, mlp_ratio * dim, bias=True),
            nn.GELU(approximate='tanh'),
            nn.Linear(mlp_ratio * dim, dim, bias=True),
        )
        self.dropout2 = nn.Dropout(dropout)

        self.adaLN_modulation = nn.Linear(cond_dim, 6 * dim, bias=True)
        nn.init.zeros_(self.adaLN_modulation.weight)
        nn.init.zeros_(self.adaLN_modulation.bias)

    def forward(self, x, rotary_cos_sin, c, attention_mask=None):
        batch_size, seq_len, dim = x.shape

        mod = self.adaLN_modulation(c)[:, None, :].chunk(6, dim=2)
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod

        x_skip = x
        x_norm = self.norm1(x) * (1 + scale_msa) + shift_msa

        q = self.q_proj(x_norm).view(batch_size, seq_len, self.n_heads, self.head_dim)
        k = self.k_proj(x_norm).view(batch_size, seq_len, self.n_heads, self.head_dim)
        v = self.v_proj(x_norm).view(batch_size, seq_len, self.n_heads, self.head_dim)

        cos, sin = rotary_cos_sin
        q, k = apply_rotary_pos_emb(q, k, cos, sin)

        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # Bidirectional attention (no causal mask)
        attn_output = F.scaled_dot_product_attention(
            q, k, v, attn_mask=attention_mask,
            dropout_p=self.dropout1.p if self.training else 0.0
        )
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, dim)

        x = x_skip + gate_msa * self.dropout1(self.attn_out(attn_output))

        x_skip = x
        x_norm = self.norm2(x) * (1 + scale_mlp) + shift_mlp
        x = x_skip + gate_mlp * self.dropout2(self.mlp(x_norm))

        return x


class DiffusionTranslator(nn.Module):
    """
    Masked Discrete Diffusion model for EN→DE translation.
    Input: [source_tokens | target_tokens] where target tokens are partially masked
    Bidirectional transformer (DiT blocks with adaLN for timestep conditioning)
    """

    def __init__(self, config: DiffusionTranslatorConfig):
        super().__init__()
        self.config = config

        self.vocab_embed = nn.Embedding(config.vocab_size, config.hidden_dim)
        self.sigma_map = TimestepEmbedder(config.cond_dim)
        self.rotary_emb = Rotary(config.hidden_dim // config.n_heads)
        self.segment_embed = nn.Embedding(2, config.hidden_dim)

        self.blocks = nn.ModuleList([
            DiTBlock(config.hidden_dim, config.n_heads, config.cond_dim, dropout=config.dropout)
            for _ in range(config.n_blocks)
        ])

        self.final_norm = LayerNorm(config.hidden_dim)
        self.final_adaLN = nn.Linear(config.cond_dim, 2 * config.hidden_dim, bias=True)
        nn.init.zeros_(self.final_adaLN.weight)
        nn.init.zeros_(self.final_adaLN.bias)

        self.output_proj = nn.Linear(config.hidden_dim, config.vocab_size, bias=False)
        self.output_proj.weight = self.vocab_embed.weight  # Weight tying

        self._init_weights()

    def _init_weights(self):
        for module in self.modules():
            if isinstance(module, nn.Linear):
                torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
                if module.bias is not None:
                    torch.nn.init.zeros_(module.bias)
            elif isinstance(module, nn.Embedding):
                torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, input_ids, segment_ids, timesteps):
        x = self.vocab_embed(input_ids) + self.segment_embed(segment_ids)
        c = F.silu(self.sigma_map(timesteps))
        rotary_cos_sin = self.rotary_emb(x)

        for block in self.blocks:
            x = block(x, rotary_cos_sin, c)

        shift, scale = self.final_adaLN(c)[:, None, :].chunk(2, dim=2)
        x = self.final_norm(x) * (1 + scale) + shift
        logits = self.output_proj(x)
        return logits

    def count_parameters(self):
        return sum(p.numel() for p in self.parameters() if p.requires_grad)


def compute_diffusion_loss(model, input_ids, segment_ids, target_ids, target_mask, config):
    """Compute masked diffusion training loss (LLaDA-style ELBO)."""
    batch_size = input_ids.shape[0]
    device = input_ids.device

    eps = 1e-5
    t = torch.rand(batch_size, device=device) * (1 - eps) + eps

    mask_prob = t[:, None].expand_as(target_mask)
    random_mask = torch.rand_like(mask_prob) < mask_prob
    diffusion_mask = random_mask & target_mask

    noised_input = input_ids.clone()
    noised_input[diffusion_mask] = config.mask_token_id

    logits = model(noised_input, segment_ids, t)

    logits_flat = logits.view(-1, config.vocab_size)
    targets_flat = target_ids.view(-1)

    if diffusion_mask.sum() == 0:
        zero = torch.tensor(0.0, device=device, requires_grad=True)
        return zero, zero

    ce_loss = F.cross_entropy(logits_flat, targets_flat, reduction='none')

    masked_loss_2d = ce_loss.view(batch_size, -1) * diffusion_mask.float()
    per_example_counts = diffusion_mask.float().sum(dim=1).clamp(min=1.0)
    per_example_loss = masked_loss_2d.sum(dim=1) / per_example_counts

    weighted_loss = (per_example_loss / t).mean()
    unweighted_loss = per_example_loss.mean()

    return weighted_loss, unweighted_loss


@torch.no_grad()
def generate(model, src_ids, src_segment_ids, config, num_steps=50, device='cuda'):
    """Generate translation using iterative unmasking."""
    model.eval()
    batch_size = src_ids.shape[0]
    tgt_len = config.max_tgt_len

    tgt_ids = torch.full((batch_size, tgt_len), config.mask_token_id, device=device)
    tgt_segment_ids = torch.ones(batch_size, tgt_len, dtype=torch.long, device=device)

    input_ids = torch.cat([src_ids, tgt_ids], dim=1)
    segment_ids = torch.cat([src_segment_ids, tgt_segment_ids], dim=1)
    src_len = src_ids.shape[1]

    for step in range(num_steps, 0, -1):
        t = torch.tensor([step / num_steps], device=device).expand(batch_size)
        s = torch.tensor([(step - 1) / num_steps], device=device).expand(batch_size)

        logits = model(input_ids, segment_ids, t)
        tgt_logits = logits[:, src_len:, :]
        predicted_tokens = tgt_logits.argmax(dim=-1)

        current_tgt = input_ids[:, src_len:]
        still_masked = (current_tgt == config.mask_token_id)

        if step > 1:
            remask_prob = s[0].item() / t[0].item() if t[0].item() > 0 else 0.0
            remask = torch.rand_like(predicted_tokens.float()) < remask_prob
            new_tgt = current_tgt.clone()
            unmask_positions = still_masked & ~remask
            new_tgt[unmask_positions] = predicted_tokens[unmask_positions]
        else:
            new_tgt = current_tgt.clone()
            new_tgt[still_masked] = predicted_tokens[still_masked]

        input_ids = torch.cat([src_ids, new_tgt], dim=1)

    return input_ids[:, src_len:]


# ═══════════════════════════════════════════════════════════════
# DATASET
# ═══════════════════════════════════════════════════════════════

class WMT14EnDeDataset(Dataset):
    def __init__(self, data, tokenizer, max_src_len=128, max_tgt_len=128):
        self.data = data
        self.tokenizer = tokenizer
        self.max_src_len = max_src_len
        self.max_tgt_len = max_tgt_len

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

    def __getitem__(self, idx):
        item = self.data[idx]
        en_text = item['translation']['en']
        de_text = item['translation']['de']

        src_enc = self.tokenizer(
            "translate English to German: " + en_text,
            max_length=self.max_src_len, truncation=True,
            padding='max_length', return_tensors=None,
        )
        tgt_enc = self.tokenizer(
            de_text,
            max_length=self.max_tgt_len, truncation=True,
            padding='max_length', return_tensors=None,
        )

        src_ids = src_enc['input_ids']
        tgt_ids = tgt_enc['input_ids']
        segment_ids = [0] * len(src_ids) + [1] * len(tgt_ids)
        full_ids = src_ids + tgt_ids
        target_mask = [0] * len(src_ids) + tgt_enc['attention_mask']

        return {
            'input_ids': torch.tensor(full_ids, dtype=torch.long),
            'segment_ids': torch.tensor(segment_ids, dtype=torch.long),
            'target_ids': torch.tensor(full_ids, dtype=torch.long),
            'target_mask': torch.tensor(target_mask, dtype=torch.bool),
        }


# ═══════════════════════════════════════════════════════════════
# CONFIGURATION
# ═══════════════════════════════════════════════════════════════

MODEL_CONFIG = dict(
    vocab_size=None,
    max_src_len=128,
    max_tgt_len=128,
    hidden_dim=512,
    n_heads=8,
    n_blocks=12,
    dropout=0.1,
    cond_dim=128,
    mask_token_id=None,
    pad_token_id=None,
)

TRAIN_CONFIG = dict(
    learning_rate=3e-4,
    weight_decay=0.01,
    warmup_steps=4000,
    max_steps=200_000,
    batch_size=64,
    gradient_accumulation_steps=4,
    eval_every=5000,
    save_every=10000,
    log_every=100,
    max_grad_norm=1.0,
    num_gen_steps=50,
    fp16=True,
    seed=42,
)

HUB_MODEL_ID = "vedkdev/text-diffusion-en-de"
TOKENIZER_NAME = "Helsinki-NLP/opus-mt-en-de"


# ═══════════════════════════════════════════════════════════════
# TRAINING
# ═══════════════════════════════════════════════════════════════

def train():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    torch.manual_seed(TRAIN_CONFIG['seed'])

    print(f"Device: {device}")
    print(f"Loading tokenizer: {TOKENIZER_NAME}")

    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
    if tokenizer.mask_token is None:
        tokenizer.add_special_tokens({'mask_token': '<mask>'})

    mask_token_id = tokenizer.mask_token_id
    pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0

    print(f"Vocab size: {len(tokenizer)}")
    print(f"Mask token ID: {mask_token_id}, Pad token ID: {pad_token_id}")

    MODEL_CONFIG['vocab_size'] = len(tokenizer)
    MODEL_CONFIG['mask_token_id'] = mask_token_id
    MODEL_CONFIG['pad_token_id'] = pad_token_id
    config = DiffusionTranslatorConfig(**MODEL_CONFIG)

    model = DiffusionTranslator(config).to(device)
    print(f"Model parameters: {model.count_parameters():,}")

    print("Loading WMT14 EN-DE dataset...")
    dataset = load_dataset("wmt/wmt14", "de-en", trust_remote_code=True)
    train_data = dataset['train']
    val_data = dataset['validation']
    print(f"Train: {len(train_data):,} | Val: {len(val_data):,}")

    train_dataset = WMT14EnDeDataset(train_data, tokenizer, config.max_src_len, config.max_tgt_len)
    val_dataset = WMT14EnDeDataset(val_data, tokenizer, config.max_src_len, config.max_tgt_len)

    train_loader = DataLoader(train_dataset, batch_size=TRAIN_CONFIG['batch_size'],
                              shuffle=True, num_workers=4, pin_memory=True, drop_last=True)
    val_loader = DataLoader(val_dataset, batch_size=TRAIN_CONFIG['batch_size'],
                            shuffle=False, num_workers=2, pin_memory=True)

    optimizer = torch.optim.AdamW(model.parameters(), lr=TRAIN_CONFIG['learning_rate'],
                                  weight_decay=TRAIN_CONFIG['weight_decay'], betas=(0.9, 0.98), eps=1e-8)
    scheduler = get_cosine_schedule_with_warmup(optimizer,
                                                num_warmup_steps=TRAIN_CONFIG['warmup_steps'],
                                                num_training_steps=TRAIN_CONFIG['max_steps'])

    scaler = torch.amp.GradScaler('cuda') if (TRAIN_CONFIG['fp16'] and device.type == 'cuda') else None

    trackio.init(project="text-diffusion-en-de", name="v1-wmt14-dit12-512d")

    global_step = 0
    best_val_loss = float('inf')
    accum_loss = 0.0
    accum_loss_uw = 0.0
    accum_count = 0

    eff_bs = TRAIN_CONFIG['batch_size'] * TRAIN_CONFIG['gradient_accumulation_steps']
    print(f"\n=== Starting Training ===")
    print(f"Effective batch size: {eff_bs} | Max steps: {TRAIN_CONFIG['max_steps']:,}")
    print(f"Warmup: {TRAIN_CONFIG['warmup_steps']:,} | LR: {TRAIN_CONFIG['learning_rate']}")

    model.train()
    optimizer.zero_grad()
    data_iter = iter(train_loader)
    start_time = time.time()

    total_micro_steps = TRAIN_CONFIG['max_steps'] * TRAIN_CONFIG['gradient_accumulation_steps']
    for step in range(1, total_micro_steps + 1):
        try:
            batch = next(data_iter)
        except StopIteration:
            data_iter = iter(train_loader)
            batch = next(data_iter)

        input_ids = batch['input_ids'].to(device)
        segment_ids = batch['segment_ids'].to(device)
        target_ids = batch['target_ids'].to(device)
        target_mask = batch['target_mask'].to(device)

        if scaler is not None:
            with torch.amp.autocast('cuda'):
                wl, uwl = compute_diffusion_loss(model, input_ids, segment_ids, target_ids, target_mask, config)
                loss = wl / TRAIN_CONFIG['gradient_accumulation_steps']
            scaler.scale(loss).backward()
        else:
            wl, uwl = compute_diffusion_loss(model, input_ids, segment_ids, target_ids, target_mask, config)
            loss = wl / TRAIN_CONFIG['gradient_accumulation_steps']
            loss.backward()

        accum_loss += wl.item()
        accum_loss_uw += uwl.item()
        accum_count += 1

        if step % TRAIN_CONFIG['gradient_accumulation_steps'] == 0:
            if scaler is not None:
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), TRAIN_CONFIG['max_grad_norm'])
                scaler.step(optimizer)
                scaler.update()
            else:
                torch.nn.utils.clip_grad_norm_(model.parameters(), TRAIN_CONFIG['max_grad_norm'])
                optimizer.step()

            scheduler.step()
            optimizer.zero_grad()
            global_step += 1

            # Log
            if global_step % TRAIN_CONFIG['log_every'] == 0:
                avg_l = accum_loss / accum_count
                avg_uw = accum_loss_uw / accum_count
                elapsed = time.time() - start_time
                sps = global_step / elapsed
                lr = scheduler.get_last_lr()[0]
                print(f"step={global_step} | loss={avg_l:.4f} | ce_loss={avg_uw:.4f} | lr={lr:.2e} | steps/s={sps:.2f}")
                trackio.log({"train/loss_weighted": avg_l, "train/loss_ce": avg_uw,
                             "train/learning_rate": lr, "train/steps_per_sec": sps}, step=global_step)
                accum_loss = 0.0
                accum_loss_uw = 0.0
                accum_count = 0

            # Eval
            if global_step % TRAIN_CONFIG['eval_every'] == 0:
                vl, vuw = evaluate(model, val_loader, config, device, scaler is not None)
                print(f"  [EVAL] step={global_step} | val_loss={vl:.4f} | val_ce={vuw:.4f}")
                trackio.log({"eval/loss_weighted": vl, "eval/loss_ce": vuw}, step=global_step)

                if global_step % (TRAIN_CONFIG['eval_every'] * 4) == 0:
                    bleu = evaluate_bleu(model, tokenizer, config, device, num_samples=100,
                                         num_steps=TRAIN_CONFIG['num_gen_steps'])
                    trackio.log({"eval/sacrebleu": bleu}, step=global_step)

                if vuw < best_val_loss:
                    best_val_loss = vuw
                    save_model(model, config, tokenizer, global_step, is_best=True)
                model.train()

            # Save + push
            if global_step % TRAIN_CONFIG['save_every'] == 0:
                save_model(model, config, tokenizer, global_step, push_to_hub=True)

    # Final
    save_model(model, config, tokenizer, global_step, push_to_hub=True)
    print("\n=== Final BLEU Evaluation ===")
    bleu = evaluate_bleu(model, tokenizer, config, device, num_samples=200,
                          num_steps=TRAIN_CONFIG['num_gen_steps'])
    trackio.log({"eval/final_sacrebleu": bleu}, step=global_step)
    print(f"\n=== Training Complete === Final BLEU: {bleu:.2f}")


def evaluate(model, val_loader, config, device, use_fp16=True):
    model.eval()
    total_l = total_uw = 0.0
    count = 0
    with torch.no_grad():
        for i, batch in enumerate(val_loader):
            if i >= 50:
                break
            ids = batch['input_ids'].to(device)
            seg = batch['segment_ids'].to(device)
            tgt = batch['target_ids'].to(device)
            mask = batch['target_mask'].to(device)
            if use_fp16:
                with torch.amp.autocast('cuda'):
                    wl, uwl = compute_diffusion_loss(model, ids, seg, tgt, mask, config)
            else:
                wl, uwl = compute_diffusion_loss(model, ids, seg, tgt, mask, config)
            total_l += wl.item()
            total_uw += uwl.item()
            count += 1
    return total_l / max(count, 1), total_uw / max(count, 1)


def evaluate_bleu(model, tokenizer, config, device, num_samples=100, num_steps=50):
    import sacrebleu
    model.eval()
    ds = load_dataset("wmt/wmt14", "de-en", split="test", trust_remote_code=True)
    refs, hyps = [], []
    for i in range(min(num_samples, len(ds))):
        en = ds[i]['translation']['en']
        de_ref = ds[i]['translation']['de']
        enc = tokenizer("translate English to German: " + en, max_length=config.max_src_len,
                        truncation=True, padding='max_length', return_tensors='pt')
        src_ids = enc['input_ids'].to(device)
        src_seg = torch.zeros_like(src_ids)
        with torch.no_grad():
            if device.type == 'cuda':
                with torch.amp.autocast('cuda'):
                    gen = generate(model, src_ids, src_seg, config, num_steps=num_steps, device=device)
            else:
                gen = generate(model, src_ids, src_seg, config, num_steps=num_steps, device=device)
        hyp = tokenizer.decode(gen[0], skip_special_tokens=True)
        refs.append(de_ref)
        hyps.append(hyp)
        if i < 5:
            print(f"  EN: {en[:100]}")
            print(f"  REF: {de_ref[:100]}")
            print(f"  GEN: {hyp[:100]}")
            print()
    bleu = sacrebleu.corpus_bleu(hyps, [refs])
    print(f"SacreBLEU: {bleu.score:.2f}")
    return bleu.score


def save_model(model, config, tokenizer, step, is_best=False, push_to_hub=False):
    save_dir = "checkpoints/best" if is_best else f"checkpoints/step-{step}"
    os.makedirs(save_dir, exist_ok=True)
    torch.save(model.state_dict(), os.path.join(save_dir, "model.pt"))
    config_dict = {k: getattr(config, k) for k in [
        'vocab_size', 'max_src_len', 'max_tgt_len', 'hidden_dim',
        'n_heads', 'n_blocks', 'dropout', 'cond_dim', 'mask_token_id', 'pad_token_id'
    ]}
    with open(os.path.join(save_dir, "config.json"), "w") as f:
        json.dump(config_dict, f, indent=2)
    tokenizer.save_pretrained(save_dir)
    if push_to_hub:
        push_model_to_hub(save_dir, step, config)
    print(f"  Saved checkpoint to {save_dir}")


def push_model_to_hub(save_dir, step, config):
    from huggingface_hub import HfApi, upload_folder
    api = HfApi()
    try:
        api.create_repo(HUB_MODEL_ID, exist_ok=True, private=False)
    except Exception as e:
        print(f"  Warning creating repo: {e}")

    readme = f"""---
tags:
- text-diffusion
- machine-translation
- en-de
- masked-diffusion
language:
- en
- de
datasets:
- wmt/wmt14
---

# Text Diffusion Model for EN→DE Translation

A **masked discrete diffusion** model for English-to-German machine translation,
trained from scratch on WMT14 EN-DE.

## Architecture
- **Type**: Masked Discrete Diffusion (inspired by MDLM + LLaDA)
- **Backbone**: DiT (Diffusion Transformer) with adaptive LayerNorm (adaLN)
- **Parameters**: ~72M
- **Blocks**: {config.n_blocks} DiT blocks, hidden_dim={config.hidden_dim}, heads={config.n_heads}
- **Tokenizer**: {TOKENIZER_NAME} (~58K vocab)
- **Max sequence**: {config.max_src_len} src + {config.max_tgt_len} tgt tokens

## Training
- **Dataset**: WMT14 EN-DE (~4.5M pairs)
- **Method**: Masked discrete diffusion with ELBO weighting (1/t)
- **Optimizer**: AdamW, lr=3e-4, cosine with 4K warmup
- **Effective batch size**: {TRAIN_CONFIG['batch_size'] * TRAIN_CONFIG['gradient_accumulation_steps']}
- **Training steps**: {step:,}

## How It Works
1. Source (EN) + target (DE) tokens concatenated: `[source | target]`
2. Training: target tokens randomly masked with prob `t ~ U(0,1)`, predict masked tokens
3. Inference: start fully masked → iteratively unmask over {TRAIN_CONFIG['num_gen_steps']} steps

## References
- [MDLM](https://arxiv.org/abs/2406.07524) | [LLaDA](https://arxiv.org/abs/2502.09992) | [DiNoiSer](https://arxiv.org/abs/2302.10025)
"""
    with open(os.path.join(save_dir, "README.md"), "w") as f:
        f.write(readme)

    upload_folder(repo_id=HUB_MODEL_ID, folder_path=save_dir,
                  commit_message=f"Checkpoint step {step}")
    print(f"  Pushed to hub: https://huggingface.co/{HUB_MODEL_ID}")


if __name__ == "__main__":
    train()