"""组合模型 9B 版 (导出包 vendored): Qwen3-Omni AuT(2048) + connector(2048→4096 投影) + Qwen3.5-9B (early fusion)。

与 audio-trans work/audio_model_9b.py 逻辑一致, 路径全部参数化
(llm_path / omni_path 由调用方传入, 不再硬编码 $AT)。

架构: wav → WhisperFeatureExtractor(128 mel) → AuT(2048) → connector(2048→4096) →
替换 <|audio_pad|> 位。connector = 列正交 W × gain (EMBED_RMS_9B/AUT_RMS/√(2048/4096),
输出 per-elem RMS ≈ 9B embed 行 RMS 0.0137); log_alpha 已删 (2026-09-11 极简口径)。

类名保持 Qwen35AudioForCausalLM 与 2B 包一致, pipeline 无需区分 2B/9B。
"""
import json
import os

import torch
import torch.nn as nn

AUDIO_PAD = 248076  # <|audio_pad|>
AUDIO_DIM = 2048    # AuT output_dim
EMBED_RMS_9B = 0.0137   # 实测 9B 基座 embed 行均值
AUT_RMS = 0.0143        # 实测 AuT 输出帧 RMS


class AudioConnector9B(nn.Module):
    """h' = W·h; W: 2048→4096 列正交 × gain。"""

    def __init__(self, in_dim=AUDIO_DIM, out_dim=4096,
                 gain=EMBED_RMS_9B / AUT_RMS / (AUDIO_DIM / 4096) ** 0.5):
        super().__init__()
        self.proj = nn.Linear(in_dim, out_dim, bias=False)
        nn.init.orthogonal_(self.proj.weight)  # out>in → 列正交, 保范数
        self.proj.weight.data *= gain

    def forward(self, h):
        return self.proj(h)


def load_audio_tower(omni_path, dtype=torch.bfloat16):
    """从 Qwen3-Omni shard 1 加载 AuT(525 keys)。omni_path: 含 config.json +
    model-00001-of-00015.safetensors 的目录 (导出包 stlm_omni/)。"""
    from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
        Qwen3OmniMoeAudioEncoder, Qwen3OmniMoeAudioEncoderConfig)
    from safetensors.torch import load_file
    cfg = json.load(open(os.path.join(omni_path, 'config.json')))
    audio_cfg = Qwen3OmniMoeAudioEncoderConfig(**{k: v for k, v in
                                                  cfg['thinker_config']['audio_config'].items()
                                                  if not k.startswith('_')})
    enc = Qwen3OmniMoeAudioEncoder(audio_cfg).eval()
    sd = load_file(os.path.join(omni_path, 'model-00001-of-00015.safetensors'))
    sd = {k.replace('thinker.audio_tower.', ''): v for k, v in sd.items()
          if k.startswith('thinker.audio_tower.')}
    missing, unexpected = enc.load_state_dict(sd, strict=False)
    assert not missing and not unexpected, f'missing={missing[:3]} unexpected={unexpected[:3]}'
    return enc.to(dtype)


def load_feature_extractor():
    from transformers import WhisperFeatureExtractor
    fe = WhisperFeatureExtractor(feature_size=128, hop_length=160, n_fft=400,
                                 sampling_rate=16000, padding_value=0.0,
                                 return_attention_mask=True)
    # 构造器忽略 n_samples; 必须事后设(默认 480000=30s 会截断 60s 语料)
    fe.n_samples = 4800000
    fe.nb_max_frames = 30000
    return fe


class Qwen35AudioForCausalLM(nn.Module):
    """AuT + connector + Qwen3.5-9B, audio embedding 原位替换 <|audio_pad|>。

    (类名与 2B 包相同; 9B 差异: LLM 4096h/32L untied, connector 为 2048→4096 投影)
    """

    def __init__(self, llm_path, omni_path, dtype=torch.bfloat16,
                 freeze_llm=False, freeze_audio=False):
        super().__init__()
        from transformers import AutoModelForCausalLM, AutoTokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(llm_path)
        self.llm = AutoModelForCausalLM.from_pretrained(llm_path, dtype=dtype)
        self.audio_tower = load_audio_tower(omni_path, dtype)
        self.fe = load_feature_extractor()
        self.connector = AudioConnector9B(out_dim=self.llm.config.hidden_size).to(dtype)
        assert self.llm.config.hidden_size == 4096
        assert self.llm.config.tie_word_embeddings is False
        if freeze_llm:
            for p in self.llm.parameters():
                p.requires_grad_(False)
        if freeze_audio:
            for p in self.audio_tower.parameters():
                p.requires_grad_(False)

    @property
    def device(self):
        return next(self.llm.parameters()).device

    def encode_audio(self, wav_paths):
        """wav paths → list of (T_i, 4096) connector 后 audio embeddings + 各段帧数。"""
        import librosa
        embs = []
        dt = next(self.audio_tower.parameters()).dtype
        for p in wav_paths:
            wav, _ = librosa.load(p, sr=16000)
            f = self.fe(wav, sampling_rate=16000, return_tensors='pt',
                        return_attention_mask=True)
            T = int(f.attention_mask.sum(-1)[0])
            feats = f.input_features[0][:, :T]
            out = self.audio_tower(input_features=feats.to(self.device, dt),
                                   feature_lens=torch.tensor([T], device=self.device))
            h = out.last_hidden_state if hasattr(out, 'last_hidden_state') else out[0]
            embs.append(self.connector(h.to(self.device, dt)))
        return embs

    def forward(self, input_ids, attention_mask=None, labels=None,
                audio_embs=None, **kwargs):
        embed = self.llm.get_input_embeddings()
        inputs_embeds = embed(input_ids)
        if audio_embs is not None:
            pad_mask = input_ids == AUDIO_PAD
            n_pad = int(pad_mask.sum())
            n_audio = sum(e.shape[0] for e in audio_embs)
            assert n_pad == n_audio, f'pad tokens {n_pad} != audio frames {n_audio}'
            inputs_embeds = inputs_embeds.clone()
            inputs_embeds[pad_mask] = torch.cat(audio_embs).to(inputs_embeds.dtype)
        return self.llm(inputs_embeds=inputs_embeds, attention_mask=attention_mask,
                        labels=labels, **kwargs)