"""oruk-loop-153M: a double-looped Qwen3-ASR student.

The Qwen3-ASR architecture is kept, but each of its two transformer stacks is replaced by ONE shared block that is
applied repeatedly:

* audio encoder: one block applied 24 times, made depth-aware by a rank-64 adapter on every linear layer plus
  per-depth biases and LayerNorm gains;
* text decoder: one block applied 28 times, with rank-16 adapters and per-depth RMSNorm gains;
* token table: the kept English vocabulary (36,974 rows); 368 frequent rows are stored exactly and the rest are
  factorised at rank 729. The same table is used for input embeddings and the (tied) output layer.

Token ids stay in the original Qwen3 id space, so the stock Qwen3-ASR processor (tokenizer, chat template,
feature extractor) works unchanged.

Usage:
    from transformers import AutoModel, AutoProcessor
    model = AutoModel.from_pretrained("oruk/oruk-loop-153M", trust_remote_code=True)
    processor = AutoProcessor.from_pretrained("oruk/oruk-loop-153M")
    print(model.transcribe("speech.wav", processor))
"""
from __future__ import annotations

import copy
import math
import os
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
from transformers import AutoConfig, GenerationConfig, Qwen3ASRForConditionalGeneration

SAMPLE_RATE = 16000
PAD_ID = 151645
EOS_IDS = (151643, 151645)


# ---------------------------------------------------------------------------------------------- looped modules
class DepthCursor:
    """Which virtual depth the shared block is computing (set by each call site before the block runs)."""

    def __init__(self, n: int):
        self.n = n
        self.depth = 0


class LoopLinear(nn.Module):
    """y = W x + b[d] + B[d] (A[d] x): one shared weight, per-depth bias and per-depth rank-r adapter."""

    def __init__(self, cursor: DepthCursor, n: int, in_features: int, out_features: int, rank: int, bias: bool):
        super().__init__()
        object.__setattr__(self, "cursor", cursor)
        self.in_features, self.out_features, self.rank = in_features, out_features, rank
        self.weight = nn.Parameter(torch.zeros(out_features, in_features))
        self.bias = nn.Parameter(torch.zeros(n, out_features)) if bias else None
        if rank > 0:
            self.lora_A = nn.Parameter(torch.zeros(n, rank, in_features))
            self.lora_B = nn.Parameter(torch.zeros(n, out_features, rank))
        object.__setattr__(self, "_merged", None)
        object.__setattr__(self, "_merged_bias", None)

    @torch.no_grad()
    def merge(self, dtype=torch.bfloat16):
        """Materialise W + B[d] A[d] for every depth (faster decoding; ~3.5 GB in bf16 for the whole model)."""
        n = self.lora_A.shape[0] if self.rank > 0 else self.cursor.n
        W = torch.empty((n, self.out_features, self.in_features), dtype=dtype, device=self.weight.device)
        for i in range(n):
            w = self.weight.float()
            if self.rank > 0:
                w = w + self.lora_B[i].float() @ self.lora_A[i].float()
            W[i] = w.to(dtype)
        object.__setattr__(self, "_merged", W)
        object.__setattr__(self, "_merged_bias", None if self.bias is None else self.bias.detach().to(dtype))

    def unmerge(self):
        object.__setattr__(self, "_merged", None)
        object.__setattr__(self, "_merged_bias", None)

    def forward(self, x):
        d = self.cursor.depth
        if self._merged is not None:
            b = None if self._merged_bias is None else self._merged_bias[d]
            return F.linear(x.to(self._merged.dtype), self._merged[d], b)
        y = F.linear(x, self.weight, None if self.bias is None else self.bias[d])
        if self.rank > 0:
            y = y + F.linear(F.linear(x, self.lora_A[d]), self.lora_B[d])
        return y


class LoopLayerNorm(nn.Module):
    def __init__(self, cursor, n, shape, eps):
        super().__init__()
        object.__setattr__(self, "cursor", cursor)
        self.shape, self.eps = tuple(shape), eps
        self.weight = nn.Parameter(torch.ones(n, *self.shape))
        self.bias = nn.Parameter(torch.zeros(n, *self.shape))

    def forward(self, x):
        d = self.cursor.depth
        w, b = self.weight[d], self.bias[d]
        if not torch.is_autocast_enabled(x.device.type):  # under autocast layer_norm runs in fp32 with fp32 gains
            w, b = w.to(x.dtype), b.to(x.dtype)
        return F.layer_norm(x, self.shape, w, b, self.eps)


class LoopRMSNorm(nn.Module):
    def __init__(self, cursor, n, dim, eps):
        super().__init__()
        object.__setattr__(self, "cursor", cursor)
        self.weight = nn.Parameter(torch.ones(n, dim))
        self.variance_epsilon = eps

    def forward(self, x):
        d = self.cursor.depth
        dtype = x.dtype
        h = x.to(torch.float32)
        h = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + self.variance_epsilon)
        return (self.weight[d] * h.to(dtype)).to(dtype)


def build_loop_block(template: nn.Module, cursor: DepthCursor, n: int, rank: int) -> nn.Module:
    """Turn one transformer layer into the shared looped block (same module paths, depth-aware leaves)."""
    block = copy.deepcopy(template)
    for name, mod in list(block.named_modules()):
        if not name:
            continue
        if isinstance(mod, nn.Linear):
            new = LoopLinear(cursor, n, mod.in_features, mod.out_features, rank, mod.bias is not None)
        elif isinstance(mod, nn.LayerNorm):
            new = LoopLayerNorm(cursor, n, mod.normalized_shape, mod.eps)
        elif type(mod).__name__.endswith("RMSNorm"):
            new = LoopRMSNorm(cursor, n, mod.weight.shape[0], mod.variance_epsilon)
        else:
            continue
        parent_name, _, child = name.rpartition(".")
        setattr(block.get_submodule(parent_name) if parent_name else block, child, new)
    return block


class LoopCall(nn.Module):
    """Parameter-free call site for virtual depth `index` of a shared block (one per original layer)."""

    def __init__(self, block, cursor, index, kv_slot=None):
        super().__init__()
        object.__setattr__(self, "_block", block)
        object.__setattr__(self, "_cursor", cursor)
        self.index, self.kv_slot = index, kv_slot

    def forward(self, hidden_states, *args, **kwargs):
        self._cursor.depth = self.index
        if self.kv_slot is not None:  # each virtual decoder depth keeps its own KV-cache slot
            self._block.self_attn.layer_idx = self.kv_slot
        return self._block(hidden_states, *args, **kwargs)


# ---------------------------------------------------------------------------------------------- token table
class TokenTable(nn.Module):
    """Tied table over the kept vocabulary: exact frequent rows + rank-r factorised rare rows (U @ P)."""

    def __init__(self, n_full, dim, n_kept, n_exact, rank):
        super().__init__()
        self.n_full, self.embedding_dim, self.rank = n_full, dim, rank
        self.register_buffer("keep", torch.zeros(n_kept, dtype=torch.long))
        self.register_buffer("lut", torch.zeros(n_full, dtype=torch.long))
        self.register_buffer("exact_idx", torch.zeros(n_exact, dtype=torch.long))
        self.register_buffer("rare_idx", torch.zeros(n_kept - n_exact, dtype=torch.long))
        self.exact_rows = nn.Parameter(torch.zeros(n_exact, dim, dtype=torch.bfloat16), requires_grad=False)
        self.U = nn.Parameter(torch.zeros(n_kept - n_exact, rank))
        self.P = nn.Parameter(torch.zeros(rank, dim))
        object.__setattr__(self, "_merged", None)

    def weight(self, dtype=None):
        dtype = dtype or self.U.dtype
        if self._merged is not None:
            return self._merged if self._merged.dtype == dtype else self._merged.to(dtype)
        W = torch.empty(self.keep.shape[0], self.embedding_dim, dtype=dtype, device=self.U.device)
        W = W.index_copy(0, self.exact_idx, self.exact_rows.to(dtype))
        return W.index_copy(0, self.rare_idx, (self.U.float() @ self.P.float()).to(dtype))

    @torch.no_grad()
    def merge(self, dtype=torch.bfloat16):
        object.__setattr__(self, "_merged", None)
        object.__setattr__(self, "_merged", self.weight(dtype).detach())

    def forward(self, ids):
        # read the merged table in its own dtype (bf16 on GPU), as evaluated; never re-cast it per call
        return F.embedding(self.lut[ids], self._merged if self._merged is not None else self.weight())


class TableEmbedding(nn.Module):
    def __init__(self, table):
        super().__init__()
        object.__setattr__(self, "table", table)
        self.embedding_dim = table.embedding_dim

    @property
    def num_embeddings(self):
        return self.table.n_full

    @property
    def weight(self):
        return self.table.weight()

    def forward(self, ids):
        return self.table(ids)


class TableHead(nn.Module):
    """Output layer tied to the token table; full-width logits with -inf at ids outside the kept vocabulary."""

    def __init__(self, table):
        super().__init__()
        object.__setattr__(self, "table", table)
        self.in_features, self.out_features = table.embedding_dim, table.n_full

    @property
    def weight(self):
        return self.table.weight()

    def forward(self, h):
        kept = F.linear(h, self.table.weight(h.dtype))
        out = torch.full(h.shape[:-1] + (self.table.n_full,), float("-inf"), dtype=kept.dtype, device=kept.device)
        return out.index_copy_(-1, self.table.keep, kept)


# ---------------------------------------------------------------------------------------------- audio helpers
def load_audio(source, sampling_rate=None):
    """path / (array, sr) / 1-D array -> float32 mono at 16 kHz."""
    if isinstance(source, (str, os.PathLike)):
        import soundfile as sf
        wav, sr = sf.read(str(source), dtype="float32", always_2d=False)
    elif isinstance(source, tuple):
        wav, sr = source
    else:
        wav, sr = source, sampling_rate or SAMPLE_RATE
    wav = np.asarray(wav, dtype=np.float32)
    if wav.ndim > 1:
        wav = wav.mean(axis=1 if wav.shape[0] > wav.shape[1] else 0)
    if sr != SAMPLE_RATE:
        try:
            import soxr
            wav = soxr.resample(wav, sr, SAMPLE_RATE, quality="HQ").astype(np.float32)
        except ImportError:
            from math import gcd
            from scipy.signal import resample_poly
            g = gcd(int(sr), SAMPLE_RATE)
            wav = resample_poly(wav, SAMPLE_RATE // g, int(sr) // g).astype(np.float32)
    # the evaluation pipeline stores 16-bit PCM; quantise the same way so results reproduce exactly
    return np.clip(np.rint(wav * 32767.0), -32768, 32767).astype(np.float32) / 32768.0


def split_points(w, sr=SAMPLE_RATE, max_s=30.0, search_s=5.0, frame=0.1):
    """Chunk boundaries for audio longer than max_s: each cut sits at the quietest 100 ms frame within the last
    search_s seconds before the limit, so words are rarely split."""
    n, cuts, start = len(w), [], 0
    L, S, Fr = int(max_s * sr), int(search_s * sr), int(frame * sr)
    while n - start > L:
        lo, hi = start + L - S, start + L
        seg = w[lo:hi].astype(np.float32)
        k = len(seg) // Fr
        e = (seg[:k * Fr].reshape(k, Fr) ** 2).mean(1)
        cut = lo + int(np.argmin(e)) * Fr + Fr // 2 if k else hi
        cuts.append(cut)
        start = cut
    return [0] + cuts + [n]


# ---------------------------------------------------------------------------------------------- model
class OrukLoopForConditionalGeneration(Qwen3ASRForConditionalGeneration):
    """Qwen3-ASR with both stacks looped. `from_pretrained` loads the 153M stored parameters; see `transcribe`."""

    _tied_weights_keys = {}

    def __init__(self, config):
        loop = dict(getattr(config, "oruk_loop", {}))
        tc, ac = config.text_config, config.audio_config
        n_dec, n_enc, vocab = tc.num_hidden_layers, ac.encoder_layers, tc.vocab_size
        types = list(tc.layer_types)
        # build the stock skeleton with ONE template layer per stack and a tiny vocabulary (cheap), then restore
        tc.num_hidden_layers, tc.layer_types, tc.vocab_size, ac.encoder_layers = 1, types[:1], 8, 1
        try:
            super().__init__(config)
        finally:
            tc.num_hidden_layers, tc.layer_types, tc.vocab_size, ac.encoder_layers = n_dec, types, vocab, n_enc
        at, lm = self.model.audio_tower, self.model.language_model
        for c in {id(at.config): at.config}.values():  # sub-modules may hold copies of the sub-configs
            c.encoder_layers = n_enc
        lm.config.num_hidden_layers, lm.config.layer_types, lm.config.vocab_size = n_dec, types, vocab

        self.enc_cursor, self.dec_cursor = DepthCursor(n_enc), DepthCursor(n_dec)
        at.loop = build_loop_block(at.layers[0], self.enc_cursor, n_enc, int(loop.get("encoder_rank", 64)))
        at.layers = nn.ModuleList(LoopCall(at.loop, self.enc_cursor, i) for i in range(n_enc))
        lm.loop = build_loop_block(lm.layers[0], self.dec_cursor, n_dec, int(loop.get("decoder_rank", 16)))
        lm.layers = nn.ModuleList(LoopCall(lm.loop, self.dec_cursor, i, kv_slot=i) for i in range(n_dec))
        lm.token_table = TokenTable(vocab, tc.hidden_size, int(loop.get("table_rows", 36974)),
                                    int(loop.get("table_exact", 368)), int(loop.get("table_rank", 729)))
        lm.embed_tokens = TableEmbedding(lm.token_table)
        self.lm_head = TableHead(lm.token_table)
        self._source = None

    # the table is tied by construction; never let transformers re-tie or resize it
    def tie_weights(self, *args, **kwargs):
        return None

    def tie_embeddings_and_encoder_decoder(self, *args, **kwargs):
        return None

    def num_stored_parameters(self):
        return sum(p.numel() for p in self.parameters())

    def merge_loops(self, dtype=torch.bfloat16):
        for m in self.modules():
            if isinstance(m, LoopLinear):
                m.merge(dtype)

    def unmerge_loops(self):
        for m in self.modules():
            if isinstance(m, LoopLinear):
                m.unmerge()

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, config=None, dtype=None, torch_dtype=None,
                        device=None, device_map=None, merge=None, revision=None, cache_dir=None, token=None,
                        local_files_only=False, **kwargs):
        """Load the release weights.

        device: "cuda", "cpu", ... (default: cuda if available). dtype: compute dtype (default bf16 on GPU, fp32 on CPU).
        merge: materialise per-depth weights for fast decoding (default: True on GPU, False on CPU).
        """
        ignored = [k for k in kwargs if not k.startswith("_") and k not in ("trust_remote_code", "low_cpu_mem_usage",
                                                                            "attn_implementation", "use_safetensors")]
        if ignored:
            import warnings
            warnings.warn(f"oruk-loop from_pretrained ignores these arguments: {ignored}")
        src = str(pretrained_model_name_or_path)
        if os.path.isdir(src):
            path = Path(src)
        else:
            from huggingface_hub import snapshot_download
            path = Path(snapshot_download(src, revision=revision, cache_dir=cache_dir, token=token,
                                          local_files_only=local_files_only,
                                          allow_patterns=["*.json", "*.safetensors", "*.py", "*.jinja", "*.txt"]))
        if config is None:
            config = AutoConfig.from_pretrained(str(path))
        if device is None:
            device = device_map if isinstance(device_map, str) and device_map != "auto" else None
        device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
        dtype = dtype or torch_dtype
        if isinstance(dtype, str) and dtype != "auto":
            dtype = getattr(torch, dtype)
        if dtype is None or dtype == "auto":
            dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
        if merge is None:
            merge = device.type == "cuda"

        model = cls(config)
        from safetensors.torch import load_file
        state = load_file(str(path / "model.safetensors"))
        missing, unexpected = model.load_state_dict(state, strict=False)
        if missing or unexpected:
            raise RuntimeError(f"oruk-loop weights do not match the model: missing={missing[:8]} "
                               f"unexpected={unexpected[:8]}")
        del state
        try:
            model.generation_config = GenerationConfig.from_pretrained(str(path))
        except Exception:
            pass
        model.to(device)
        model.model.language_model.token_table.merge(dtype)
        if merge:
            model.merge_loops(dtype)
        if dtype == torch.bfloat16 and device.type == "cuda":
            # evaluated precision: modules kept from Qwen3-ASR in bf16, looped weights merged to bf16, per-depth norms
            # and biases in fp32, everything run under bf16 autocast
            at, lm = model.model.audio_tower, model.model.language_model
            for mod in (at.conv2d1, at.conv2d2, at.conv2d3, at.conv_out, at.ln_post, model.model.multi_modal_projector,
                        lm.norm):
                mod.to(torch.bfloat16)
            model._autocast = torch.bfloat16
        elif dtype != torch.float32:
            model.to(dtype)
        model.eval()
        model._source = str(path)
        return model

    _autocast = None

    def generate(self, *args, **kwargs):
        if self._autocast is not None and not torch.is_autocast_enabled(self.device.type):
            with torch.autocast(self.device.type, dtype=self._autocast):
                return super().generate(*args, **kwargs)
        return super().generate(*args, **kwargs)

    # ------------------------------------------------------------------------------------------ transcription
    @torch.no_grad()
    def transcribe(self, audio, processor=None, *, sampling_rate=None, batch_size=16, num_beams=1, language="English",
                   length_penalty=1.0, return_chunks=False):
        """Transcribe one recording or a list of recordings (paths, 1-D arrays at `sampling_rate`, or (array, sr)).

        Audio longer than 30 s is cut at quiet points into <=30 s chunks; chunk texts are joined with spaces.
        Returns a string (or a list of strings for a list input).
        """
        if processor is None:
            from transformers import AutoProcessor
            processor = AutoProcessor.from_pretrained(self._source or "oruk/oruk-loop-153M")
        single = not isinstance(audio, list)
        clips = [load_audio(a, sampling_rate) for a in ([audio] if single else audio)]
        items, owners = [], []
        for w in clips:
            b = split_points(w)
            ids = []
            for a, z in zip(b[:-1], b[1:]):
                piece = w[a:z]
                if len(piece) < 1600:  # <0.1 s tail
                    continue
                enc = processor.apply_transcription_request(audio=piece, language=language, return_tensors="pt",
                                                            processor_kwargs={"audio_kwargs": {"sampling_rate": SAMPLE_RATE}})
                ids.append(len(items))
                items.append({"input_ids": enc["input_ids"][0], "input_features": enc["input_features"][0].float(),
                              "input_features_mask": enc["input_features_mask"][0], "dur": len(piece) / SAMPLE_RATE})
            owners.append(ids)
        texts = [""] * len(items)
        order = sorted(range(len(items)), key=lambda i: items[i]["dur"])
        dev, dt = self.device, next(p for p in self.model.audio_tower.conv2d1.parameters()).dtype
        for k in range(0, len(order), batch_size):
            idx = order[k:k + batch_size]
            batch = self._collate([items[i] for i in idx], dev, dt)
            max_new = int(min(448, math.ceil(10 * max(items[i]["dur"] for i in idx)) + 24))
            out = self.generate(**batch, do_sample=False, num_beams=num_beams, length_penalty=length_penalty,
                                max_new_tokens=max_new, use_cache=True)
            cont = out[:, batch["input_ids"].shape[1]:]
            for j, i in enumerate(idx):
                toks = cont[j].tolist()
                for e_pos, t in enumerate(toks):
                    if t in EOS_IDS:
                        toks = toks[:e_pos + 1]
                        break
                texts[i] = processor.decode(toks, return_format="transcription_only").strip()
        results = [[texts[i] for i in own] for own in owners]
        joined = [" ".join(t for t in r if t) for r in results]
        if return_chunks:
            return (results[0] if single else results), (joined[0] if single else joined)
        return joined[0] if single else joined

    @staticmethod
    def _collate(items, device, dtype):
        B = len(items)
        P = max(len(it["input_ids"]) for it in items)
        Fm = max(it["input_features"].shape[-1] for it in items)
        Fm = (Fm + 99) // 100 * 100
        ids = torch.full((B, P), PAD_ID, dtype=torch.long)
        am = torch.zeros((B, P), dtype=torch.long)
        feats = torch.zeros((B, items[0]["input_features"].shape[0], Fm), dtype=torch.float32)
        fm = torch.zeros((B, Fm), dtype=torch.long)
        for b, it in enumerate(items):
            n = len(it["input_ids"])
            ids[b, P - n:] = it["input_ids"]
            am[b, P - n:] = 1
            f = it["input_features"].shape[-1]
            feats[b, :, :f] = it["input_features"]
            fm[b, :f] = it["input_features_mask"]
        return {"input_ids": ids.to(device), "attention_mask": am.to(device),
                "input_features": feats.to(device, dtype=dtype), "input_features_mask": fm.to(device)}