import os

import onnx
import onnxruntime
import torch
import torch.nn as nn
from onnxruntime.quantization import (
    CalibrationDataReader,
    QuantFormat,
    QuantType,
    quantize_static,
)
from torchvision import datasets, models, transforms

# Configuration
NUM_CLASSES = 10
DATA_DIR = "dataset"
PYTORCH_MODEL_PATH = "best_thai_food_model.pth"
ONNX_MODEL_PATH = "efficientnet-b0.onnx"
QUANTIZED_MODEL_PATH = "efficientnet-b0_quant.onnx"

print("Starting model conversion process...")

# 1. Load PyTorch Model Structure and Weights
model = models.efficientnet_b0(pretrained=False)
model.classifier[1] = nn.Linear(model.classifier[1].in_features, NUM_CLASSES)
model.load_state_dict(torch.load(PYTORCH_MODEL_PATH, map_location='cpu'))
model.eval()

print("PyTorch model loaded successfully.")

# 2. Export to ONNX Format
dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model, 
    dummy_input, 
    ONNX_MODEL_PATH,
    export_params=True, 
    opset_version=13, 
    do_constant_folding=True,
    input_names=['input'], 
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
)

print(f"ONNX export complete -> Size: {os.path.getsize(ONNX_MODEL_PATH) / 1e6:.2f} MB")

# 3. Create Calibration Data Reader for Static Quantization
class ThaiFoodCalibrationDataReader(CalibrationDataReader):
    def __init__(self, data_dir, batch_size=1, num_samples=100):
        transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        
        dataset = datasets.ImageFolder(os.path.join(data_dir, 'val'), transform)
        
        self.num_samples = min(num_samples, len(dataset))
        subset = torch.utils.data.Subset(dataset, range(self.num_samples))
        
        self.dataloader = torch.utils.data.DataLoader(subset, batch_size=batch_size, shuffle=False)
        self.iterator = iter(self.dataloader)
        self.input_name = "input"

    def get_next(self):
        try:
            inputs, _ = next(self.iterator)
            return {self.input_name: inputs.numpy()}
        except StopIteration:
            return None

# 4. Perform Static Quantization
print("Preparing calibration data...")
calibration_reader = ThaiFoodCalibrationDataReader(data_dir=DATA_DIR)

print("Starting static quantization (this may take a few minutes)...")
quantize_static(
    model_input=ONNX_MODEL_PATH,
    model_output=QUANTIZED_MODEL_PATH,
    calibration_data_reader=calibration_reader,
    quant_format=QuantFormat.QDQ,
    activation_type=QuantType.QUInt8,
    weight_type=QuantType.QInt8
)

print(f"Static quantization complete -> Size: {os.path.getsize(QUANTIZED_MODEL_PATH) / 1e6:.2f} MB")
print("Quantized ONNX model is ready for inference.")