import onnx

from ..quant_utils import TENSOR_NAME_QUANT_SUFFIX, QuantizedValue, QuantizedValueType, attribute_to_kwarg, ms_domain
from .base_operator import QuantOperatorBase
from .qdq_base_operator import QDQOperatorBase


class QLinearWhere(QuantOperatorBase):
    def should_quantize(self):
        return True

    def quantize(self):
        node = self.node
        assert node.op_type == "Where"
        if not self.quantizer.force_quantize_no_input_check:
            self.quantizer.new_nodes += [node]
            return
        (
            data_found,
            output_scale_name,
            output_zp_name,
            _,
            _,
        ) = self.quantizer._get_quantization_params(node.output[0])
        (
            q_input_names,
            zero_point_names,
            scale_names,
            nodes,
        ) = self.quantizer.quantize_activation(node, [1, 2])
        if not data_found or q_input_names is None:
            return super().quantize()
        qlinear_output = node.output[0] + TENSOR_NAME_QUANT_SUFFIX
        qlinear_output_name = node.name + "_quant" if node.name else ""

        q_output = QuantizedValue(
            node.output[0],
            qlinear_output,
            output_scale_name,
            output_zp_name,
            QuantizedValueType.Input,
        )
        self.quantizer.quantized_value_map[node.output[0]] = q_output

        kwargs = {}
        for attribute in node.attribute:
            kwargs.update(attribute_to_kwarg(attribute))
        kwargs["domain"] = ms_domain

        qlwhere_inputs = [
            node.input[0],
            q_input_names[0],
            scale_names[0],
            zero_point_names[0],
            q_input_names[1],
            scale_names[1],
            zero_point_names[1],
            output_scale_name,
            output_zp_name,
        ]
        qlwhere_node = onnx.helper.make_node(
            "QLinearWhere", qlwhere_inputs, [qlinear_output], qlinear_output_name, **kwargs
        )

        self.quantizer.new_nodes += nodes
        self.quantizer.new_nodes += [qlwhere_node]


class QDQWhere(QDQOperatorBase):
    def quantize(self):
        node = self.node
        assert node.op_type == "Where"
        if self.quantizer.force_quantize_no_input_check:
            if not self.quantizer.is_tensor_quantized(node.input[1]):
                self.quantizer.quantize_activation_tensor(node.input[1])
            if not self.quantizer.is_tensor_quantized(node.input[2]):
                self.quantizer.quantize_activation_tensor(node.input[2])
            if not self.disable_qdq_for_node_output:
                for output in node.output:
                    self.quantizer.quantize_activation_tensor(output)
        elif (
            self.quantizer.is_tensor_quantized(node.input[1])
            and self.quantizer.is_tensor_quantized(node.input[2])
            and not self.disable_qdq_for_node_output
        ):
            for output in node.output:
                self.quantizer.quantize_activation_tensor(output)