"""
Quantize vit_large_patch16_224_fp32.onnx using Nvidia Model Optimizer.

Quantization modes attempted:
  - FP16:  convert_to_f16
  - INT8:  entropy / max calibration
  - FP8:   entropy / max calibration
  - INT4:  awq_clip / awq_lite  (symmetric + asymmetric) / awq_full / rtn_dq

Calibration data: vit_large_patch16_224_calibration.npy (20 images, shape 20x3x224x224)

Usage:
    python3 vit_large_patch16_224_quantize_all.py                # run all modes
    python3 vit_large_patch16_224_quantize_all.py --mode fp16    # run FP16 only
    python3 vit_large_patch16_224_quantize_all.py --mode int8    # run INT8 entropy + max
    python3 vit_large_patch16_224_quantize_all.py --mode fp8     # run FP8 entropy + max
    python3 vit_large_patch16_224_quantize_all.py --mode int4    # run all INT4 variants

Running one mode at a time avoids OOM kills on memory-constrained systems.
"""

import argparse
import os
import sys
import time
import numpy as np
import onnx
import onnxruntime as ort
import gc


# ---------------------------------------------------------------------------
# Post-quantization repair for INT4 models
# ---------------------------------------------------------------------------

def _fix_int4_dq_axis(output_path):
    """Fix ModelOpt bug: classifier.weight DequantizeLinear has axis=0 but
    scale shape is correct for axis=1 with block_size=128.

    ORT validates ceil(Di/block_size) on the declared axis and rejects the
    model at runtime.  Patching axis 0 -> 1 makes the scale shape valid.
    """
    model = onnx.load(output_path)
    fixed = False

    for node in model.graph.node:
        if node.op_type != "DequantizeLinear":
            continue
        if "classifier.weight" not in node.name:
            continue

        axis = None
        block_size = None
        for attr in node.attribute:
            if attr.name == "axis":
                axis = attr.i
            elif attr.name == "block_size":
                block_size = attr.i

        if axis != 0 or block_size is None:
            continue

        weight_name = node.input[0]
        scale_name = node.input[1]
        weight_shape = scale_shape = None
        for init in model.graph.initializer:
            if init.name == weight_name:
                weight_shape = list(init.dims)
            if init.name == scale_name:
                scale_shape = list(init.dims)

        if weight_shape is None or scale_shape is None:
            continue

        expected_axis1 = list(weight_shape)
        expected_axis1[1] = (expected_axis1[1] + block_size - 1) // block_size

        if scale_shape == expected_axis1:
            for attr in node.attribute:
                if attr.name == "axis":
                    attr.i = 1
                    fixed = True
                    print(f"  [fix] Patched {node.name}: axis 0 -> 1 "
                          f"(scale {scale_shape} matches axis=1, block_size={block_size})")

    if fixed:
        onnx.save(model, output_path)
        print(f"  [fix] Saved repaired model to {output_path}")
    else:
        print(f"  [fix] No classifier DQ axis issues found — model unchanged")
# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
ONNX_PATH = "vit_large_patch16_224_fp32.onnx"
CALIB_PATH = "vit_large_patch16_224_calibration.npy"
MODEL_NAME = "vit_large_patch16_224"

# Load calibration data
calib_data = np.load(CALIB_PATH)
print(f"Calibration data: shape={calib_data.shape}, dtype={calib_data.dtype}")

# ---------------------------------------------------------------------------
# FP16 Quantization
# ---------------------------------------------------------------------------
def run_fp16():
    print("\n" + "=" * 70)
    print("FP16 Quantization")
    print("=" * 70)

    from modelopt.onnx.quantization.int8 import convert_to_f16

    onnx_model = onnx.load(ONNX_PATH, load_external_data=True)
    fp16_model = convert_to_f16(onnx_model, keep_io_types=True)

    output_path = "fp16/vit_large_patch16_224_fp16.onnx"
    if not os.path.exists("fp16"):
        os.mkdir("fp16")
        
    onnx.save(fp16_model, output_path)
    size_mb = os.path.getsize(output_path) / 1e6
    print(f"  Saved: {output_path} ({size_mb:.1f} MB)")
    return output_path


# ---------------------------------------------------------------------------
# ViT-specific: find nodes that must stay in FP32 for accurate quantization
# ---------------------------------------------------------------------------

def _find_vit_nodes_to_exclude(onnx_path):
    """Find all MatMul and Add nodes that must stay in FP32 for ViT.

    ViT-Large has 192 MatMuls per 24 transformer layers:
      - 72 Q/K/V projections (fed by LayerNorm, 3/layer)
      - 24 Q@K^T attention score MatMuls (feed into Softmax)
      - 24 attn@V attention apply MatMuls
      - 24 attention output projection MatMuls
      - 24 MLP input MatMuls (fed by LayerNorm)
      - 24 MLP hidden MatMuls (SAFE to quantize)

    Additionally, 49 Add nodes feed directly into LayerNorm (residual→Pre-LN).
    Quantizing these inserts Q/DQ before LayerNorm, amplifying noise via x/std.

    Only the 24 MLP hidden MatMuls are safe for INT8 quantization.
    All 168 attention-path MatMuls must stay FP32 because:
      - LayerNorm→MatMul: Q/DQ on LN output amplifies noise via x/std
      - Q@K^T: INT8 error in scores gets exponentially amplified by Softmax
      - attn@V: error in attention weights corrupts the value aggregation
      - Output proj: receives attn@V result which must stay accurate
    """
    model = onnx.load(onnx_path, load_external_data=False)

    # Build maps
    output_to_producer = {}
    output_to_consumers = {}
    for node in model.graph.node:
        for out in node.output:
            output_to_producer[out] = node
        for inp in node.input:
            output_to_consumers.setdefault(inp, []).append(node)

    all_matmuls = [n for n in model.graph.node if n.op_type == "MatMul"]
    excluded = set()

    # 1. MatMuls fed directly by LayerNorm (Q/K/V projections + MLP input)
    for mm in all_matmuls:
        for inp in mm.input:
            prod = output_to_producer.get(inp)
            if prod and prod.op_type == "LayerNormalization":
                excluded.add(mm.name)
                break

    # 2. MatMuls that feed into Softmax (Q@K^T attention scores)
    #    Trace backwards from each Softmax through arithmetic/shape ops
    _passthrough = {"Mul", "Div", "Reshape", "Transpose", "Cast", "Add",
                    "Slice", "Where", "IsNaN", "Concat", "Shape", "Expand"}
    for sm in (n for n in model.graph.node if n.op_type == "Softmax"):
        visited, queue = set(), [sm]
        while queue:
            node = queue.pop(0)
            if node.name in visited:
                continue
            visited.add(node.name)
            for inp in node.input:
                prod = output_to_producer.get(inp)
                if prod:
                    if prod.op_type == "MatMul":
                        excluded.add(prod.name)
                    elif prod.op_type in _passthrough:
                        queue.append(prod)

    # 3. MatMuls that apply attention weights (attn@V)
    #    Trace forwards from each Softmax through arithmetic/shape ops
    for sm in (n for n in model.graph.node if n.op_type == "Softmax"):
        visited, queue = set(), list(sm.output)
        for _ in range(10):
            next_q = []
            for out_name in queue:
                for consumer in output_to_consumers.get(out_name, []):
                    if consumer.name in visited:
                        continue
                    visited.add(consumer.name)
                    if consumer.op_type == "MatMul":
                        excluded.add(consumer.name)
                    elif consumer.op_type in _passthrough:
                        next_q.extend(consumer.output)
            queue = next_q
            if not queue:
                break

    # 4. Attention output projection MatMuls
    #    Trace forward from each attn@V MatMul to find the next MatMul
    attn_v_names = {n.name for n in all_matmuls if n.name in excluded
                    and n.name not in {mm.name for mm in all_matmuls
                                       if any(output_to_producer.get(inp, None) is not None
                                              and output_to_producer[inp].op_type == "LayerNormalization"
                                              for inp in mm.input)}}
    for mm in all_matmuls:
        if mm.name in attn_v_names:
            visited, queue = set(), list(mm.output)
            for _ in range(5):
                next_q = []
                for out_name in queue:
                    for consumer in output_to_consumers.get(out_name, []):
                        if consumer.name in visited:
                            continue
                        visited.add(consumer.name)
                        if consumer.op_type == "MatMul" and consumer.name not in excluded:
                            excluded.add(consumer.name)
                        elif consumer.op_type in _passthrough:
                            next_q.extend(consumer.output)
                queue = next_q
                if not queue:
                    break

    # 5. Add nodes whose output feeds into LayerNorm (residual→Pre-LN path)
    #    Quantizing these inserts Q/DQ before LayerNorm, same problem as #1.
    add_to_ln = 0
    for node in model.graph.node:
        if node.op_type == "Add":
            for out in node.output:
                for consumer in output_to_consumers.get(out, []):
                    if consumer.op_type == "LayerNormalization":
                        excluded.add(node.name)
                        add_to_ln += 1
                        break

    del model
    matmul_excluded = len(excluded & {n.name for n in all_matmuls})
    matmul_safe = len(all_matmuls) - matmul_excluded
    print(f"  Excluded: {matmul_excluded} attention MatMuls + {add_to_ln} residual→LN Adds "
          f"= {len(excluded)} nodes total")
    print(f"  Safe to quantize: {matmul_safe} MLP MatMuls + remaining Adds + Conv")
    return sorted(excluded)


# ---------------------------------------------------------------------------
# INT8 Quantization
# ---------------------------------------------------------------------------
def run_int8(method, label=None):
    """Run INT8 quantization with a given calibration method."""
    if label is None:
        label = f"int8_{method}{zp_tag}"

    print("\n" + "=" * 70)
    print(f"INT8 Quantization — method={method}")
    print("=" * 70)

    from modelopt.onnx.quantization import quantize

    # Find attention-path MatMul nodes that must stay in FP32
    nodes_to_exclude = _find_vit_nodes_to_exclude(ONNX_PATH)

    output_path = f"int8/vit_large_patch16_224_{label}.onnx"
    if not os.path.exists("int8"):
        os.mkdir("int8")
    t0 = time.time()
    quantize(
        onnx_path=ONNX_PATH,
        quantize_mode="int8",
        calibration_data=calib_data,
        calibration_method=method,
        output_path=output_path,
        use_external_data_format=True,
        calibration_eps=["cuda:0", "cpu"],
        opset=19,
        # ViT-specific fixes: the decomposed attention mechanism and LayerNorm
        # are extremely sensitive to INT8 quantization.
        # Without these exclusions, Top-1 drops from 82% to 0%.
        #
        # Root cause: ModelOpt's QDQ pattern inserts Q/DQ on the activation
        # inputs of quantized MatMuls. For attention-path MatMuls, this
        # quantizes LayerNorm outputs (amplified by x/std normalization) and
        # attention scores (exponentially amplified by Softmax).
        #
        # Fix: allowlist only MatMul+Add+Conv, and exclude all 168 attention-
        # path MatMuls via nodes_to_exclude. Only 24 MLP hidden MatMuls are
        # quantized.
        op_types_to_quantize=["MatMul", "Add", "Conv"],
        nodes_to_exclude=nodes_to_exclude,
    )
    elapsed = time.time() - t0

    # Calculate total size (onnx + .data)
    size_mb = os.path.getsize(output_path) / 1e6
    data_path = output_path + ".data"
    if os.path.exists(data_path):
        size_mb += os.path.getsize(data_path) / 1e6
    print(f"  Saved: {output_path} ({size_mb:.1f} MB) in {elapsed:.1f}s")
    print("ir_version: " + str(onnx.load(output_path).ir_version))
    return output_path


# ---------------------------------------------------------------------------
# FP8 Quantization
# ---------------------------------------------------------------------------
def run_fp8(method):
    print("\n" + "=" * 70)
    print(f"FP8 Quantization — method={method}")
    print("=" * 70)

    from modelopt.onnx.quantization import quantize

    # Find attention-path MatMul nodes that must stay in FP32
    nodes_to_exclude = _find_vit_nodes_to_exclude(ONNX_PATH)

    output_path = f"fp8/vit_large_patch16_224_fp8_{method}.onnx"
    if not os.path.exists("fp8"):
        os.mkdir("fp8")
    t0 = time.time()
    quantize(
        onnx_path=ONNX_PATH,
        quantize_mode="fp8",
        calibration_data=calib_data,
        calibration_method=method,
        output_path=output_path,
        use_external_data_format=True,
        calibration_eps=["cuda:0", "cpu"],
        opset=19,
        # ViT-specific fixes (same rationale as INT8)        op_types_to_quantize=["MatMul", "Add", "Conv"],
        nodes_to_exclude=nodes_to_exclude,
    )
    elapsed = time.time() - t0

    size_mb = os.path.getsize(output_path) / 1e6
    data_path = output_path + ".data"
    if os.path.exists(data_path):
        size_mb += os.path.getsize(data_path) / 1e6
    print(f"  Saved: {output_path} ({size_mb:.1f} MB) in {elapsed:.1f}s")
    return output_path


# ---------------------------------------------------------------------------
# INT4 Quantization
# ---------------------------------------------------------------------------
def run_int4(method, use_zero_point=False):
    print("\n" + "=" * 70)
    print(f"INT4 Quantization — method={method}, zero_point={use_zero_point}")
    print("=" * 70)

    from modelopt.onnx.quantization import quantize

    zp_tag = "_asym" if use_zero_point else ""
    output_path = f"int4/vit_large_patch16_224_int4_{method}{zp_tag}.onnx"
    if not os.path.exists("int4"):
        os.mkdir("int4")
    t0 = time.time()
    quantize(
        onnx_path=ONNX_PATH,
        quantize_mode="int4",
        calibration_data=calib_data,
        calibration_method=method,
        use_zero_point=use_zero_point,
        output_path=output_path,
        use_external_data_format=True,
        calibration_eps=["cuda:0", "cpu"],
        opset=21,
    )
    elapsed = time.time() - t0

    size_mb = os.path.getsize(output_path) / 1e6
    data_path = output_path + ".data"
    if os.path.exists(data_path):
        size_mb += os.path.getsize(data_path) / 1e6
    print(f"  Saved: {output_path} ({size_mb:.1f} MB) in {elapsed:.1f}s")

    # Post-quantization repair: fix classifier DequantizeLinear axis bug
    _fix_int4_dq_axis(output_path)

    return output_path


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
ALL_MODES = ["fp16", "int8", "fp8", "int4"]

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Quantize vit_large_patch16_224 using NVIDIA ModelOpt"
    )
    parser.add_argument(
        "--mode", type=str, nargs="*", default=ALL_MODES,
        choices=ALL_MODES,
        help=f"Quantization mode(s) to run (default: all). "
             f"Run one at a time to avoid OOM on large models. "
             f"Choices: {ALL_MODES}",
    )
    args = parser.parse_args()
    modes = args.mode
    print(f"Running modes: {modes}")

    results = {}

    # --- FP16 ---
    if "fp16" in modes:
        try:
            results["fp16"] = run_fp16()
        except Exception as e:
            print(f"  FP16 FAILED: {e}")
        finally:
            gc.collect()

    # --- INT8 variants ---
    if "int8" in modes:
        for method in ["entropy", "max"]:
            label = f"int8_{method}"
            try:
                results[label] = run_int8(method, label=label)
            except Exception as e:
                print(f"  {label} FAILED: {e}")
            finally:
                gc.collect()

    # --- FP8 variants ---
    if "fp8" in modes:
        for method in ["entropy", "max"]:
            label = f"fp8_{method}"
            try:
                results[label] = run_fp8(method)
            except Exception as e:
                print(f"  {label} FAILED: {e}")
            finally:
                gc.collect()

    # --- INT4 variants ---
    if "int4" in modes:
        for method in ["awq_clip", "awq_full", "rtn_dq"]:
            label = f"int4_{method}"
            try:
                results[label] = run_int4(method, use_zero_point=False)
            except Exception as e:
                print(f"  {label} FAILED: {e}")
            finally:
                gc.collect()
        # special treat for awq_lite
        method = "awq_lite"
        for zp in [False, True]:
            label = f"int4_{method}" + ("_asym" if zp else "")
            try:
                results[label] = run_int4(method, use_zero_point=zp)
            except Exception as e:
                print(f"  {label} FAILED: {e}")
            finally:
                gc.collect()

    # --- Summary ---
    print("\n" + "=" * 70)
    print("Quantization Summary")
    print("=" * 70)
    for label, path in results.items():
        size_mb = os.path.getsize(path) / 1e6
        data_path = path + ".data"
        if os.path.exists(data_path):
            size_mb += os.path.getsize(data_path) / 1e6
        print(f"  {label:30s} → {path:50s} ({size_mb:.1f} MB)")