import pickle
import onnx
from onnxruntime import quantization
from onnxruntime.quantization import QuantFormat, QuantType, quantize_dynamic
from warnings import simplefilter

# Suppress all future warnings
simplefilter(action='ignore', category=FutureWarning)

class InputDataReader(quantization.CalibrationDataReader):
    """ONNX Runtime calibration reader backed by data_directory/input.pkl."""
    def __init__(self, input_path):
        # Calibration uses the same pickled input format as normal inference.
        with open(input_path, 'rb') as f:
            self.input_data = pickle.load(f)

        self.enum_data = iter(self.preprocess(self.input_data))
        


    def preprocess(self, inputs):
        
        for input in inputs:
            # Prep input
            padding_mask = (input == -1).all(dim=2)

            yield {
                'hits': input.cpu().numpy(),
                'padding_mask': padding_mask.cpu().numpy()
            }


    def get_next(self):
        try:
            return next(self.enum_data)
        except StopIteration:
            return None

def quantize_model(data_dir):
    """Create model_quant.onnx from model.onnx using calibration inputs."""
    
    with open(f"{data_dir}input.pkl", 'rb') as f:
        input_data = pickle.load(f)

    data_reader = InputDataReader(f"{data_dir}input.pkl")

    config = quantization.StaticQuantConfig(
        calibration_data_reader=data_reader,
        activation_type=QuantType.QInt16,
        weight_type=QuantType.QInt16,
        quant_format=QuantFormat.QDQ,
        extra_options={"ActivationSymmetric": False}
    )

    quantization.quantize(f"{data_dir}model.onnx", f"{data_dir}model_quant.onnx", config)

    #quantize_dynamic(f"{data_dir}model.onnx", f"{data_dir}model_quant.onnx", weight_type=QuantType.QInt16)

    
