"""
Export Cohere Transcribe to ONNX with feature extraction BAKED INTO the encoder.

The encoder takes raw audio waveform (batch, samples) and internally computes:
  preemphasis -> STFT -> mel filterbank -> log -> per-feature normalization
  -> Conformer encoder -> projection -> cross K/V pre-computation

This means NO external feature extraction is needed at inference time.

Requirements:
    pip install torch transformers safetensors sentencepiece onnx onnxruntime librosa soundfile protobuf onnxscript

Usage:
    python export_onnx_baked.py --model-dir . --output-dir ./onnx-baked
"""

import argparse
import math
import os

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


# ---------------------------------------------------------------------------
# ONNX-compatible STFT via Conv1d
#
# torch.stft doesn't export via TorchScript ONNX. We implement STFT as a
# Conv1d with DFT basis filters, which is fully traceable.
# ---------------------------------------------------------------------------

class OnnxSTFT(nn.Module):
    """STFT implemented as Conv1d with DFT basis -- ONNX-exportable."""

    def __init__(self, n_fft=512, hop_length=160, win_length=400):
        super().__init__()
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.win_length = win_length
        self.pad_amount = n_fft // 2  # center=True padding

        # Build DFT basis as Conv1d filters
        n_freqs = n_fft // 2 + 1
        window = torch.hann_window(win_length, periodic=False)

        # Pad window to n_fft if needed
        if win_length < n_fft:
            left_pad = (n_fft - win_length) // 2
            right_pad = n_fft - win_length - left_pad
            window = F.pad(window, (left_pad, right_pad))

        # DFT basis: real (cosine) and imaginary (negative sine)
        n_range = torch.arange(0, n_fft, dtype=torch.float32)
        k_range = torch.arange(0, n_freqs, dtype=torch.float32)
        # angles: (n_freqs, n_fft)
        angles = 2.0 * math.pi * k_range.unsqueeze(1) * n_range.unsqueeze(0) / n_fft

        real_filters = torch.cos(angles) * window.unsqueeze(0)
        imag_filters = -torch.sin(angles) * window.unsqueeze(0)

        # Stack: (2 * n_freqs, 1, n_fft) for Conv1d
        filters = torch.cat([real_filters, imag_filters], dim=0).unsqueeze(1)
        self.register_buffer("filters", filters)

    def forward(self, x):
        # x: (batch, samples)
        # Pad for center=True
        x = F.pad(x, (self.pad_amount, self.pad_amount), mode="constant", value=0.0)
        # Conv1d: (batch, 1, samples) -> (batch, 2*n_freqs, frames)
        x = x.unsqueeze(1)
        out = F.conv1d(x, self.filters, stride=self.hop_length)
        # Split into real and imaginary, compute power spectrum
        n_freqs = self.n_fft // 2 + 1
        real_part = out[:, :n_freqs, :]
        imag_part = out[:, n_freqs:, :]
        power = real_part.pow(2) + imag_part.pow(2)
        return power  # (batch, n_freqs, frames)


# ---------------------------------------------------------------------------
# Baked feature extraction: raw audio -> log-mel features
# Matches Cohere's FilterbankFeatures exactly (minus dither, which is ~0 impact)
# ---------------------------------------------------------------------------

class BakedFeatureExtractor(nn.Module):
    """ONNX-traceable mel feature extraction matching Cohere's pipeline."""

    def __init__(self, fb_weights, sample_rate=16000, n_fft=512, win_length=400,
                 hop_length=160, n_mels=128, preemph=0.97, log_zero_guard=2**-24):
        super().__init__()
        self.preemph = preemph
        self.log_zero_guard = log_zero_guard
        self.hop_length = hop_length

        # ONNX-compatible STFT
        self.stft = OnnxSTFT(n_fft=n_fft, hop_length=hop_length, win_length=win_length)

        # Mel filterbank: (1, n_mels, n_fft//2+1) from checkpoint
        if fb_weights.dim() == 2:
            fb_weights = fb_weights.unsqueeze(0)
        self.register_buffer("fb", fb_weights.float())

    def forward(self, audio):
        # audio: (batch, samples) float32

        # 1. Pre-emphasis: y[n] = x[n] - 0.97 * x[n-1]
        x = torch.cat([audio[:, :1], audio[:, 1:] - self.preemph * audio[:, :-1]], dim=1)

        # 2. STFT -> power spectrum: (batch, n_freqs, frames)
        power = self.stft(x)

        # 3. Mel filterbank: (batch, n_mels, frames)
        mel = torch.matmul(self.fb, power)

        # 4. Log with guard value
        mel = torch.log(mel + self.log_zero_guard)

        # 5. Per-feature normalization (mean=0, std=1 per mel bin)
        mean = mel.mean(dim=2, keepdim=True)
        var = ((mel - mean) ** 2).mean(dim=2, keepdim=True)
        n_frames = mel.shape[2]
        # Unbiased std (ddof=1) to match original + dither constant for stability
        std = torch.sqrt(var * n_frames / (n_frames - 1.0) + 1e-5)
        mel = (mel - mean) / std

        return mel  # (batch, n_mels=128, frames)


# ---------------------------------------------------------------------------
# Full baked encoder: raw audio -> cross K/V
# ---------------------------------------------------------------------------

class BakedEncoderForONNX(nn.Module):
    """
    Full pipeline: raw audio -> mel features -> Conformer encoder -> cross K/V.

    Input:  audio (batch, samples) float32 at 16kHz
    Output: n_layer_cross_k (n_dec_layers, batch, T', hidden_size)
            n_layer_cross_v (n_dec_layers, batch, T', hidden_size)
    """

    def __init__(self, model, fb_weights):
        super().__init__()
        self.features = BakedFeatureExtractor(fb_weights=fb_weights)
        self.encoder = model.encoder
        self.encoder_decoder_proj = model.encoder_decoder_proj

        dec_layers = model.transf_decoder._decoder.layers
        self.cross_k_projs = nn.ModuleList([layer.second_sub_layer.key_net for layer in dec_layers])
        self.cross_v_projs = nn.ModuleList([layer.second_sub_layer.value_net for layer in dec_layers])

    def forward(self, audio):
        # audio: (batch, samples) float32

        # 1. Extract mel features: (batch, 128, frames)
        mel = self.features(audio)

        # 2. Run Conformer encoder
        batch = mel.shape[0]
        length = torch.full((batch,), mel.shape[2], dtype=torch.long, device=mel.device)
        enc_out, _enc_len = self.encoder(mel, length)

        # 3. Project encoder hidden to decoder hidden size
        if self.encoder_decoder_proj is not None:
            enc_out = self.encoder_decoder_proj(enc_out)

        # 4. Pre-compute cross-attention K/V per decoder layer
        cross_k_list = []
        cross_v_list = []
        for k_proj, v_proj in zip(self.cross_k_projs, self.cross_v_projs):
            cross_k_list.append(k_proj(enc_out))
            cross_v_list.append(v_proj(enc_out))

        n_layer_cross_k = torch.stack(cross_k_list, dim=0)
        n_layer_cross_v = torch.stack(cross_v_list, dim=0)

        return n_layer_cross_k, n_layer_cross_v


# ---------------------------------------------------------------------------
# Import decoder + tokens from original export script
# ---------------------------------------------------------------------------

from export_onnx import DecoderForONNX, generate_tokens_txt


# ---------------------------------------------------------------------------
# Export
# ---------------------------------------------------------------------------

def export_baked_encoder(model, fb_weights, output_dir, opset_version=17):
    """Export baked encoder (audio -> cross K/V) to ONNX."""
    print("Wrapping baked encoder...")
    encoder = BakedEncoderForONNX(model, fb_weights)
    encoder.eval()

    # Dummy: 10 seconds of 16kHz audio
    dummy_audio = torch.randn(1, 160000, dtype=torch.float32)

    output_path = os.path.join(output_dir, "cohere-encoder.onnx")
    print(f"Exporting baked encoder to {output_path}...")

    with torch.no_grad():
        torch.onnx.export(
            encoder,
            (dummy_audio,),
            output_path,
            input_names=["audio"],
            output_names=["n_layer_cross_k", "n_layer_cross_v"],
            dynamic_axes={
                "audio": {0: "batch", 1: "samples"},
                "n_layer_cross_k": {1: "batch", 2: "T_enc"},
                "n_layer_cross_v": {1: "batch", 2: "T_enc"},
            },
            opset_version=opset_version,
            do_constant_folding=True,
            dynamo=False,
        )
    print(f"Baked encoder exported: {output_path}")
    return output_path


def main():
    parser = argparse.ArgumentParser(description="Export Cohere Transcribe to ONNX (baked features)")
    parser.add_argument("--model-dir", type=str, default=".", help="Directory with model files")
    parser.add_argument("--output-dir", type=str, default="./onnx-baked", help="Output directory")
    parser.add_argument("--opset", type=int, default=17, help="ONNX opset version")
    parser.add_argument("--skip-encoder", action="store_true")
    parser.add_argument("--skip-decoder", action="store_true")
    parser.add_argument("--skip-tokens", action="store_true")
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    if not args.skip_tokens:
        print("=== Generating tokens.txt ===")
        generate_tokens_txt(args.model_dir, os.path.join(args.output_dir, "tokens.txt"))

    if args.skip_encoder and args.skip_decoder:
        print("Done.")
        return

    # Load mel filterbank from checkpoint
    print("=== Loading filterbank weights from checkpoint ===")
    from safetensors.torch import load_file
    state = load_file(os.path.join(args.model_dir, "model.safetensors"))
    fb_weights = state["preprocessor.featurizer.fb"].float()
    print(f"Filterbank shape: {fb_weights.shape}")

    # Load model
    print("=== Loading model ===")
    from transformers import AutoModelForSpeechSeq2Seq
    model = AutoModelForSpeechSeq2Seq.from_pretrained(
        args.model_dir,
        trust_remote_code=True,
        dtype=torch.bfloat16,
    )
    model.eval()
    model = model.float()

    if not args.skip_encoder:
        print("\n=== Exporting Baked Encoder ===")
        export_baked_encoder(model, fb_weights, args.output_dir, args.opset)

    if not args.skip_decoder:
        print("\n=== Exporting Decoder ===")
        from export_onnx import export_decoder
        export_decoder(model, args.output_dir, args.opset)

    print("\n=== Done ===")
    print(f"Files in: {args.output_dir}/")
    print("Encoder input: raw audio waveform (batch, samples) float32 @ 16kHz")
    print("No external feature extraction needed!")
    print()
    print("Next steps:")
    print("  1. Quantize: python quantize.py --input-dir ./onnx-baked")
    print("  2. Use in C# with ONNX Runtime directly")


if __name__ == "__main__":
    main()