# mypy: allow-untyped-defs
from collections import OrderedDict
from collections.abc import Callable
from typing import Any

import torch
from torch.ao.quantization.fx._equalize import EqualizationQConfig
from torch.ao.quantization.fx._model_report.detector import (
    DETECTOR_IS_POST_OBS_KEY,
    DETECTOR_OBS_ARGS_KEY,
    DETECTOR_OBS_TO_INSERT_KEY,
    DETECTOR_TARGET_NODE_KEY,
    DetectorBase,
    DetectorQConfigInfo,
)
from torch.ao.quantization.fx._model_report.model_report_visualizer import (
    ModelReportVisualizer,
)
from torch.ao.quantization.fx.graph_module import GraphModule
from torch.ao.quantization.observer import ObserverBase
from torch.ao.quantization.qconfig_mapping import QConfig, QConfigMapping


class ModelReport:
    r"""
    The ModelReport class aims to provide users an easy way to diagnose issues that they run into
    with their models. The class works with all traceable GraphModules to help diagnose issues,
    though the requirements on the type of model more-so depends on the specific report the user
    is trying to generate. With respect to the reports, the ModelReport class is initialized with
    a set of Detector classes, each of which generate reports on quantization configuration
    issues a use might have.

    Currently supports generating reports on:
    - Suggestions for per-channel vs. per-tensor quantization (nn.Module)
    - Suggestions for dynamic vs static quantization for linear layers (Graph Modules)
    - Suggestions for input-weight equalization for linear and conv layers (Graph Modules)
    - Suggestions for outlier detection for all layers (Graph Modules)

    The ModelReport class has the primary functionality of inserting observers (primarily the ModelReportObserver)
    where needed for each detector to gather the information it needs, and then after calibration, the ModelReport
    class compiles the report generated by each Detector class into a single report to return to the user. It also
    has the capability to remove all the observers it inserted as well.

    * :attr:`_model` The model we wish to generate the report for. Must be a traceable GraphModule

    * :attr:`_desired_report_detectors` The set of Detectors representing desired reports from the ModelReport class
        Make sure that these are all unique types of detectors [do not have more than 1 of the same class]

    * :attr:`_desired_detector_names` The set of detector names of the _desired_report_detectors.
        This set is generated by calling the get_detector_name() of each detector

    * :attr:`_detector_name_to_observer_fqns` The mapping from each detector to fqns of observers of interest
        The purpose of this is to keep track of what observers were inserted for each detector, so that they
        can be removed at the end if desired

    * :attr:`_prepared_flag` A boolean flag that keeps track of whether we have prepared the model or not
        This is to ensure we only insert observers once with the ModelReport instance

    * :attr:`_removed_observers` A boolean to track if we have removed observers already
        The purpose is to ensure we don't attempt to remove observers twice with the same ModelReport
        instance. This also allows the functionality where we can generate the report multiple times
        as long as we haven't removed the observers yet.

    Note:
        This class was initially designed to work with the Fx Graph Mode workflow in mind. However,
        full functionality is available as long as there is a traceable GraphModule that is being used.
        One method to get a traceable GraphModule without going through the Fx workflow is to use
        the QuantizationTracer class.

    General Flow for Fx workflow:
    1.) Initialize ModelReport object with reports of interest by passing in initialized detector objects and model
    2.) Prepare your model with prepare_fx
    3.) Call model_report.prepare_detailed_calibration to add relevant observers
    4.) Calibrate your model with data
    5.) Call model_report.generate_report on your model to generate report and optionally remove added observers
    Optional
        6.) Call model_report.generate_visualizer to get a ModelReportVisualizer instance
        7.) To help in parsing report information and debugging, view report info as a:
            - Table
            - Histogram
            - Line plot
    8.) Call model_report.generate_qconfigs to generate the qconfigs based on the report suggestions

    Example (with QuantizationTracer):
        >>> # xdoctest: +SKIP
        >>> # get the necessary qconfig
        >>> config = PrepareCustomConfig()
        >>> skipped_module_names, skipped_module_classes = (
        ...     get_skipped_module_name_and_classes(config, False)
        ... )

        >>> # initialize our model and get GraphModule
        >>> model = SomeModel()
        >>> tracer = QuantizationTracer(skipped_module_names, skipped_module_classes)
        >>> graph_module = GraphModule(model, tracer.trace(model))

        >>> # get our set of detectors and ModelReport instance
        >>> detector_set = set(
        ...     [
        ...         DynamicStaticDetector(tolerance=0.5),
        ...         InputWeightEqualizationDetector(ratio_threshold=0.7),
        ...     ]
        ... )
        >>> tracer_reporter = ModelReport(graph_module, tracer_detector_set)

        >>> # now we insert the observers and calibrate the model
        >>> tracer_model_with_observers = tracer_reporter.prepare_detailed_calibration()
        >>> for i in range(num_callibration_batches):
        >>>     example_input = get_callibration_input()
        >>>     tracer_model_with_observers(example_input)

        >>> # finally we generate the reports and optionally remove the observers we inserted
        >>> reports = tracer_reporter.generate_model_report(
        ...     remove_inserted_observers=True
        ... )

        >>> # Optional: we can generate the qconfig mapping based on the suggestions
        >>> qconfigs = model_report.generate_qconfig_mapping()

        >>> # Optional: we can generate the equalization mapping based on the suggestions
        >>> qconfigs = model_report.generate_equalization_mapping()

        >>> # Optional: we get a ModelReportVisualizer instance to do any visualizations desired
        >>> model_report_visualizer = tracer_reporter.generate_visualizer()

    """

    def __init__(self, model: GraphModule, desired_report_detectors: set[DetectorBase]):
        if len(desired_report_detectors) == 0:
            raise ValueError("Should include at least 1 desired report")

        # keep track of the model we wish to generate report for
        self._model: GraphModule = model

        # keep the reports private so they can't be modified
        self._desired_report_detectors = desired_report_detectors
        self._desired_detector_names = {
            detector.get_detector_name() for detector in desired_report_detectors
        }

        # keep a mapping of desired reports to observers of interest
        # this is to get the readings, and to remove them, can create a large set
        # this set can then be used to traverse the graph and remove added observers
        self._detector_name_to_observer_fqns: dict[str, set[str]] = {}

        # initialize each report to have empty set of observers of interest
        for desired_report in self._desired_detector_names:
            self._detector_name_to_observer_fqns[desired_report] = set()

        # flags to ensure that we can only prepare and remove observers once
        self._prepared_flag = False
        self._removed_observers = False

        # store the reports that we generated for visualization purposes
        # initially empty since no reports generated
        self._generated_reports: dict[str, dict] = {}

    def get_desired_reports_names(self) -> set[str]:
        """Returns a copy of the desired reports for viewing"""
        return self._desired_detector_names.copy()

    def get_observers_of_interest(self) -> dict[str, set[str]]:
        """Returns a copy of the observers of interest for viewing"""
        return self._detector_name_to_observer_fqns.copy()

    def prepare_detailed_calibration(self) -> GraphModule:
        r"""
        Takes in a graph model and inserts the following observers:
        - ModelReportObserver

        Each observer is inserted based on the desired_reports into the relevant locations

        Right now, each report in self._desired_detector_names has independent insertions
            However, if a module already has a Observer of the same type, the insertion will not occur
            This is because all of the same type of Observer collect same information, so redundant

        Returns the same GraphModule with the observers inserted
        """

        # if already prepared once, cannot prepare again
        if self._prepared_flag:
            raise ValueError(
                "Already ran preparing detailed calibration. Run the report generation next after calibration."
            )

        # loop through each detector, find where placements should be, and keep track
        insert_observers_fqns: dict[str, Any] = {}

        for detector in self._desired_report_detectors:
            # determine observer points for each detector
            obs_fqn_to_info = detector.determine_observer_insert_points(self._model)
            # map each insert point to the observer to use
            insert_observers_fqns.update(obs_fqn_to_info)
            # update the set of observers this report cares about
            self._detector_name_to_observer_fqns[detector.get_detector_name()] = set(
                obs_fqn_to_info.keys()
            )

        # now insert all the observers at their desired locations
        for observer_fqn in insert_observers_fqns:
            target_node = insert_observers_fqns[observer_fqn][DETECTOR_TARGET_NODE_KEY]
            insert_obs = insert_observers_fqns[observer_fqn][DETECTOR_OBS_TO_INSERT_KEY]
            insert_post = insert_observers_fqns[observer_fqn][DETECTOR_IS_POST_OBS_KEY]
            observer_args = insert_observers_fqns[observer_fqn][DETECTOR_OBS_ARGS_KEY]
            self._insert_observer_around_module(
                observer_fqn, target_node, insert_obs, observer_args, insert_post
            )

        self._prepared_flag = True

        return self._model

    def _insert_observer_around_module(
        self,
        obs_fqn: str,
        target_node: torch.fx.node.Node,
        obs_to_insert: ObserverBase,
        observer_args: tuple,
        insert_post: bool,
    ):
        r"""
        Helper function that inserts the observer into both the graph structure and the module of the model

        Args
            node_fqn (str): The fully qualified name of the observer we want to insert
            target_node (torch.fx.node.Node): The node in model we are inserting observers around
            obs_to_insert (ObserverBase): The observer we are inserting around target_node
            observer_args (Tuple): The arguments we want to pass into the observer
            insert_post (bool): whether this is meant to be a post observer for this node
        """
        # if we are inserting post, then our target node is the next node
        if insert_post:
            target_node = target_node.next

        with self._model.graph.inserting_before(target_node):
            self._model.add_submodule(obs_fqn, obs_to_insert)
            self._model.graph.create_node(
                op="call_module", target=obs_fqn, args=observer_args
            )

        # recompile model after inserts are made
        self._model.recompile()

    def _get_node_from_fqn(self, node_fqn: str) -> torch.fx.node.Node:
        r"""
        Takes in a node fqn and returns the node based on the fqn

        Args
            node_fqn (str): The fully qualified name of the node we want to find in model

        Returns the Node object of the given node_fqn otherwise returns None
        """
        node_to_return = None
        for node in self._model.graph.nodes:
            # if the target matches the fqn, it's the node we are looking for
            if node.target == node_fqn:
                node_to_return = node
                break

        if node_to_return is None:
            raise ValueError("The node_fqn is was not found within the module.")

        # assert for MyPy
        if not isinstance(node_to_return, torch.fx.node.Node):
            raise AssertionError("node_to_return must be a torch.fx.node.Node")

        return node_to_return

    def generate_model_report(
        self, remove_inserted_observers: bool
    ) -> dict[str, tuple[str, dict]]:
        r"""
        Generates all the requested reports.

        Note:
            You should have calibrated the model with relevant data before calling this

        The reports generated are specified by the desired_reports specified in desired_reports

        Can optionally remove all the observers inserted by the ModelReport instance

        Args:
            remove_inserted_observers (bool): True to remove the observers inserted by this ModelReport instance

        Returns a mapping of each desired report name to a tuple with:
            The textual summary of that report information
            A dictionary containing relevant statistics or information for that report

        Note:
            Throws exception if we try to generate report on model we already removed observers from
            Throws exception if we try to generate report without preparing for calibration
        """
        # if we haven't prepped model for calibration, then we shouldn't generate report yet
        if not self._prepared_flag:
            raise Exception(  # noqa: TRY002
                "Cannot generate report without preparing model for calibration"
            )

        # if we already removed the observers, we cannot generate report
        if self._removed_observers:
            raise Exception(  # noqa: TRY002
                "Cannot generate report on model you already removed observers from"
            )

        # keep track of all the reports of interest and their outputs
        reports_of_interest = {}

        for detector in self._desired_report_detectors:
            # generate the individual report for the detector
            report_output = detector.generate_detector_report(self._model)
            reports_of_interest[detector.get_detector_name()] = report_output

        # if user wishes to remove inserted observers, go ahead and remove
        if remove_inserted_observers:
            self._removed_observers = True
            # get the set of all Observers inserted by this instance of ModelReport
            all_observers_of_interest: set[str] = set()
            for desired_report in self._detector_name_to_observer_fqns:
                observers_of_interest = self._detector_name_to_observer_fqns[
                    desired_report
                ]
                all_observers_of_interest.update(observers_of_interest)

            # go through all_observers_of_interest and remove them from the graph and model
            for observer_fqn in all_observers_of_interest:
                # remove the observer from the model
                self._model.delete_submodule(observer_fqn)

                # remove the observer from the graph structure
                node_obj = self._get_node_from_fqn(observer_fqn)

                if node_obj:
                    self._model.graph.erase_node(node_obj)
                else:
                    raise ValueError("Node no longer exists in GraphModule structure")

            # remember to recompile the model
            self._model.recompile()

        # save the generated reports for visualization purposes
        saved_reports: dict[str, dict] = {
            report_name: report_tuple[1]
            for report_name, report_tuple in reports_of_interest.items()
        }

        self._generated_reports = saved_reports

        # return the reports of interest
        return reports_of_interest

    def _is_same_info_for_same_key(self, info_dict_a: dict, info_dict_b: dict) -> bool:
        r"""
        Takes in two dictionaries and ensures that any common keys between the two have the same
        values.

        Args:
            info_dict_a (Dict): First dictionary we wish to compare
            info_dict_b (Dict): Second dictionary we wish to compare

        Returns True if all shared keys have same values, false otherwise
        """
        # get the set of keys for both
        dict_a_keys: set = set(info_dict_a.keys())
        dict_b_keys: set = set(info_dict_b.keys())

        # get the insersection keys and check if same value for both dicts
        intersecting_keys: set = dict_a_keys.intersection(dict_b_keys)

        for key in intersecting_keys:
            dict_a_val = info_dict_a[key]
            dict_b_val = info_dict_b[key]

            # if it's a tensor we have to handle separately
            if type(dict_a_val) is torch.Tensor:
                # if dict_b_val not tensor, automatically false
                if (
                    type(dict_b_val) is not torch.Tensor
                    or sum(dict_a_val != dict_b_val) != 0
                ):
                    return False
            else:
                # for non-tensor vals
                if dict_a_val != dict_b_val:
                    return False

        # if no non matching shared keys found, return true
        return True

    def _reformat_reports_for_visualizer(self) -> OrderedDict:
        r"""
        Takes the generated reports and reformats them into the format that is desired by the
        ModelReportVisualizer

        Returns an OrderedDict mapping module_fqns to their features
        """
        # we want to reorder and reformat the information so it is ordered in terms of order
        # found in the model

        # first create new dict with all modules as keys and features under respective module
        module_fqns_to_features: dict[str, dict] = {}

        for report_name in self._generated_reports:
            # get mod -> feature dict and go through
            module_info = self._generated_reports[report_name]

            for module_fqn in module_info:
                # check if already in our accumulation dict
                if module_fqn in module_fqns_to_features:
                    # we merge all the features together
                    new_info: dict = module_info[module_fqn]
                    present_info: dict = module_fqns_to_features[module_fqn]

                    # merge them together into the new unioned dict
                    # same features keys -> same info, so okay if override

                    # do safety check to make sure shared keys have same info
                    if self._is_same_info_for_same_key(new_info, present_info):
                        module_fqns_to_features[module_fqn] = {
                            **new_info,
                            **present_info,
                        }
                    else:
                        error_str = "You have the same key with different values across detectors. "
                        error_str += "Someone incorrectly implemented a detector with conflicting keys to existing detectors."
                        raise ValueError(error_str)
                else:
                    # we just set it
                    module_fqns_to_features[module_fqn] = module_info[module_fqn]

        # our ordered dict so that modules can be ordered in order of how they appear in model
        features_by_module: OrderedDict[str, dict] = OrderedDict()

        # we loop through modules in graph in order
        for fqn, _module in self._model.named_modules():
            # find that fqn in fqns_to_features
            if fqn in module_fqns_to_features:
                # add it to our ordered dict
                features_by_module[fqn] = module_fqns_to_features[fqn]

        # return the ordered dict of info we created
        return features_by_module

    def generate_visualizer(self) -> ModelReportVisualizer:
        r"""
        Generates a ModelReportVisualizer instance using the reports generated
        by the generate_model_report() method.

        Returns the generated ModelReportVisualizer instance initialized

        Note:
            Throws exception if attempt to get visualizers without generating report
        """
        # check if user has generated reports at least once
        if len(self._generated_reports) == 0:
            raise Exception(  # noqa: TRY002
                "Unable to generate visualizers without first generating reports"
            )

        # get the ordered dict mapping modules to their full set of collected features / stats
        module_fqns_to_features: OrderedDict = self._reformat_reports_for_visualizer()

        # create and return ModelReportVisualizer instance
        visualizer: ModelReportVisualizer = ModelReportVisualizer(
            module_fqns_to_features
        )

        return visualizer

    def _generate_qconfig_mapping_helper(
        self,
        detector_qconfig_info_combined: dict[str, DetectorQConfigInfo],
        generation_function: Callable,
    ) -> QConfigMapping:
        r"""
        This helper takes in the compiled detector qconfig info that
        has been compiled together and merges it into a QConfigMapping
        """
        # keep track of the qconfigmapping
        qconfig_mapping = QConfigMapping()

        # loop through each module / fqn and attempt to create QConfigMapping
        for fqn, module in self._model.named_modules():
            # if we have a qconfig info for this module
            if fqn in detector_qconfig_info_combined:
                qconfig_info_compiled = detector_qconfig_info_combined[fqn]

                # now generate the qconfig and add it to the mapping
                generated_qconfig = generation_function(qconfig_info_compiled, module)

                # add to our config
                qconfig_mapping.set_module_name(fqn, generated_qconfig)

        # return compiled mapping
        return qconfig_mapping

    def _update_detector_quantizaiton_qconfig_info(
        self, combined_info: DetectorQConfigInfo, new_info: DetectorQConfigInfo
    ):
        r"""
        Takes in the old and new information and updates the combined information.

        Args:
            combined_info (DetectorQConfigInfo): The DetectorQConfigInfo we are compiling all of the information in
            new_info (DetectorQConfigInfo): The DetectorQConfigInfo with the information we are trying to merge the new info
                into it
        """
        combined_info.is_activation_dynamic = (
            combined_info.is_activation_dynamic or new_info.is_activation_dynamic
        )
        combined_info.is_weight_per_channel = (
            combined_info.is_weight_per_channel or new_info.is_weight_per_channel
        )

    def _update_detector_equalization_qconfig_info(
        self, combined_info: DetectorQConfigInfo, new_info: DetectorQConfigInfo
    ):
        r"""
        Takes in the old and new information and updates the combined information.

        Args:
            combined_info (DetectorQConfigInfo): The DetectorQConfigInfo we are compiling all of the information in
            new_info (DetectorQConfigInfo): The DetectorQConfigInfo with the information we are trying to merge the new info
                into it
        """
        is_equalization_recommended = (
            combined_info.is_equalization_recommended
            or new_info.is_equalization_recommended
        )
        combined_info.is_equalization_recommended = is_equalization_recommended

    def _generate_module_fqn_to_detector_info_mapping(
        self, update_qconfig_info_function: Callable
    ) -> dict[str, DetectorQConfigInfo]:
        r"""
        Generates a QConfigMapping based on the suggestions of the
        ModelReport API. The generated mapping encompasses all the
        different types of feedback from the different detectors
        all into one place.

        These configs are based on the suggestions provided by the ModelReport API
        and can only be generated once the reports have been generated.

        Args:
            update_qconfig_info_function (Callable) takes in a function that takes in two DetectorQConfigInfo
            and updates the one that is being compiled

        Returns a Dict mapping module_fqns to DetectorQConfigInfo objects

        Note:
            Throws exception if we try to generate mapping on model we already removed observers from
            Throws exception if we try to generate mapping without preparing for calibration
        """
        # if we haven't prepped model for calibration, then we shouldn't generate mapping yet
        if not self._prepared_flag:
            raise Exception(  # noqa: TRY002
                "Cannot generate report without preparing model for calibration"
            )

        # if we already removed the observers, we cannot mapping
        if self._removed_observers:
            raise Exception(  # noqa: TRY002
                "Cannot generate report on model you already removed observers from"
            )

        # keep track of qconfig info for each module across detectors
        detector_qconfig_info_combined: dict[str, DetectorQConfigInfo] = {}

        for detector in self._desired_report_detectors:
            # get the info from the detector
            detector_info: dict[str, DetectorQConfigInfo] = detector.get_qconfig_info(
                self._model
            )

            # we go through the modules
            for module_fqn in detector_info:
                # see if we already have info on it
                if module_fqn in detector_qconfig_info_combined:
                    # we combine the current options with what is there
                    current_options = detector_qconfig_info_combined[module_fqn]
                    detector_options = detector_info[module_fqn]

                    update_qconfig_info_function(current_options, detector_options)
                else:
                    # we just use this for now
                    detector_qconfig_info_combined[module_fqn] = detector_info[
                        module_fqn
                    ]

        return detector_qconfig_info_combined

    def generate_qconfig_mapping(self) -> QConfigMapping:
        r"""
        Generates a QConfigMapping based on the suggestions of the
        ModelReport API. The generated mapping encompasses all the
        different types of feedback from the different detectors
        all into one place.

        These configs are based on the suggestions provided by the ModelReport API
        and can only be generated once the reports have been generated.

        Returns a QConfigMapping for the quantization configuration

        Note:
            Throws exception if we try to generate mapping on model we already removed observers from
            Throws exception if we try to generate mapping without preparing for calibration
        """
        # get the mapping info
        detector_qconfig_info_combined = (
            self._generate_module_fqn_to_detector_info_mapping(
                self._update_detector_quantizaiton_qconfig_info
            )
        )

        # we will do a bit of processing and remove fqns that don't have input weight recommended

        # now we generate the QConfig for each of the options
        mapping: QConfigMapping = self._generate_qconfig_mapping_helper(
            detector_qconfig_info_combined, self._quantization_config_generator
        )

        # return the generated mapping
        return mapping

    def _quantization_config_generator(
        self, detector_qconfig_info: DetectorQConfigInfo, module: torch.nn.Module
    ) -> QConfig:
        r"""
        Returns the quantization configuration generated by the DetectorQConfigInfo object
        """
        return detector_qconfig_info.generate_quantization_qconfig(module)

    def _equalization_config_generator(
        self, detector_qconfig_info: DetectorQConfigInfo, module: torch.nn.Module
    ) -> EqualizationQConfig:
        r"""
        We ignore the module argument here, and only focus on thedetector_qconfig_info

        Returns the equalization configuration generated by the DetectorQConfigInfo object
        """
        return detector_qconfig_info.generate_equalization_qconfig()

    def generate_equalization_mapping(self) -> QConfigMapping:
        r"""
        Generates a QConfigMapping based on the suggestions of the
        ModelReport API for equalization. The generated mapping encompasses all the
        different types of feedback from the input-weight equalization detector.

        These configs are based on the suggestions provided by the ModelReport API
        and can only be generated once the reports have been generated.

        Returns a QConfigMapping for the equalization configuration
        """
        # get the mapping info
        detector_qconfig_info_combined = (
            self._generate_module_fqn_to_detector_info_mapping(
                self._update_detector_equalization_qconfig_info
            )
        )

        # now we generate the QConfig for each of the options
        mapping: QConfigMapping = self._generate_qconfig_mapping_helper(
            detector_qconfig_info_combined, self._equalization_config_generator
        )

        # return the generated mapping
        return mapping
