"""
Arabic zero-shot router with a real routing head.

Two earlier formulations, and why this one:

  fixed-slot head  - `num_labels = max_lanes` over one pooled vector. Scored exactly 1/n at every
                     lane count. The head is positional (slot i's logit is w_i . h) and the corpus
                     randomises lane order, so there is nothing to learn.
  pairwise         - score each (text, category) separately, argmax. Works (0.842 on unseen lane
                     sets) but costs n forward passes and each category is scored in isolation,
                     never against its competitors.

Here the categories live in the sequence and get their *own* pooled vectors, which a shared linear
head turns into one logit each; the softmax is over the categories present. The head is shared
across positions, so it scores category *content*, not slot index - which is what makes it
position-invariant and zero-shot.

Layout is text-first, categories-second, deliberately. The reference encoder is bidirectional so
order does not matter there; our base is causal, and this ordering is what lets every category
token attend to the whole text. Reversed, the categories would be encoded blind to the text.

  النص:
  {text}

  الفئات:
  - فئة أولى
  - فئة ثانية

Each category's character span is mapped to token indices via the tokenizer's offset mapping and
mean-pooled. Slots beyond a row's category count are masked to -inf before the loss.
"""
import json
import os
from collections import defaultdict
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset
from transformers import AutoModel, AutoTokenizer, Trainer, TrainingArguments

os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0")

BASE_MODEL = os.environ.get("BASE_MODEL", "oddadmix/50M-2048-Emhotob")
OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "./Nawah-Router-head-50M")
DATA = Path(os.environ.get("DATA_DIR", "data"))
MAX_ROUTES = int(os.environ.get("MAX_ROUTES", 9))
MAX_LENGTH = int(os.environ.get("MAX_LENGTH", 320))
LR = float(os.environ.get("LR", 3e-4))
EPOCHS = float(os.environ.get("EPOCHS", 3))
BATCH_SIZE = int(os.environ.get("BATCH_SIZE", 32))
WARMUP = int(os.environ.get("WARMUP", 500))
SEED = 42


def build_text(row):
    """-> (full string, [(start, end) char span per category])."""
    head = f"النص:\n{row['text']}\n\nالفئات:\n"
    s = head
    spans = []
    for c in row["routes"]:
        s += "- "
        spans.append((len(s), len(s) + len(c)))
        s += c + "\n"
    return s, spans


class RouterModel(nn.Module):
    """Backbone + a shared linear scorer applied to each category's pooled span."""

    def __init__(self, base_model):
        super().__init__()
        self.backbone = AutoModel.from_pretrained(base_model, dtype=torch.float32)
        h = self.backbone.config.hidden_size
        self.score = nn.Sequential(nn.Linear(h, h), nn.GELU(), nn.Linear(h, 1))
        self.config = self.backbone.config

    def forward(self, input_ids, attention_mask, cat_pool, n_routes, labels=None):
        hs = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
        # cat_pool: (B, K, T) row-normalised selector -> (B, K, H)
        cat_vecs = torch.bmm(cat_pool.to(hs.dtype), hs)
        logits = self.score(cat_vecs).squeeze(-1)                       # (B, K)
        ar = torch.arange(logits.size(1), device=logits.device)[None, :]
        logits = logits.masked_fill(ar >= n_routes[:, None], torch.finfo(logits.dtype).min)
        loss = F.cross_entropy(logits, labels) if labels is not None else None
        return {"loss": loss, "logits": logits}


def load(name):
    p = DATA / f"{name}.jsonl"
    return [json.loads(l) for l in open(p, encoding="utf-8")] if p.exists() else []


class RouterDataset(Dataset):
    def __init__(self, rows, tok, max_length):
        self.rows, self.tok, self.max_length = rows, tok, max_length

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

    def __getitem__(self, i):
        r = self.rows[i]
        text, spans = build_text(r)
        enc = self.tok(text, return_offsets_mapping=True, add_special_tokens=False,
                       truncation=True, max_length=self.max_length)
        ids, offs = enc["input_ids"], enc["offset_mapping"]
        pool = torch.zeros(MAX_ROUTES, len(ids))
        for ci, (s, e) in enumerate(spans[:MAX_ROUTES]):
            idxs = [t for t, (a, b) in enumerate(offs) if a < e and b > s and a != b]
            if idxs:
                pool[ci, idxs] = 1.0 / len(idxs)
        return {"input_ids": torch.tensor(ids, dtype=torch.long), "cat_pool": pool,
                "n_routes": torch.tensor(min(len(r["routes"]), MAX_ROUTES), dtype=torch.long),
                "labels": torch.tensor(min(r["label"], MAX_ROUTES - 1), dtype=torch.long)}


class Collator:
    def __init__(self, pad_id):
        self.pad_id = pad_id

    def __call__(self, feats):
        T = max(f["input_ids"].size(0) for f in feats)
        ids, att, pools = [], [], []
        for f in feats:
            n = f["input_ids"].size(0); pad = T - n
            ids.append(torch.cat([f["input_ids"], torch.full((pad,), self.pad_id, dtype=torch.long)]))
            att.append(torch.cat([torch.ones(n, dtype=torch.long), torch.zeros(pad, dtype=torch.long)]))
            pools.append(F.pad(f["cat_pool"], (0, pad)))
        return {"input_ids": torch.stack(ids), "attention_mask": torch.stack(att),
                "cat_pool": torch.stack(pools),
                "n_routes": torch.stack([f["n_routes"] for f in feats]),
                "labels": torch.stack([f["labels"] for f in feats])}


def metrics_fn(p):
    return {"accuracy": float((np.asarray(p.predictions).argmax(-1) ==
                               np.asarray(p.label_ids)).mean())}


@torch.no_grad()
def report(model, tok, rows, name, batch=64):
    model.eval()
    ds = RouterDataset(rows, tok, MAX_LENGTH); coll = Collator(tok.pad_token_id)
    preds = []
    for i in range(0, len(rows), batch):
        b = coll([ds[j] for j in range(i, min(i + batch, len(rows)))])
        b = {k: v.to(next(model.parameters()).device) for k, v in b.items()}
        b.pop("labels")
        preds += model(**b)["logits"].argmax(-1).tolist()
    correct = [int(p == r["label"]) for p, r in zip(preds, rows)]
    acc = sum(correct) / len(rows)
    rand = sum(1 / len(r["routes"]) for r in rows) / len(rows)
    print(f"\n[{name}]  n={len(rows):,}  route accuracy {acc:.4f}   random {rand:.4f}")
    out = {"accuracy": acc, "random_baseline": rand, "n": len(rows)}
    for key in ("n_routes", "difficulty", "mode"):
        if key == "mode" and not all("mode" in r for r in rows):
            continue
        b = defaultdict(lambda: [0, 0])
        for c, r in zip(correct, rows):
            k = len(r["routes"]) if key == "n_routes" else r[key]
            b[k][1] += 1; b[k][0] += c
        print(f"    by {key:<10} " +
              "  ".join(f"{k}:{v[0]/v[1]:.3f}(n={v[1]})" for k, v in sorted(b.items(), key=str)))
        out[f"by_{key}"] = {str(k): {"acc": v[0]/v[1], "n": v[1]} for k, v in b.items()}
    return out


def main():
    tok = AutoTokenizer.from_pretrained(BASE_MODEL)
    if tok.pad_token_id is None:
        tok.pad_token = tok.eos_token
    model = RouterModel(BASE_MODEL)
    print(f"[*] {BASE_MODEL} | params {sum(p.numel() for p in model.parameters())/1e6:.2f}M "
          f"| hidden {model.config.hidden_size}")

    train = load("train")
    evals = {"unseen_lanes": load("eval_unseen_lanes"),
             "unseen_domain": load("eval_unseen_domain"),
             # v2 only: axes held out of training entirely. The strongest zero-shot test, so the
             # checkpoint is selected on it when it exists.
             "unseen_axis": load("eval_unseen_axis"),
             "hard": load("eval_hard")}
    evals = {k: v for k, v in evals.items() if v}
    print(f"[*] train {len(train):,} | " + " | ".join(f"{k} {len(v):,}" for k, v in evals.items()))

    args = TrainingArguments(
        output_dir=OUTPUT_DIR, num_train_epochs=EPOCHS,
        per_device_train_batch_size=BATCH_SIZE, per_device_eval_batch_size=64,
        learning_rate=LR, lr_scheduler_type="cosine", warmup_steps=WARMUP,
        max_grad_norm=1.0, bf16=True, logging_steps=200,
        eval_strategy="steps", eval_steps=1000, save_strategy="steps", save_steps=1000,
        save_total_limit=2, load_best_model_at_end=True,
        metric_for_best_model=("eval_unseen_axis_accuracy" if "unseen_axis" in evals
                               else "eval_unseen_domain_accuracy"), greater_is_better=True,
        report_to=[], seed=SEED, dataloader_num_workers=4, remove_unused_columns=False,
        label_names=["labels"])

    trainer = Trainer(model=model, args=args,
                      train_dataset=RouterDataset(train, tok, MAX_LENGTH),
                      eval_dataset={k: RouterDataset(v, tok, MAX_LENGTH)
                                    for k, v in evals.items() if v},
                      data_collator=Collator(tok.pad_token_id), compute_metrics=metrics_fn)
    trainer.train()

    Path(OUTPUT_DIR).mkdir(exist_ok=True)
    torch.save(model.state_dict(), Path(OUTPUT_DIR, "router_model.pt"))
    model.backbone.save_pretrained(OUTPUT_DIR); tok.save_pretrained(OUTPUT_DIR)
    results = {k: report(model, tok, v, k) for k, v in evals.items() if v}
    Path(OUTPUT_DIR, "train_metrics.json").write_text(json.dumps(
        {"results": results, "base_model": BASE_MODEL, "formulation": "routing_head",
         "log_history": trainer.state.log_history}, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"\n[+] done -> {OUTPUT_DIR}")


if __name__ == "__main__":
    main()