"""© KAND CA 2026 - Nawah-BERT-6M-bilingual: BertForMaskedLM, 50/50 Arabic/English, from scratch.

Same 6M rung as train-bert-ladder.py (hidden 128, 8 layers, 2 heads, 5,993,600 params) and the
same sparse-MLM-head / O(1)-resume machinery, but three things differ:

1. **A tokenizer trained for both languages.** custom_llama_tokenizer (32K, Arabic-only fit) costs
   English ~2.02 tok/word; train_tokenizer_bilingual.py fixes that at the source by interleaving
   Arabic and English documents before learning BPE merges, at a real but modest cost to Arabic
   (1.20 -> 1.37 tok/word) in exchange for English falling to 1.15. Every bin below is tokenized
   with `bilingual32k_tokenizer`, wrapped for AutoTokenizer in `bilingual32k-tokenizer-hf/`.

2. **Two 5B-token bins, one stream.** `DualMemmapDataset` (forked from train_mlm_dialect.py's
   class of the same name - see that module's docstring for why the fixed-permutation design
   keeps resume O(1)) draws a 50/50 share of tokens_ar_5B_bi32k.bin and tokens_en_5B_bi32k.bin.
   Both bins are exactly 5B tokens, so share=0.5 needs no truncation on either side - a genuinely
   balanced single pass, unlike the dialect run's 70/30 with dialect repeated to fill the budget.

3. **BertForMaskedLM, not ModernBERT** - this is the same architecture as the existing 6M rung,
   so a downstream comparison (router head, RuleCheck, Guard) is a clean like-for-like swap of the
   backbone, holding architecture and task fixed and varying only the pretraining corpus and
   tokenizer.

    python train_bert6m_bilingual.py --bench 20
    python train_bert6m_bilingual.py
"""
import argparse
import os

import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, SequentialSampler
from transformers import (BertConfig, BertForMaskedLM, Trainer, TrainingArguments)
from transformers.modeling_outputs import MaskedLMOutput
from transformers.trainer_utils import get_last_checkpoint

ROOT = os.path.dirname(os.path.abspath(__file__))

AR_BIN, AR_TOKENS = os.path.join(ROOT, "tokens_ar_5B_bi32k.bin"), 5_000_000_000
EN_BIN, EN_TOKENS = os.path.join(ROOT, "tokens_en_5B_bi32k.bin"), 5_000_000_000
TOKENIZER_DIR = os.path.join(ROOT, "bilingual32k-tokenizer-hf")
OUTPUT_DIR = os.path.join(ROOT, "Nawah-BERT-6M-bilingual-pretrain")
FINAL_DIR = os.path.join(ROOT, "Nawah-BERT-6M-bilingual-pretrain-FINAL")

SEQ_LEN = 2048
VOCAB = 32_000
LAYERS = 8
HEAD_DIM = 64
HIDDEN = 128
EXPECTED_PARAMS = 5_993_600
PERM_SEED = 1234
HOLDOUT_CHUNKS = 2_048          # reserved for eval, never trained on, per language

MICRO_BATCH = 16
GRAD_ACCUM = 8
CHUNKS_PER_STEP = MICRO_BATCH * GRAD_ACCUM      # 128
TOKENS_PER_STEP = CHUNKS_PER_STEP * SEQ_LEN     # 262,144

MLM_PROB_MAIN, MLM_PROB_FINAL = 0.30, 0.15
MLM_SWITCH_FRAC = 0.89
LR = 6e-4
WARMUP_STEPS = 2_000


class DualMemmapDataset(Dataset):
    """50/50 Arabic/English chunks in one fixed, resumable order.

    See train_mlm_dialect.py's class of the same name for the full rationale; the only change
    here is share=0.5 over two EQUAL-SIZE bins, so neither side needs truncating or repeating -
    every chunk of both bins is used exactly once.
    """

    def __init__(self, ar_bin, ar_tokens, en_bin, en_tokens, share=0.5, seq_len=SEQ_LEN,
                start_chunk=0, seed=PERM_SEED, chunks_per_step=CHUNKS_PER_STEP, max_steps=None):
        self.seq_len = seq_len
        self.paths = [ar_bin, en_bin]
        self._data = [None, None]
        self.n_tok = [ar_tokens, en_tokens]

        for p, n in zip(self.paths, self.n_tok):
            if not os.path.exists(p):
                raise SystemExit(f"[!] token bin not found: {p}")
            size = os.path.getsize(p)
            if size != n * 2:
                raise SystemExit(f"[!] {p} is {size:,} bytes; expected {n:,} tokens "
                                 f"({n * 2:,} bytes). Refusing to guess.")

        n_ar = ar_tokens // seq_len
        n_en = en_tokens // seq_len
        # Both bins are the same size and share=0.5, so this reduces to n_ar == n_en with no
        # truncation on either side - written generally so a future unequal-size bin still works.
        want_ar = n_ar
        want_en = min(n_en, int(want_ar * (1 - share) / share))
        idx = [(0, i) for i in range(want_ar)] + [(1, j) for j in range(want_en)]
        rng = np.random.default_rng(seed)
        order = rng.permutation(len(idx))
        self.index = np.asarray(idx, dtype=np.int64)[order]

        self.n_total = len(self.index)
        self.n_train = self.n_total - HOLDOUT_CHUNKS
        self.start_chunk = start_chunk
        self.chunks_per_step = chunks_per_step
        self.total_steps = max_steps or (self.n_train // chunks_per_step)
        self.switch_step = int(self.total_steps * MLM_SWITCH_FRAC)

        print(f"[=] Arabic  {ar_bin} ({ar_tokens / 1e9:.2f}B tokens, {n_ar:,} chunks -> "
              f"{want_ar:,} used)")
        print(f"[=] English {en_bin} ({en_tokens / 1e9:.2f}B tokens, {n_en:,} chunks -> "
              f"{want_en:,} used)")
        print(f"[+] {self.n_total:,} chunks, realised Arabic share "
              f"{want_ar / max(self.n_total, 1):.3f}; {HOLDOUT_CHUNKS:,} held out; "
              f"{self.total_steps:,} steps; mask 30%->15% at step {self.switch_step:,}")

    def data(self, src):
        if self._data[src] is None:
            n = self.n_tok[src] // self.seq_len * self.seq_len
            self._data[src] = np.memmap(self.paths[src], dtype=np.uint16, mode="r", shape=(n,))
        return self._data[src]

    def holdout_indices(self):
        return self.index[self.n_train:]

    def chunk(self, src, chunk_id):
        s = int(chunk_id) * self.seq_len
        arr = np.asarray(self.data(src)[s:s + self.seq_len], dtype=np.int64)
        return torch.from_numpy(arr)

    def __len__(self):
        return self.n_train

    def __getitem__(self, i):
        pos = (self.start_chunk + i) % self.n_train
        step = (self.start_chunk + i) // self.chunks_per_step
        src, cid = self.index[pos]
        return {"input_ids": self.chunk(int(src), int(cid)),
                "mlm_prob": MLM_PROB_MAIN if step < self.switch_step else MLM_PROB_FINAL}


def mlm_collate(features):
    """80/10/10 BERT masking, vectorised, per-item mlm_prob (carries the anneal)."""
    from transformers import AutoTokenizer
    tok = mlm_collate.tok
    ids = torch.stack([f["input_ids"] for f in features])
    probs = torch.tensor([f["mlm_prob"] for f in features]).unsqueeze(1)
    special = torch.zeros_like(ids, dtype=torch.bool)
    for tid in (tok.bos_token_id, tok.eos_token_id, tok.pad_token_id):
        special |= ids == tid
    prob_matrix = probs.expand_as(ids).clone()
    prob_matrix.masked_fill_(special, 0.0)
    masked = torch.bernoulli(prob_matrix).bool()
    labels = ids.clone()
    labels[~masked] = -100

    replace = torch.bernoulli(torch.full(ids.shape, 0.8)).bool() & masked
    ids = ids.clone()
    ids[replace] = tok.mask_token_id
    random_repl = torch.bernoulli(torch.full(ids.shape, 0.5)).bool() & masked & ~replace
    rand_ids = torch.randint(5, tok.vocab_size, ids.shape, dtype=ids.dtype)
    ids[random_repl] = rand_ids[random_repl]
    return {"input_ids": ids, "labels": labels}


class SparseBertMLM(BertForMaskedLM):
    """BertForMaskedLM, LM head run only at masked positions - see train-bert-ladder.py."""

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None,
                position_ids=None, labels=None, num_items_in_batch=None, **kw):
        out = self.bert(input_ids=input_ids, attention_mask=attention_mask,
                        token_type_ids=token_type_ids, position_ids=position_ids, **kw)
        hidden = out.last_hidden_state
        if labels is None:
            return MaskedLMOutput(logits=self.cls(hidden))
        flat = labels.reshape(-1)
        idx = (flat != -100).nonzero(as_tuple=True)[0]
        sel = hidden.reshape(-1, hidden.size(-1)).index_select(0, idx)
        logits = self.cls(sel).float()
        tgt = flat.index_select(0, idx)
        if num_items_in_batch is not None:
            loss = F.cross_entropy(logits, tgt, reduction="sum") / num_items_in_batch
        else:
            loss = F.cross_entropy(logits, tgt)
        return MaskedLMOutput(loss=loss, logits=logits)


SparseBertMLM.__name__ = "BertForMaskedLM"


class OffsetTrainer(Trainer):
    def _get_train_sampler(self, *a, **kw):
        return SequentialSampler(self.train_dataset)


def reconcile_tf32():
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True


def latest_checkpoint(d):
    return get_last_checkpoint(d) if os.path.isdir(d) else None


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--bench", type=int, default=0)
    p.add_argument("--micro-batch", type=int, default=MICRO_BATCH)
    p.add_argument("--grad-accum", type=int, default=GRAD_ACCUM)
    a = p.parse_args()
    if a.micro_batch * a.grad_accum != CHUNKS_PER_STEP:
        raise SystemExit(f"[!] micro-batch x grad-accum must be {CHUNKS_PER_STEP}")

    reconcile_tf32()
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(TOKENIZER_DIR)
    mlm_collate.tok = tok

    ds = DualMemmapDataset(AR_BIN, AR_TOKENS, EN_BIN, EN_TOKENS,
                           chunks_per_step=a.micro_batch * a.grad_accum,
                           max_steps=a.bench or None)
    max_steps = a.bench or ds.total_steps

    ckpt = latest_checkpoint(OUTPUT_DIR) if not a.bench else None
    if ckpt:
        model = SparseBertMLM.from_pretrained(ckpt, dtype=torch.float32)
        step = int(os.path.basename(ckpt).split("-")[1])
        ds.start_chunk = step * ds.chunks_per_step
        print(f"[*] resuming {ckpt} at step {step:,} -> chunk {ds.start_chunk:,}", flush=True)
    else:
        cfg = BertConfig(vocab_size=VOCAB, hidden_size=HIDDEN, num_hidden_layers=LAYERS,
                         num_attention_heads=HIDDEN // HEAD_DIM, intermediate_size=4 * HIDDEN,
                         max_position_embeddings=SEQ_LEN, tie_word_embeddings=True,
                         pad_token_id=tok.pad_token_id)
        model = SparseBertMLM(cfg)
    n = sum(q.numel() for q in model.parameters())
    print(f"[*] parameters {n:,} (expected {EXPECTED_PARAMS:,})", flush=True)
    if not ckpt and n != EXPECTED_PARAMS:
        raise SystemExit(f"[!] parameter count {n:,} != expected {EXPECTED_PARAMS:,}")

    args = TrainingArguments(
        output_dir=OUTPUT_DIR, max_steps=max_steps,
        per_device_train_batch_size=a.micro_batch, gradient_accumulation_steps=a.grad_accum,
        learning_rate=LR, lr_scheduler_type="cosine",
        warmup_steps=0 if a.bench else WARMUP_STEPS,
        weight_decay=0.1, adam_beta1=0.9, adam_beta2=0.95, max_grad_norm=1.0,
        optim="adamw_torch_fused", bf16=True, torch_compile=True,
        save_strategy="no" if a.bench else "steps", save_steps=250, save_total_limit=2,
        logging_steps=10 if a.bench else 100,
        report_to="none" if a.bench else "tensorboard",
        dataloader_num_workers=min(8, (os.cpu_count() or 8) // 2),
        dataloader_pin_memory=True, remove_unused_columns=False, ignore_data_skip=True,
        seed=PERM_SEED, disable_tqdm=False)

    trainer = OffsetTrainer(model=model, args=args, train_dataset=ds, data_collator=mlm_collate,
                            processing_class=tok)
    trainer.train(resume_from_checkpoint=ckpt)

    if a.bench:
        print(f"[*] peak VRAM: {torch.cuda.max_memory_allocated() / 2**30:.2f} GiB")
        print("[*] benchmark finished - nothing saved.")
        return
    trainer.save_model(FINAL_DIR)
    tok.save_pretrained(FINAL_DIR)
    print(f"[+] saved -> {FINAL_DIR}")


if __name__ == "__main__":
    main()