"""Arm A data: memmapped Vocos mels + byte tokens with per-token language ids; frame-budget bucketing.
Beta-binomial alignment priors are computed in DataLoader workers."""
import json, os, random
import numpy as np
import torch
from torch.utils.data import Dataset, Sampler

from .text import encode, LANG2ID
from .model import beta_binomial_prior

N_MELS = 100


def load_meta(root, datasets, split="train", max_frames=1400, max_tokens=420):
    items, mm = [], {}
    for ds in datasets:
        p = f"{root}/{ds}/meta.jsonl"
        if not os.path.exists(p) or os.path.getsize(p) == 0:
            continue
        mm[ds] = np.memmap(f"{root}/{ds}/mels.f16", dtype=np.float16, mode="r").reshape(-1, N_MELS)
        for l in open(p, encoding="utf-8"):
            m = json.loads(l)
            if m["split"] != split or m["nf"] > max_frames:
                continue
            if len(m["text"].encode("utf-8")) + 2 > max_tokens:
                continue
            items.append(m)
    return items, mm


class TTSData(Dataset):
    def __init__(self, items, mm, spk2id):
        self.items, self.mm, self.spk2id = items, mm, spk2id

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

    def __getitem__(self, i):
        m = self.items[i]
        mel = np.asarray(self.mm[m["ds"]][m["off"]: m["off"] + m["nf"]], dtype=np.float32)
        tok, _ = encode(f"<{m['lang']}>{m['text']}</{m['lang']}>")
        lang = [LANG2ID[m["lang"]]] * len(tok)
        return dict(tok=tok, lang=lang, mel=mel, spk=self.spk2id[m["spk"]], id=m["id"])


def collate(batch):
    B = len(batch)
    N = max(len(b["tok"]) for b in batch)
    T = max(len(b["mel"]) for b in batch)
    tok = torch.zeros(B, N, dtype=torch.long)
    lang = torch.zeros(B, N, dtype=torch.long)
    mel = torch.zeros(B, T, N_MELS)
    prior = torch.full((B, T, N), -1e4)
    tl, ml = torch.zeros(B, dtype=torch.long), torch.zeros(B, dtype=torch.long)
    for i, b in enumerate(batch):
        n, t = len(b["tok"]), len(b["mel"])
        tok[i, :n] = torch.tensor(b["tok"]); lang[i, :n] = torch.tensor(b["lang"])
        mel[i, :t] = torch.from_numpy(b["mel"]); tl[i], ml[i] = n, t
        prior[i, :t, :n] = beta_binomial_prior(n, t)
    spk = torch.tensor([b["spk"] for b in batch])
    return dict(tok=tok, lang=lang, tlens=tl, mel=mel, mlens=ml, spk=spk, prior=prior)


class FrameBucketSampler(Sampler):
    """Batches of similar length under a padded-frame budget; dataset weights by repetition."""

    def __init__(self, items, frame_budget=40000, weights=None, seed=0):
        self.items, self.budget, self.seed, self.epoch = items, frame_budget, seed, 0
        self.weights = weights or {}

    def _batches(self):
        rng = random.Random(self.seed + self.epoch)
        idx = []
        for i, m in enumerate(self.items):
            w = self.weights.get(m["ds"], 1.0)
            idx += [i] * int(w) + ([i] if rng.random() < (w - int(w)) else [])
        idx.sort(key=lambda i: self.items[i]["nf"] + rng.random() * 30)
        batches, cur, mx = [], [], 0
        for i in idx:
            nf = self.items[i]["nf"]
            if cur and max(mx, nf) * (len(cur) + 1) > self.budget:
                batches.append(cur); cur, mx = [], 0
            cur.append(i); mx = max(mx, nf)
        if cur:
            batches.append(cur)
        rng.shuffle(batches)
        return batches

    def __iter__(self):
        b = self._batches()
        self.epoch += 1
        return iter(b)

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