#!/usr/bin/env python3
"""Export Qwen3-0.6B Q8_0 with all optimization passes for Metal benchmarking.

Usage (from repo root):
    source .venv/bin/activate
    python runner/export_qwen3_q8_optimized.py
"""

import os
import sys
import time
import types

import torch

# Monkey-patch to work around torchao/torch version mismatch
# Must be done AFTER torch import so torch.ao.quantization exists
import torch.ao.quantization as _taoq
if not hasattr(_taoq, 'quantizer'):
    _mod = types.ModuleType('torch.ao.quantization.quantizer')
    _inner = types.ModuleType('torch.ao.quantization.quantizer.quantizer')
    class _FakeQuantizer: pass
    _inner.Quantizer = _FakeQuantizer
    _mod.quantizer = _inner
    _taoq.quantizer = _mod
    sys.modules['torch.ao.quantization.quantizer'] = _mod
    sys.modules['torch.ao.quantization.quantizer.quantizer'] = _inner

# Must import executorch_ggml first for RTLD_GLOBAL symbol resolution
import executorch_ggml  # noqa: F401, I001  # isort: skip


def export_qwen3_q8_optimized(max_seq_len: int = 128):
    """Export Qwen3-0.6B Q8_0 with swap_rms_norm, fold_rms_norm_weights,
    fuse_rope_in_graph, strip_gqa_expand."""
    from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower
    from executorch_ggml import GgmlPartitioner, GgmlQuantConfig
    from executorch_ggml.passes import RemoveGraphAssertsPass
    from executorch_ggml.passes.replace_copy_ops_pass import ReplaceCopyOpsPass
    from executorch_ggml.modules.rms_norm import swap_rms_norm
    from executorch_ggml.passes.fold_rms_norm_weights import fold_rms_norm_weights
    from executorch_ggml.passes.fuse_rope_pass import fuse_rope_in_graph
    from executorch_ggml.passes.strip_gqa_expand_pass import strip_gqa_expand
    from optimum.exporters.executorch.integrations import CausalLMExportableModule
    from transformers import AutoConfig, AutoModelForCausalLM, GenerationConfig

    model_id = "Qwen/Qwen3-0.6B"
    config = AutoConfig.from_pretrained(model_id)
    if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
        config.rope_scaling["type"] = "default"

    head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
    freq_base = getattr(config, "rope_theta", 10000.0)
    print(f"head_dim={head_dim}, freq_base={freq_base}")

    print(f"Loading {model_id}...")
    eager_model = AutoModelForCausalLM.from_pretrained(
        model_id,
        device_map="cpu",
        torch_dtype=torch.float32,
        config=config,
        attn_implementation="sdpa",
        generation_config=GenerationConfig(
            use_cache=True,
            cache_implementation="static",
            max_length=max_seq_len,
            cache_config={"batch_size": 1, "max_cache_len": max_seq_len},
        ),
    )

    # 1. Module swaps (before export)
    n_swapped = swap_rms_norm(eager_model)
    print(f"swap_rms_norm: swapped {n_swapped} modules")
    fold_rms_norm_weights(eager_model)
    print("fold_rms_norm_weights: done")

    # 2. Export
    print("Exporting with torch.export...")
    exportable = CausalLMExportableModule(
        eager_model,
        max_seq_len=max_seq_len,
        use_custom_kv_cache=False,
        use_custom_sdpa=False,
        disable_dynamic_shapes=False,
    )
    ep = exportable.export()["model"]

    # 3. AOT graph passes
    fuse_rope_in_graph(ep.graph_module, head_dim, freq_base)
    print("fuse_rope_in_graph: done")
    strip_gqa_expand(ep.graph_module)
    print("strip_gqa_expand: done")

    # 4. Edge lowering with Q8_0 quantization
    print("Lowering with Q8_0 quantization...")
    t0 = time.time()
    quant_config = GgmlQuantConfig()
    edge_mgr = to_edge_transform_and_lower(
        ep,
        partitioner=[GgmlPartitioner(quant_config=quant_config)],
        compile_config=EdgeCompileConfig(
            _check_ir_validity=False,
            _skip_dim_order=True,
        ),
        transform_passes=[ReplaceCopyOpsPass(), RemoveGraphAssertsPass()],
        constant_methods=exportable.metadata,
    )
    t1 = time.time()
    print(f"Lowering took {t1 - t0:.1f}s")

    return edge_mgr


def main():
    max_seq_len = 128
    out_dir = "/Users/larryliu/executorch-ggml/qwen3"
    q8_path = os.path.join(out_dir, "qwen3_q8_0_optimized.pte")

    edge_mgr = export_qwen3_q8_optimized(max_seq_len=max_seq_len)
    et = edge_mgr.to_executorch()
    pte_bytes = et.buffer

    os.makedirs(out_dir, exist_ok=True)
    with open(q8_path, "wb") as f:
        f.write(pte_bytes)

    q8_size_mb = len(pte_bytes) / (1024 * 1024)
    print(f"\nSaved optimized Q8_0 .pte to {q8_path} ({q8_size_mb:.1f} MB)")


if __name__ == "__main__":
    main()