import torch
import torch.onnx
from JiRackTernaryPyTorch_1b_inf import TernaryTransformer1B, TernaryConfig
from onnxruntime.quantization import quantize_dynamic, QuantType
import os
import types

def export_and_quantize():
    device = torch.device("cpu")
    checkpoint_path = "/root/JiRackTernary1/new/model_packed.safetensors"
    onnx_raw = "/root/JiRackTernary1/new/model/jirack_1b_kv_raw.onnx"
    onnx_final = "/root/JiRackTernary1/new/model/jirack_1b_kv_int8.onnx"

    config = TernaryConfig()
    model = TernaryTransformer1B(config).to(device)
    model.load_prod_weights(checkpoint_path, device)
    model.eval()

    # Сохраняем оригинальный forward
    original_forward = model.forward

    # Временная замена forward для экспорта
    def forward_with_past(self, input_ids, *past_key_values):
        # Проверяем количество past тензоров
        expected_count = config.num_hidden_layers * 2
        if len(past_key_values) != expected_count:
            raise ValueError(f"Ожидалось {expected_count} past тензоров, получено {len(past_key_values)}")

        # Вызываем оригинальный forward БЕЗ лишнего self
        logits, _ = original_forward(input_ids)  # ← исправлено здесь

        # Dummy present для экспорта (ONNX увидит форму и свяжет граф)
        head_dim = config.hidden_size // config.num_attention_heads  # 64
        seq_len = input_ids.shape[1]
        present = []
        for _ in range(config.num_hidden_layers):
            k_dummy = torch.zeros(1, config.num_attention_heads, seq_len, head_dim,
                                  dtype=torch.float16, device=device)
            v_dummy = torch.zeros(1, config.num_attention_heads, seq_len, head_dim,
                                  dtype=torch.float16, device=device)
            present.append(k_dummy)
            present.append(v_dummy)

        return (logits, *present)

    # Применяем monkey-patch
    model.forward = types.MethodType(forward_with_past, model)

    # Dummy входы
    dummy_input_ids = torch.randint(0, config.vocab_size, (1, 1)).to(device)  # seq_len=1

    head_dim = config.hidden_size // config.num_attention_heads
    dummy_past = []
    for _ in range(config.num_hidden_layers):
        k = torch.zeros(1, config.num_attention_heads, 0, head_dim, dtype=torch.float16, device=device)
        v = torch.zeros(1, config.num_attention_heads, 0, head_dim, dtype=torch.float16, device=device)
        dummy_past.extend([k, v])

    dummy_inputs = (dummy_input_ids, *dummy_past)

    # Имена входов и выходов
    input_names = ['input_ids']
    output_names = ['logits']

    for i in range(config.num_hidden_layers):
        input_names.extend([f'past_key_values.{i}.key', f'past_key_values.{i}.value'])
        output_names.extend([f'present.{i}.key', f'present.{i}.value'])

    print("🚀 Экспорт ONNX с KV-cache...")

    torch.onnx.export(
        model,
        dummy_inputs,
        onnx_raw,
        export_params=True,
        opset_version=17,
        do_constant_folding=True,
        input_names=input_names,
        output_names=output_names,
        dynamic_axes={
            'input_ids': {0: 'batch_size', 1: 'sequence_length'},
            'logits': {0: 'batch_size', 1: 'sequence_length'},
            **{f'past_key_values.{i}.key':   {0: 'batch_size', 2: 'past_sequence_length'} for i in range(config.num_hidden_layers)},
            **{f'past_key_values.{i}.value': {0: 'batch_size', 2: 'past_sequence_length'} for i in range(config.num_hidden_layers)},
            **{f'present.{i}.key':   {0: 'batch_size', 2: 'total_sequence_length'} for i in range(config.num_hidden_layers)},
            **{f'present.{i}.value': {0: 'batch_size', 2: 'total_sequence_length'} for i in range(config.num_hidden_layers)},
        }
    )

    print(f"✅ Сырой ONNX сохранён: {onnx_raw}")

    # Восстанавливаем оригинальный forward (на всякий случай)
    model.forward = original_forward

    # Квантизация
    print("🔥 Квантизация в int8...")
    quantize_dynamic(
        onnx_raw,
        onnx_final,
        weight_type=QuantType.QInt8,
        per_channel=False,
        reduce_range=True
    )

    print(f"✅ Готовая модель с KV-cache и int8: {onnx_final}")

    # Если не нужен промежуточный файл — раскомментируй
    # os.remove(onnx_raw)

if __name__ == "__main__":
    export_and_quantize()