import json
import os
from typing import Dict

import torch
import torch.fx as fx
import torch.nn as nn
from core.frameworks.pytorch.quantization.sdk.graph_utils import ModelGraphUtils
from core.frameworks.pytorch.quantization.sdk.model_preparer import (
    _get_module_for_dotted_name,
    prepare_model,
)
from core.frameworks.pytorch.quantization.sdk.modules.fake_quantize import (
    FakeQuantize as CustomFakeQuantize,
)
from core.frameworks.pytorch.quantization.sdk.modules.module_replacement_registry import (
    ModuleReplacements,
)
from core.frameworks.pytorch.quantization.sdk.modules.quantized_modules import (
    QuantizationMixin,
)
from core.frameworks.pytorch.quantization.sdk.quantization_utils import (
    build_clean_onnx_path,
    build_onnx_path_from_torch_name,
    get_bitwidth_from_string_dtype,
    normalize_class_name,
    print_quant_summary,
)
from core.frameworks.pytorch.utils.utils import is_leaf_module
from logger.logger_registry import get_logger
from utils.config_utils import (
    get_model_configs,
    get_task_configs,
)
from utils.helpers import set_nested_attr

from .modules.function_modules import REPLACEMENTS
from .quantsim_config.quantsim_parser import QuantSimConfigParser

logger = get_logger("main")


class QuantizerEnginer:
    def __init__(
        self,
        model: torch.nn.Module,
        configs,
        device: torch.device,
    ):
        self._model: torch.nn.Module = model
        self.config = configs
        self.logger = get_logger("main")
        self.device = device
        self.task_type = (
            "qft" if self.config["qft"] else "ptq" if self.config["ptq"] else "pruning"
        )

        self.quantsim_config = QuantSimConfigParser()

    def update_model(self, model: torch.nn.Module):
        self._model = model

    def fuse_all_conv_bn(self, module):
        fusion_list = []
        module_list = list(module.named_modules())

        for (name1, mod1), (name2, mod2) in zip(module_list, module_list[1:]):
            if isinstance(mod1, torch.nn.Conv2d) and isinstance(
                mod2, torch.nn.BatchNorm2d
            ):
                fusion_list.append([name1, name2])

        if fusion_list:
            torch.ao.quantization.fuse_modules(module, fusion_list, inplace=True)

    def _check_super_groups_module(self, layer_name: str) -> bool:
        super_groups = self.quantsim_config.get_super_groups_ops()
        next_node_name = self.graphutils.get_next_non_identity_module(layer_name)
        if not next_node_name:
            return self.quantsim_config.is_default_ops_output_quantized()
        layer_name = self.graphutils.get_normalized_layer_name(layer_name)
        next_node_name = self.graphutils.get_normalized_layer_name(next_node_name)
        if next_node_name:
            next_node_name_lower = next_node_name.lower()
            node_name_lower = layer_name.lower()
            for super_group in super_groups:
                super_group_first = super_group[0].lower()
                super_group_second = super_group[1].lower()
                if (
                    node_name_lower in super_group_first
                    and next_node_name_lower in super_group_second
                ):
                    return False
                elif next_node_name_lower in super_group_second:
                    return True
        return self.quantsim_config.is_default_ops_output_quantized()

    def apply_io_quantizer_flags_to_graph(
        self, graph: torch.fx.GraphModule
    ) -> dict[str, tuple[bool, bool]]:
        """
        Iterate over all nodes in the graph and determine IO quantizer flags.

        Returns:
            A dictionary mapping node names to (add_input_quantizer, add_output_quantizer) flags.
        """
        io_flags = {}
        for node in graph.graph.nodes:
            if node.op == "call_module":
                flags = self.__get_io_quantizer_flags(graph, node, None)
                io_flags[node.name] = flags
                print(
                    f"[IO Quantizer Flags] {node.name}: Input={flags[0]}, Output={flags[1]}"
                )
                module: QuantizationMixin = _get_module_for_dotted_name(
                    graph, node.target
                )
                if not isinstance(module, QuantizationMixin):
                    continue
                if not flags[0]:
                    module.disable_input_quantization()
                if not flags[1]:
                    module.disable_output_quantization()
                if flags[0]:
                    module.enable_input_quantization()
                if flags[1]:
                    module.enable_output_quantization()

        return io_flags

    def _check_super_groups_node(
        self,
        layer_name: str,
        node: torch.fx.Node,
        next_node: torch.fx.Node,
        next_node_name: str,
    ) -> bool:
        super_groups = self.quantsim_config.get_super_groups_ops()
        for super_group in super_groups:
            if (
                layer_name in super_group[0].lower()
                and next_node
                and next_node_name in super_group[1].lower()
            ):
                print("#" * 20, layer_name, next_node_name, "#" * 20)
                return False
            elif layer_name in super_group[1].lower():
                return True
        return True

    def __is_quantizable_module(self, module: torch.nn.Module):
        return ModuleReplacements.get_replacement(module) is not None

    @classmethod
    def __is_quantized_module(cls, module: torch.nn.Module):
        return isinstance(module, QuantizationMixin)

    def _normalized_op_name(self, module: torch.nn.Module) -> str:
        """Map a module to the normalized op name used in your config."""
        return self.graphutils.get_normalized_layer_name(module)

    def _match_pattern_from(
        self,
        graph: torch.fx.GraphModule,
        start_node: torch.fx.Node,
        pattern_ops: list[str],
    ):
        """
        Try to match a supergroup pattern starting at start_node.
        Returns list of nodes if matched, else None.
        """
        if start_node.op != "call_module":
            return None

        try:
            first_module = _get_module_for_dotted_name(graph, start_node.target)
        except Exception:
            return None

        first_name = self._normalized_op_name(first_module)
        if first_name != pattern_ops[0]:
            return None

        matched_nodes = [start_node]
        curr = start_node
        for expected in pattern_ops[1:]:
            next_node = self.graphutils.get_next_non_identity_node(curr)
            if not next_node or next_node.op != "call_module":
                return None
            try:
                next_module = _get_module_for_dotted_name(graph, next_node.target)
            except Exception:
                return None
            next_name = self._normalized_op_name(next_module)
            if next_name != expected:
                return None
            matched_nodes.append(next_node)
            curr = next_node

        return matched_nodes

    def _collect_supergroup_patterns(self):
        """
        Read patterns from your quantsim config.
        Expected structure: [["Conv", "BatchNorm", "Relu"], ["MatMul", "Add"]]
        """
        cfg = self.quantsim_config.get_model_quantization_config()
        patterns = cfg.get("supergroups", [])
        cleaned = []
        for pat in patterns:
            if isinstance(pat, (list, tuple)) and len(pat) >= 2:
                cleaned.append([str(x) for x in pat])
        return cleaned

    def _apply_super_groups_config(self, graph: torch.fx.GraphModule):
        """
        Find and apply super-group quantization configuration.
        Works for groups of any length >= 2. Prevents overlaps.
        """
        patterns = self._collect_supergroup_patterns()
        self._supergroup_members = {}
        claimed = set()
        group_id = 0

        for node in graph.graph.nodes:
            if node.op != "call_module" or node in claimed:
                continue

            for pat in patterns:
                match = self._match_pattern_from(graph, node, pat)
                if not match:
                    continue
                if any(n in claimed for n in match):
                    continue

                size = len(match)
                for idx, member in enumerate(match):
                    self._supergroup_members[member] = (group_id, idx, size)
                    claimed.add(member)

                # Gather actual modules
                modules = []
                for mnode in match:
                    m = _get_module_for_dotted_name(graph, mnode.target)
                    modules.append(m)

                # Apply quantizer sharing
                self._apply_super_group_action_general(modules)

                if getattr(self, "verbose", False):
                    names = [self._normalized_op_name(m) for m in modules]
                    print(f"[SuperGroup] id={group_id} matched: {' -> '.join(names)}")

                group_id += 1
                break

    def _belongs_to_super_group(self, node: torch.fx.Node) -> bool:
        return hasattr(self, "_supergroup_members") and node in self._supergroup_members

    def _supergroup_position(self, node: torch.fx.Node):
        """Return (group_id, idx, size) if node is in a supergroup, else None."""
        if self._belongs_to_super_group(node):
            return self._supergroup_members[node]
        return None

    def _apply_super_group_action_general(self, modules: list):
        n = len(modules)
        if n < 2:
            return

        last_module = modules[-1]
        if not hasattr(last_module, "output_quantizer"):
            return

        shared_output_q = last_module.output_quantizer

        for m in modules[:-1]:
            if hasattr(m, "output_quantizer"):
                m.output_quantizer = shared_output_q

    def __get_io_quantizer_flags(
        self, graph, node: torch.fx.Node, layer_name: str | None
    ) -> tuple[bool, bool]:
        """
        Decide whether to add input/output quantizers for a node,
        respecting model IO policy and per-layer overrides.
        """
        layer_name = layer_name or self.graphutils.get_normalized_layer_name(node)
        model_quant_config = self.quantsim_config.get_model_quantization_config()

        add_input_quantizer = self.quantsim_config.is_default_ops_input_quantized()
        add_output_quantizer = self.quantsim_config.is_default_ops_output_quantized()

        if self.graphutils._is_first(node, layer_name):
            add_input_quantizer = model_quant_config.get("input_quantized", False)
        if self.graphutils._is_last(node, layer_name):
            add_output_quantizer = model_quant_config.get("output_quantized", False)

        # No need for explicit supergroup override:
        # all sharing is already handled in _apply_super_group_action_general.

        add_input_quantizer = self._apply_layer_override(
            layer_name, add_input_quantizer, is_input=True
        )
        add_output_quantizer = self._apply_layer_override(
            layer_name, add_output_quantizer, is_input=False
        )
        _module = _get_module_for_dotted_name(graph, node.target)
        return add_input_quantizer, add_output_quantizer

    def _apply_layer_override(
        self, layer_name: str, current_value: bool, is_input: bool
    ) -> bool:
        if is_input:
            if (
                layer_name
                in self.quantsim_config.get_layers_to_skip_from_input_quantizers()
            ):
                return False
            if layer_name in self.quantsim_config.get_layers_to_add_input_quantizers():
                return True
        else:
            if (
                layer_name
                in self.quantsim_config.get_layers_to_skip_from_output_quantizers()
            ):
                return False
            if layer_name in self.quantsim_config.get_layers_to_add_output_quantizers():
                return True
        return current_value

    def _add_quantization_wrappers(self, module, prefix=""):

        if self.__is_quantized_module(module):
            return
        for module_name, module_ref in module.named_children():
            full_name = f"{prefix}.{module_name}" if prefix else module_name
            self.logger.info("nn.Module found : %s", module_ref)
            print("nn.Module found : %s", module_ref)

            if self.__is_quantizable_module(module_ref) and is_leaf_module(module_ref):
                quantized_module = self._create_quantizer_module(module_ref, full_name)
                if not quantized_module:
                    self.logger.info(f"Please register {full_name}")
                    continue
                setattr(module, module_name, quantized_module)
            else:
                self._add_quantization_wrappers(module_ref, prefix=full_name)

    def _create_quantizer_module(
        self, module_to_quantize: torch.nn.Module, module_name: str
    ):

        param_per_channel = get_task_configs(
            self.config, "ptq", "parameter_per_channel", False
        )
        act_per_channel = get_task_configs(
            self.config, "ptq", "activation_per_channel", False
        )
        param_per_tensor = get_task_configs(
            self.config, "ptq", "parameter_per_tensor", True
        )
        act_per_tensor = get_task_configs(
            self.config, "ptq", "activation_per_tensor", True
        )
        param_is_symmetric = get_task_configs(
            self.config, "ptq", "parameter_is_symmetric", True
        )
        act_is_symmetric = get_task_configs(
            self.config, "ptq", "activation_is_symmetric", False
        )

        global_activation_dtype = get_task_configs(
            self.config, "ptq", "activation_dtype", "int4"
        )
        global_param_dtype = get_task_configs(
            self.config, "ptq", "parameter_dtype", "int4"
        )
        global_activation_observer = get_task_configs(
            self.config, "ptq", "activation_observer"
        )
        global_weight_observer = get_task_configs(self.config, "ptq", "weight_observer")

        quantizer = ModuleReplacements.get_replacement(module_to_quantize)

        if not quantizer:
            self.logger.info(f"Please register {type(module_to_quantize)}")
            return None

        # registering parameter and activation dtype
        setattr(module_to_quantize, "activation_dtype", global_activation_dtype)
        setattr(module_to_quantize, "parameter_dtype", global_param_dtype)
        setattr(module_to_quantize, "parameter_observer", global_weight_observer)
        setattr(module_to_quantize, "activation_observer", global_activation_observer)
        setattr(module_to_quantize, "param_per_channel", param_per_channel)
        setattr(module_to_quantize, "act_per_channel", act_per_channel)
        setattr(module_to_quantize, "param_per_tensor", param_per_tensor)
        setattr(module_to_quantize, "act_per_tensor", act_per_tensor)
        setattr(module_to_quantize, "param_is_symmetric", param_is_symmetric)
        setattr(module_to_quantize, "act_is_symmetric", act_is_symmetric)

        self.logger.info(
            f"Replacing {type(module_to_quantize)} with {quantizer.__name__} for quantization"
        )
        quantized_module = quantizer(
            _module_to_wrap=module_to_quantize,
            # add_input_quantizers=add_input_quantizers,
            # add_output_quantizers=add_output_quantizers,
        )
        return quantized_module

    def prepare(
        self,
    ):
        """
        Recursively replace every weight-bearing layer with a QuantWrapper.

        Args:
            layer_types: if provided, only wrap these types; else wrap all with 'weight'.
            model:       internal use for recursion (initially None → uses self._model).
        """
        self.logger.info("=" * 60)
        self.logger.info("Preparing model for QFT (Quantization-Aware Fine-Tuning)")
        self.logger.info("=" * 60)
        self._model.eval()
        self.fuse_all_conv_bn(self._model)

        # try:
        self.traced_graph = fx.symbolic_trace(self._model)
        self.graphutils = ModelGraphUtils(self._model, self.traced_graph)
        self._model = prepare_model(self._model)
        self.graphutils.update_model(self._model)
        self._add_quantization_wrappers(self._model)
        self.graphutils.update_graph(self._model)
        self.apply_io_quantizer_flags_to_graph(self._model)
        self.logger.info(self._model)

        # except Exception as e:
        #     print(e)
        #     self.graphutils = ModelGraphUtils(self._model, None)
        #     self.logger.info("Model is not graph tracable. Cannot replace math ops.")
        #     self._add_quantization_wrappers(self._model)

        self.logger.info("Model after preparing")
        self.logger.info(self._model)
        print_quant_summary(self._model)
        return self._model

    def convert(self, model) -> None:
        """
        Convert the model to a quantized model.
        This method is called after the training and preparation steps.
        """

        for name, child in list(model.named_children()):
            if isinstance(child, QuantizationMixin):
                self.logger.info(f"Replacing Quantized module: {name}")
                module: QuantizationMixin = child
                scale_fp = module.get_scale_fp()
                wrapped_module: nn.Module = module._module_to_wrap
                for k, v in scale_fp.items():
                    wrapped_module.register_buffer(f"{k}_scale", v["scale"])
                    wrapped_module.register_buffer(f"{k}_zero_point", v["zero_point"])
                    setattr(wrapped_module, f"{k}_scale", v["scale"])
                    setattr(wrapped_module, f"{k}_zero_point", v["zero_point"])

                setattr(model, name, wrapped_module)
                continue

            self.convert(child)

    def export_model(self, model: torch.nn.Module, task_type: str) -> None:
        """
        Export the quantized model to a format suitable for deployment.
        """
        self.model_name = get_model_configs(self.config)["name"]
        export_format = self.config["export"]["format"]
        output_dir = self.config["export"]["output_dir"] + "/" + self.model_name
        if export_format == "onnx":
            export_model = model.apply(torch.ao.quantization.disable_observer)
            export_model.cpu()
            if not os.path.exists(output_dir):
                os.makedirs(output_dir, exist_ok=True)

            param_dtype = get_task_configs(
                self.config, task_type, "parameter_dtype", "int4"
            )
            opset = get_task_configs(self.config, "export", "opset", 13)
            attribute = get_task_configs(self.config, "export", "attribute", None)
            if attribute is not None:
                if not hasattr(export_model, attribute):
                    raise ValueError(f"Model has no attribute '{attribute}' for export")
                export_model = getattr(export_model, attribute)

            onnx_file = (
                f"{self.model_name}_quantized_model_{task_type}_{param_dtype}.onnx"
            )

            output_path = os.path.join(output_dir, onnx_file)

            self.logger.info("=" * 60)
            self.logger.info("Exporting the quantized model")
            self.logger.info("=" * 60)
            self.logger.info(f"Model Name      : {self.model_name}")
            self.logger.info(f"Export Format   : {export_format}")
            self.logger.info(f"Output Path     : {output_path}")
            self.logger.info("=" * 60)
            try:
                torch.onnx.export(
                    export_model,
                    torch.randn(self._input_shape),  # type: ignore
                    output_path,
                    export_params=True,
                    opset_version=opset,
                    do_constant_folding=True,
                    input_names=["input_image"],
                    output_names=["output"],
                )
                self.logger.info("Model export completed successfully.")
            except Exception as e:
                self.logger.error("Model export failed.")
                self.logger.exception(e)

    def __extract_encoding(self, module: CustomFakeQuantize) -> Dict:
        scale = module.scale
        zero_point = module.zero_point
        quant_min = module.quant_min
        quant_max = module.quant_max
        qscheme = module.qscheme
        dtype = str(module.dtype)

        bitwidth = get_bitwidth_from_string_dtype(dtype)

        is_symmetric = qscheme in [
            torch.per_tensor_symmetric,
            torch.per_channel_symmetric,
        ]

        if is_symmetric:
            encoding_min = -scale * ((quant_max - quant_min) / 2)
            encoding_max = scale * ((quant_max - quant_min) / 2)
        else:
            encoding_min = scale * (quant_min - zero_point)
            encoding_max = scale * (quant_max - zero_point)

        base_info = {
            "bitwidth": bitwidth,
            "quant_min": quant_min,
            "quant_max": quant_max,
            "qscheme": str(qscheme),
            "dtype": dtype,
            "is_symmetric": is_symmetric,
        }

        if scale.numel() == 1:
            return {
                **base_info,
                "encodings": [
                    {
                        "scale": scale.item(),
                        "offset": zero_point.item(),
                        "min": encoding_min.item(),
                        "max": encoding_max.item(),
                    }
                ],
            }

        else:
            encodings = []
            for i in range(scale.numel()):
                encodings.append(
                    {
                        "scale": scale[i].item(),
                        "offset": zero_point[i].item(),
                        "min": encoding_min[i].item(),
                        "max": encoding_max[i].item(),
                    }
                )
        return {**base_info, "encodings": encodings}

    def __get_quant_min_max(self, dtype):
        dtype = dtype.lower()
        if dtype == "int4":
            return -8, 7, torch.qint8
        elif dtype == "uint4":
            return 0, 15, torch.quint8
        elif dtype == "int8":
            return -128, 127, torch.qint8
        elif dtype == "uint8":
            return 0, 255, torch.quint8
        elif dtype == "int16":
            return -32768, 32767, torch.qint32
        #     return 0, 65535, torch.qint32
        else:
            raise ValueError(f"Unsupported dtype: {dtype}")

    def __get_encodings_from_model(self, model, encoding_dict=None) -> Dict[str, Dict]:
        """
        Extracts quantization parameters (scale, zero_point, min, max, etc.)
        from a model prepared using torch.ao.quantization.prepare_qat.

        Returns a hierarchical dictionary of encodings per FakeQuantize module.
        """
        encoding_dict = {}

        for name, child in model.named_modules():
            if not name:
                continue
            if "quant" in name:
                continue
            onnx_path = build_onnx_path_from_torch_name(name)
            cls_name = normalize_class_name(child.__class__.__name__)
            output_name = build_clean_onnx_path(f"{onnx_path}/{cls_name}")
            output_name = output_name.replace("Quantized", "")
            if "module_" in output_name:
                output_name.replace("module_", "")
            if isinstance(child, QuantizationMixin):
                data = {}
                for sub_name, sub_module in child.named_modules():
                    if isinstance(sub_module, (CustomFakeQuantize)):
                        encoding_info = self.__extract_encoding(sub_module)
                        data[f"{sub_name.replace('_quantizers','')}"] = encoding_info
                encoding_dict[output_name] = data

        return encoding_dict

    def generate_embeddings(self, attribute):
        """
        Iterates through model modules, collects quantization encodings from QuantizationMixin modules,
        and writes them to a JSON file as a list.
        """
        output_dir = (
            self.config["export"]["output_dir"]
            + "/"
            + get_model_configs(self.config, "name")
        )
        # if attribute is not None:
        #     if not hasattr(model, attribute):
        #         raise ValueError(f"Model has no attribute '{attribute}' for export")
        #     model = getattr(model, attribute)

        all_encodings = self.__get_encodings_from_model(
            self._model
            # if  isinstance(attribute, str) and not hasattr(self._model, attribute)
            # else  isinstance(attribute, str) and getattr(self._model, attribute)
        )

        task_type = (
            "qft"
            if get_model_configs(self.config, "qft")
            else "ptq" if get_model_configs(self.config, "ptq") else "pruning"
        )
        output_path = os.path.join(
            output_dir,
            f'{get_model_configs(self.config, "name")}_{task_type}_quantization_encodings.json',
        )
        os.makedirs(output_dir, exist_ok=True)

        with open(output_path, "w") as f:
            json.dump(all_encodings, f, indent=4)

        print(f"Quantization encodings saved to: {output_path}")
Update all method according to my quantsimparse
ignore aimet
i have sent one quantsimparser class right
some methods that you have used here is wrong