#!/usr/bin/env python3
"""Graft a Whisper encoder onto Emhotob-50M via the stock Qwen2Audio shell.

Qwen2AudioConfig declares `sub_configs = {"audio_config": AutoConfig, "text_config": AutoConfig}`,
which is the same property that let the vision project compose a plain llama text config into
Lfm2Vl with no custom modelling code. So there is nothing to write here beyond the graft itself:
`generate()`, the KV cache, variable-length audio masking and `save_pretrained` all come for free.

Two things make this fit better than the hand-rolled scatter the plan allowed for:

  * `Qwen2AudioEncoder` is structurally a `WhisperEncoder` -- identical conv1/conv2/embed_positions/
    layers/layer_norm names -- so whisper-small's encoder weights load by prefix remap alone.
  * `_get_feat_extract_output_lengths` + per-audio masking means variable-length audio is handled
    natively. Whisper pads every clip to 30s, and without this every 3-second clip would spend
    ~750 tokens on silence.

Frame rate: whisper conv2 (stride 2) gives 50 Hz, then Qwen2Audio's avg_pooler (stride 2) halves it
to 25 Hz. MASC's median 3.0s clip -> ~75 audio tokens, p95 7.8s -> ~195. Well inside 2048.
"""
from __future__ import annotations

import argparse

from pathlib import Path

import torch
from transformers import (AutoConfig, AutoModelForCausalLM, AutoTokenizer, Qwen2AudioConfig,
                          Qwen2AudioEncoderConfig, Qwen2AudioForConditionalGeneration,
                          WhisperConfig, WhisperForConditionalGeneration)

WHISPER = "openai/whisper-small"
LM = "oddadmix/50M-2048-Emhotob"

AUDIO_TOKEN = "<audio>"
AUDIO_START = "<|audio_start|>"
AUDIO_END = "<|audio_end|>"
SPECIALS = [AUDIO_TOKEN, AUDIO_START, AUDIO_END]
PAD_TO_MULTIPLE = 64          # vision precedent: keep the tied embedding matrix 64-aligned


def build_tokenizer(lm_repo: str = LM):
    tok = AutoTokenizer.from_pretrained(lm_repo)
    # Emhotob's 32000-token vocab is fully packed -- there are no reserved slots to steal, so the
    # vocab genuinely has to grow. 512 params per new token, tied to lm_head.
    tok.add_special_tokens({"additional_special_tokens": SPECIALS})
    return tok


def make_projector(d_in: int, d_out: int, kind: str, mult: int = 4):
    """The stock Qwen2Audio projector is a single Linear -- 262K params and no depth is the entire
    audio->text bridge. Whisper's own decoder, which transcribes the same features correctly, uses
    stacked cross-attention. With one matmul the LM receives a weak signal and falls back on its
    text prior: feed the trained model pure silence and it still emits fluent Arabic.

    "mlp" gives it depth and an input LayerNorm (whisper hidden states and Emhotob embeddings live
    at very different scales; the single scalar in calibrate_projector was the only thing bridging
    that before). Qwen2AudioMultiModalProjector.forward just calls self.linear(x), so swapping the
    attribute for a Sequential is transparent to the rest of the stack -- but the state_dict keys
    change, so a repo trained this way no longer round-trips through vanilla from_pretrained."""
    import torch.nn as nn
    if kind == "linear":
        return nn.Linear(d_in, d_out)
    if kind == "mlp":
        h = d_out * mult
        return nn.Sequential(nn.LayerNorm(d_in), nn.Linear(d_in, h), nn.GELU(), nn.Linear(h, d_out))
    raise ValueError(f"unknown projector {kind!r}")


def _last_linear(mod):
    import torch.nn as nn
    if isinstance(mod, nn.Linear):
        return mod
    return [m for m in mod.modules() if isinstance(m, nn.Linear)][-1]


def set_audio_frame_rate(model, pool_stride: int = 2):
    """Set the rate of the audio tokens the LM consumes. 2 = 25 Hz (v1/v2), 1 = 50 Hz (v3).

    Qwen2Audio halves the Whisper encoder's native 50 Hz with a stride-2 average pool right
    before the final layer_norm. That pooler is built in `__init__` from nothing -- it is not
    driven by config and holds no parameters -- so a checkpoint reloaded with vanilla
    transformers ALWAYS comes back at stride 2 regardless of how it was trained. Weights would
    load without complaint and the model would silently run at half the rate it learned.

    Everything here goes through build(), so patching both the module and the length function in
    one place keeps the collator honest too: it derives the number of audio placeholder tokens
    from `_get_feat_extract_output_lengths`, and a disagreement between that and the encoder's
    real output is a shape error at best and silent misalignment at worst.

    The default stays 2 so existing v1/v2 checkpoints keep their trained behaviour; v3 opts in.
    """
    tower = model.model.audio_tower
    tower.avg_pooler = torch.nn.AvgPool1d(pool_stride, stride=pool_stride)

    def _lens(input_lengths):
        input_lengths = (input_lengths - 1) // 2 + 1          # whisper conv2, stride 2
        output_lengths = (input_lengths - pool_stride) // pool_stride + 1
        return input_lengths, output_lengths

    tower._get_feat_extract_output_lengths = _lens
    tower.config.pool_stride = pool_stride                     # recorded into config.json
    return model


def build(whisper_repo: str = WHISPER, lm_repo: str = LM, dtype=torch.float32,
          projector: str = "linear", pool_stride: int = 2, encoder_ckpt: str | None = None):
    """encoder_ckpt: a local encoder trained by train_encoder.py, used INSTEAD of whisper_repo's.

    Its meta.pt carries the WhisperConfig it was trained with, so the audio tower is shaped from
    that rather than from whisper-small -- the 50M encoder is d576, not d768, and the projector
    picks up the narrower input automatically since it is built from the tower's real width."""
    tok = build_tokenizer(lm_repo)
    audio_token_id = tok.convert_tokens_to_ids(AUDIO_TOKEN)

    if encoder_ckpt:
        meta = torch.load(Path(encoder_ckpt).parent / "meta.pt", map_location="cpu",
                          weights_only=False)
        wcfg = WhisperConfig(**meta["config"])
    else:
        wcfg = WhisperConfig.from_pretrained(whisper_repo)
    audio_config = Qwen2AudioEncoderConfig(
        d_model=wcfg.d_model,
        encoder_layers=wcfg.encoder_layers,
        encoder_attention_heads=wcfg.encoder_attention_heads,
        encoder_ffn_dim=wcfg.encoder_ffn_dim,
        num_mel_bins=wcfg.num_mel_bins,
        max_source_positions=wcfg.max_source_positions,
        activation_function=wcfg.activation_function,
        scale_embedding=wcfg.scale_embedding,
    )
    text_config = AutoConfig.from_pretrained(lm_repo)
    # Emhotob ships `use_cache: false` (a pretraining leftover). Left as-is, generate() runs
    # cacheless and Qwen2Audio re-encodes the audio through the 88M tower on EVERY decoded token.
    text_config.use_cache = True
    # VoxtralConfig-style shells inject Mistral defaults incl. rope_theta 1e8. Qwen2Audio does not,
    # but assert rather than trust: a checkpoint using the legacy flat key would inherit silently.
    assert text_config.rope_parameters["rope_theta"] == 10000, text_config.rope_parameters

    cfg = Qwen2AudioConfig(audio_config=audio_config, text_config=text_config,
                           audio_token_index=audio_token_id)
    model = Qwen2AudioForConditionalGeneration(cfg).to(dtype)

    # --- graft the encoder ---------------------------------------------------------------
    # Load donors *through* from_pretrained rather than reading safetensors, so transformers'
    # checkpoint key-conversion map applies -- the same reason the vision graft did it this way.
    if encoder_ckpt:
        # train_encoder.py saves the bare WhisperEncoder state_dict, already in tower key space.
        enc_sd = torch.load(encoder_ckpt, map_location="cpu", weights_only=False)
        enc_sd = {k: v.to(dtype) for k, v in enc_sd.items()}
        whisper = None
    else:
        whisper = WhisperForConditionalGeneration.from_pretrained(whisper_repo, dtype=dtype)
        enc_sd = {k[len("model.encoder."):]: v
                  for k, v in whisper.state_dict().items() if k.startswith("model.encoder.")}
    missing, unexpected = model.model.audio_tower.load_state_dict(enc_sd, strict=False)
    # avg_pooler is parameter-free, so a clean graft leaves nothing missing and nothing unexpected
    assert not unexpected, f"whisper keys the audio tower did not want: {unexpected[:5]}"
    assert not missing, f"audio tower keys whisper did not provide: {missing[:5]}"
    del whisper

    # --- graft the language model --------------------------------------------------------
    lm = AutoModelForCausalLM.from_pretrained(lm_repo, dtype=dtype)
    lm_sd = {k[len("model."):]: v for k, v in lm.state_dict().items() if k.startswith("model.")}
    missing, unexpected = model.model.language_model.load_state_dict(lm_sd, strict=False)
    assert not unexpected, f"LM keys the shell did not want: {unexpected[:5]}"
    assert not missing, f"LM keys Emhotob did not provide: {missing[:5]}"
    # Emhotob ties lm_head to embed_tokens, so the causal head has no separate tensor to copy --
    # initialise it FROM the embeddings but leave it untied. Tying would save 16.4M, but
    # save_pretrained then omits lm_head.weight, and any reload through the stock Qwen2Audio class
    # silently gets a randomly-initialised head. Untied also lets the output head specialise for
    # the ASR distribution, which differs from the LM one.
    model.lm_head.weight.data.copy_(lm.get_input_embeddings().weight.data)
    del lm

    # --- grow the vocab for the audio placeholders ----------------------------------------
    model.resize_token_embeddings(len(tok), pad_to_multiple_of=PAD_TO_MULTIPLE)
    model.config.text_config.vocab_size = model.lm_head.out_features
    model.config.text_config.use_cache = True
    model.generation_config.use_cache = True
    model.generation_config.pad_token_id = tok.pad_token_id
    model.generation_config.bos_token_id = tok.bos_token_id
    model.generation_config.eos_token_id = tok.eos_token_id

    if projector != "linear":
        stock = model.model.multi_modal_projector.linear
        model.model.multi_modal_projector.linear = make_projector(
            stock.in_features, stock.out_features, projector).to(stock.weight.dtype)
    set_audio_frame_rate(model, pool_stride)
    return model, tok


@torch.no_grad()
def calibrate_projector(model, audio_features, feature_lengths=None):
    """Scale the projector so audio tokens land at the same RMS as Emhotob's text embeddings.

    Ported verbatim in spirit from the vision graft, where a freshly-initialised projector emitted
    modality tokens ~36x the magnitude of the token embeddings and alignment simply never started.
    One scalar, measured once on a real batch. If training stalls, check this before the LR.
    """
    tower = model.model.audio_tower
    hidden = tower(audio_features).last_hidden_state
    projected = model.model.multi_modal_projector(hidden)

    target = model.get_input_embeddings().weight.float().pow(2).mean().sqrt()
    measured = projected.float().pow(2).mean().sqrt()
    factor = (target / measured).item()
    proj = _last_linear(model.model.multi_modal_projector.linear)
    proj.weight.data.mul_(factor)
    if proj.bias is not None:
        proj.bias.data.mul_(factor)

    after = model.model.multi_modal_projector(hidden).float().pow(2).mean().sqrt()
    return {"target_rms": target.item(), "before_rms": measured.item(),
            "factor": factor, "after_ratio": (after / target).item()}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--whisper", default=WHISPER)
    ap.add_argument("--lm", default=LM)
    ap.add_argument("--save", default=None)
    ap.add_argument("--projector", default="linear", choices=["linear", "mlp"])
    a = ap.parse_args()

    model, tok = build(a.whisper, a.lm, projector=a.projector)
    n = lambda m: sum(p.numel() for p in m.parameters())
    tower, proj, lang = model.model.audio_tower, model.model.multi_modal_projector, model.model.language_model
    print(f"[+] audio_tower  {n(tower)/1e6:8.2f}M   (frozen in training)")
    print(f"[+] projector    {n(proj)/1e6:8.2f}M   {proj.linear}")
    print(f"[+] language     {n(lang)/1e6:8.2f}M   + lm_head {n(model.lm_head)/1e6:.2f}M")
    print(f"[+] TOTAL        {n(model)/1e6:8.2f}M")
    print(f"[+] vocab {len(tok)} -> embedding {model.get_input_embeddings().weight.shape[0]}"
          f"  audio_token_id={tok.convert_tokens_to_ids(AUDIO_TOKEN)}")

    # 3 seconds of audio through the real feature pipeline, to prove the wiring and the frame math
    from transformers import WhisperFeatureExtractor
    fe = WhisperFeatureExtractor.from_pretrained(a.whisper)
    import numpy as np
    feats = fe([np.zeros(16000 * 3, dtype=np.float32)], sampling_rate=16000, return_tensors="pt")
    stats = calibrate_projector(model, feats.input_features)
    print(f"[+] calibration  target_rms {stats['target_rms']:.4f}  before {stats['before_rms']:.4f}"
          f"  x{stats['factor']:.3f}  -> ratio {stats['after_ratio']:.3f}")
    _, out_len = model.model.audio_tower._get_feat_extract_output_lengths(torch.tensor([16000 * 3 // 160]))
    print(f"[+] 3.0s audio -> {int(out_len[0])} audio tokens  ({int(out_len[0])/3.0:.1f} tokens/sec)")

    if a.save:
        model.save_pretrained(a.save); tok.save_pretrained(a.save)
        print(f"[+] saved -> {a.save}")


if __name__ == "__main__":
    main()