# import torch
# import torch.nn as nn
# from torch.ao.quantization.observer import MinMaxObserver

# # Step 1: Define a dummy nn.Conv2d
# x = torch.randn(1, 3, 32, 32)
# conv = nn.Conv2d(in_channels=3, out_channels=8, kernel_size=3, padding=1)
# main_output = conv(x)
# # Step 2: Attach quantization metadata to the original module
# conv.activation_dtype = "uint8"
# conv.parameter_dtype = "uint8"
# conv.activation_observer = MinMaxObserver
# conv.parameter_observer = MinMaxObserver

# # Step 3: Wrap it with QuantizedConv2d
# qconv = QuantizedConv2d(old_module=conv)

# # Step 4: Run dummy input
# output = qconv(x)

# # Step 5: Check output
# print("Output shape:", output, main_output)


# import torch
# import torch.nn as nn
# import torch.nn.functional as F
# from torch.ao.quantization import FakeQuantize, MovingAverageMinMaxObserver

# # ======================================================================================
# # 1. Helper Function for Quantization Configuration
# # ======================================================================================

# def get_quant_config(bits: int, is_symmetric: bool):
#     """Returns quant_min, quant_max, and torch.dtype for a given bitwidth."""
#     if is_symmetric:
#         # For symmetric quantization (typically for weights)
#         if bits == 8:
#             return -128, 127, torch.qint8
#         elif bits == 4:
#             return -8, 7, torch.qint8
#         else:
#             raise ValueError(f"Unsupported symmetric bitwidth: {bits}")
#     else:
#         # For asymmetric quantization (typically for activations)
#         if bits == 8:
#             return 0, 255, torch.quint8
#         elif bits == 4:
#             return 0, 15, torch.quint8
#         else:
#             raise ValueError(f"Unsupported asymmetric bitwidth: {bits}")

# # ======================================================================================
# # 2. Base Class for All Quantized Modules
# # ======================================================================================

# class QuantizationMixin(nn.Module):
#     """
#     Base mixin for custom quantized modules.
#     Handles the creation of FakeQuantize modules for inputs, outputs, and parameters.
#     """
#     def __init__(self, old_module: nn.Module, activation_bits: int = 8, weight_bits: int = 8):
#         super().__init__()
#         self.old_module = old_module

#         # Activation Quantizer (asymmetric)
#         act_qmin, act_qmax, act_dtype = get_quant_config(bits=activation_bits, is_symmetric=False)
#         self.input_quantizer = FakeQuantize(
#             observer=MovingAverageMinMaxObserver, quant_min=act_qmin, quant_max=act_qmax,
#             dtype=act_dtype, qscheme=torch.per_tensor_affine, reduce_range=False
#         )
#         self.output_quantizer = FakeQuantize(
#             observer=MovingAverageMinMaxObserver, quant_min=act_qmin, quant_max=act_qmax,
#             dtype=act_dtype, qscheme=torch.per_tensor_affine, reduce_range=False
#         )

#         # Weight Quantizer (symmetric)
#         self.param_quantizers = nn.ModuleDict()
#         if not list(self.old_module.named_parameters(recurse=False)):
#             return

#         weight_qmin, weight_qmax, weight_dtype = get_quant_config(bits=weight_bits, is_symmetric=True)
#         weight_qscheme = torch.per_tensor_symmetric
#         if isinstance(self.old_module, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
#             weight_qscheme = torch.per_channel_affine

#         for name, _ in self.old_module.named_parameters(recurse=False):
#             if 'weight' in name:
#                 self.param_quantizers[name] = FakeQuantize(
#                     observer=MovingAverageMinMaxObserver, quant_min=weight_qmin, quant_max=weight_qmax,
#                     dtype=weight_dtype, qscheme=weight_qscheme, reduce_range=False,
#                     ch_axis=0 if weight_qscheme == torch.per_channel_affine else -1
#                 )

#     def forward(self, *args, **kwargs):
#         raise NotImplementedError("Forward pass must be implemented by subclasses.")

# # ======================================================================================
# # 3. CONVOLUTIONAL AND LINEAR LAYERS
# # ======================================================================================

# class QuantizedConv1d(QuantizationMixin):
#     def forward(self, x):
#         qx = self.input_quantizer(x)
#         qw = self.param_quantizers['weight'](self.old_module.weight)
#         out = F.conv1d(qx, qw, self.old_module.bias, self.old_module.stride, self.old_module.padding, self.old_module.dilation, self.old_module.groups)
#         return self.output_quantizer(out)

# class QuantizedConv2d(QuantizationMixin):
#     def forward(self, x):
#         qx = self.input_quantizer(x)
#         qw = self.param_quantizers['weight'](self.old_module.weight)
#         out = F.conv2d(qx, qw, self.old_module.bias, self.old_module.stride, self.old_module.padding, self.old_module.dilation, self.old_module.groups)
#         return self.output_quantizer(out)

# class QuantizedConv3d(QuantizationMixin):
#     def forward(self, x):
#         qx = self.input_quantizer(x)
#         qw = self.param_quantizers['weight'](self.old_module.weight)
#         out = F.conv3d(qx, qw, self.old_module.bias, self.old_module.stride, self.old_module.padding, self.old_module.dilation, self.old_module.groups)
#         return self.output_quantizer(out)

# class QuantizedLinear(QuantizationMixin):
#     def forward(self, x):
#         qx = self.input_quantizer(x)
#         qw = self.param_quantizers['weight'](self.old_module.weight)
#         out = F.linear(qx, qw, self.old_module.bias)
#         return self.output_quantizer(out)

# # ======================================================================================
# # 4. ACTIVATION FUNCTIONS
# # ======================================================================================

# class QuantizedReLU(QuantizationMixin):
#     def forward(self, x):
#         # Note: In a fused block (Conv-BN-ReLU), the input 'x' is already quantized.
#         # Calling input_quantizer again is idempotent and harmless.
#         return self.output_quantizer(F.relu(self.input_quantizer(x)))

# class QuantizedGELU(QuantizationMixin):
#     def forward(self, x):
#         return self.output_quantizer(F.gelu(self.input_quantizer(x)))

# class QuantizedSiLU(QuantizationMixin):
#     def forward(self, x):
#         return self.output_quantizer(F.silu(self.input_quantizer(x)))

# # ======================================================================================
# # 5. POOLING AND PASSTHROUGH LAYERS (No quantization needed, just passthrough)
# # ======================================================================================

# class PassthroughWrapper(nn.Module):
#     """A simple wrapper for layers that don't need quantization logic."""
#     def __init__(self, old_module, **kwargs):
#         super().__init__()
#         self.old_module = old_module

#     def forward(self, x):
#         return self.old_module(x)

# QuantizedMaxPool2d = PassthroughWrapper
# QuantizedAdaptiveAvgPool2d = PassthroughWrapper
# QuantizedDropout = PassthroughWrapper
# QuantizedIdentity = PassthroughWrapper

# # ======================================================================================
# # 6. NORMALIZATION LAYERS
# # ======================================================================================

# #
# # !!! IMPORTANT NOTE ON BATCHNORM !!!
# #
# # A `QuantizedBatchNorm` is INTENTIONALLY OMITTED. During inference (and PTQ),
# # BatchNorm layers should be "fused" or "folded" into the preceding Conv/Linear
# # layer. You must perform this fusion on the FP32 model BEFORE applying these
# # quantization wrappers. Quantizing BatchNorm as a standalone module is a known
# # anti-pattern that severely degrades accuracy.
# #
# # Example of fusion:
# # >>> from torch.ao.quantization import fuse_modules
# # >>> model_fp32 = ...
# # >>> fuse_modules(model_fp32.conv1, model_fp32.bn1, inplace=True)
# #

# class QuantizedLayerNorm(QuantizationMixin):
#     def forward(self, x):
#         qx = self.input_quantizer(x)
#         qw = self.param_quantizers['weight'](self.old_module.weight)
#         # LayerNorm is unique; it uses the functional form with parameters
#         out = F.layer_norm(qx, self.old_module.normalized_shape, qw, self.old_module.bias, self.old_module.eps)
#         return self.output_quantizer(out)

# # ======================================================================================
# # 7. ELEMENT-WISE OPERATIONS
# # ======================================================================================

# class QuantizedAdd(QuantizationMixin):
#     """Wrapper for element-wise addition, crucial for residual connections."""
#     def __init__(self, old_module=nn.Identity(), activation_bits=8, **kwargs):
#         super().__init__(old_module, activation_bits=activation_bits)
#         # Need a second input quantizer
#         act_qmin, act_qmax, act_dtype = get_quant_config(bits=activation_bits, is_symmetric=False)
#         self.input_quantizer_2 = FakeQuantize(
#             observer=MovingAverageMinMaxObserver, quant_min=act_qmin, quant_max=act_qmax, dtype=act_dtype
#         )

#     def forward(self, x1, x2):
#         # NOTE: For perfect accuracy, both inputs should have the same scale/zero-point.
#         # This requires sharing observers, which adds complexity to the replacement logic.
#         qx1 = self.input_quantizer(x1)
#         qx2 = self.input_quantizer_2(x2)
#         return self.output_quantizer(torch.add(qx1, qx2))

# class QuantizedMul(QuantizationMixin):
#     """Wrapper for element-wise multiplication."""
#     def __init__(self, old_module=nn.Identity(), activation_bits=8, **kwargs):
#         super().__init__(old_module, activation_bits=activation_bits)
#         act_qmin, act_qmax, act_dtype = get_quant_config(bits=activation_bits, is_symmetric=False)
#         self.input_quantizer_2 = FakeQuantize(
#             observer=MovingAverageMinMaxObserver, quant_min=act_qmin, quant_max=act_qmax, dtype=act_dtype
#         )

#     def forward(self, x1, x2):
#         qx1 = self.input_quantizer(x1)
#         qx2 = self.input_quantizer_2(x2)
#         return self.output_quantizer(torch.mul(qx1, qx2))