#!/usr/bin/env python3
"""Convert a tuned vanilla Qwen3 checkpoint into a quality-safe AHA hot start.

All vanilla weights are preserved. Dynamic gate weights start at zero and a
trainable bias initializes every gate to the requested full-attention
probability, so hard routing is initially exactly vanilla full attention.
"""

from __future__ import annotations

import argparse
import json
import math
import sys
from pathlib import Path

import torch
from torch import nn
from transformers import AutoTokenizer

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from modeling_aha_qwen3 import AHAQwen3Config, AHAQwen3ForCausalLM


def add_zero_bias(linear: nn.Linear) -> nn.Linear:
    new = nn.Linear(
        linear.in_features,
        linear.out_features,
        bias=True,
        device=linear.weight.device,
        dtype=linear.weight.dtype,
    )
    with torch.no_grad():
        new.weight.copy_(linear.weight)
        new.bias.zero_()
    return new


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--vanilla-path", required=True)
    parser.add_argument("--output-path", required=True)
    parser.add_argument("--window-size", type=int, default=128)
    parser.add_argument("--gate-init-full-prob", type=float, default=0.90)
    parser.add_argument(
        "--local-kind", choices=("sliding_window", "sink_recent"), default="sliding_window"
    )
    parser.add_argument(
        "--router-granularity",
        choices=("token", "token_kv_head"),
        default="token_kv_head",
        help="Native dynamic gate shape. token is one shared gate per layer/token.",
    )
    args = parser.parse_args()

    if not 0.5 < args.gate_init_full_prob < 1.0:
        raise ValueError("--gate-init-full-prob must be strictly between 0.5 and 1")
    output = Path(args.output_path)
    output.mkdir(parents=True, exist_ok=True)

    AHAQwen3Config.register_for_auto_class()
    AHAQwen3ForCausalLM.register_for_auto_class("AutoModelForCausalLM")
    model = AHAQwen3ForCausalLM.from_pretrained_qwen3(
        args.vanilla_path,
        aha_window_size=args.window_size,
        aha_local_kind=args.local_kind,
        aha_router_granularity=args.router_granularity,
        aha_mode="dynamic",
        aha_gate_target=0.70,
        aha_reg_weight=-1.0,
        torch_dtype=torch.bfloat16,
        attn_implementation="sdpa",
    )

    # Qwen3 normally has bias-free projections. A gate bias gives a stable,
    # token-independent full-attention hot start while zero biases on q/k/v/o
    # keep the original attention computation bit-for-bit unchanged.
    if not model.config.attention_bias:
        for layer in model.model.layers:
            attn = layer.self_attn
            attn.q_proj = add_zero_bias(attn.q_proj)
            attn.k_proj = add_zero_bias(attn.k_proj)
            attn.v_proj = add_zero_bias(attn.v_proj)
            attn.o_proj = add_zero_bias(attn.o_proj)
        model.config.attention_bias = True

    q_rows = model.config.num_attention_heads * model.config.head_dim
    gate_logit = math.log(args.gate_init_full_prob / (1.0 - args.gate_init_full_prob))
    with torch.no_grad():
        for layer in model.model.layers:
            q_proj = layer.self_attn.q_proj
            q_proj.weight[q_rows:].zero_()
            q_proj.bias[q_rows:].fill_(gate_logit)

    # Keep a distinct LM head so the custom checkpoint has an explicit,
    # self-contained state dict under safetensors.
    model.config.tie_word_embeddings = False
    if model.lm_head.weight.data_ptr() == model.model.embed_tokens.weight.data_ptr():
        model.lm_head.weight = nn.Parameter(model.lm_head.weight.detach().clone())
    model.config.aha_hotstart_source = str(Path(args.vanilla_path).resolve())
    model.config.aha_gate_init_full_prob = args.gate_init_full_prob

    tokenizer = AutoTokenizer.from_pretrained(args.vanilla_path, trust_remote_code=True)
    model.save_pretrained(output, safe_serialization=True)
    tokenizer.save_pretrained(output)
    gate_rows = model.model.layers[0].self_attn.aha_router_outputs
    gate_parameters = len(model.model.layers) * gate_rows * (
        model.config.hidden_size + 1
    )
    (output / "aha_hotstart_manifest.json").write_text(
        json.dumps(
            {
                "source_checkpoint": str(Path(args.vanilla_path).resolve()),
                "router_granularity": args.router_granularity,
                "native_gate_rows_per_layer": gate_rows,
                "effective_gate_parameters": gate_parameters,
                "gate_init_full_probability": args.gate_init_full_prob,
                "gate_weight_init": "zeros",
                "gate_bias_logit": gate_logit,
                "local_attention": {
                    "kind": args.local_kind,
                    "sink_size": model.config.duo_sink_size,
                    "recent_size": model.config.duo_recent_size,
                },
            },
            indent=2,
        )
        + "\n"
    )
    print(
        f"saved={output} init_full_prob={args.gate_init_full_prob:.4f} "
        f"gate_logit={gate_logit:.6f} hard_sparsity=0.0 "
        f"router_granularity={args.router_granularity} gate_params={gate_parameters:,}"
    )


if __name__ == "__main__":
    main()