"""Depolarizing error generation and application.

This module provides utilities for generating and applying depolarizing
errors to quantum circuits, including error rate calculations and
conditional error application based on circuit structure.
"""

# Copyright 2021 The PECOS Developers
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
# the License.You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
# specific language governing permissions and limitations under the License.

from __future__ import annotations

from typing import TYPE_CHECKING

import pecos as pc
from pecos.circuits import QuantumCircuit
from pecos.engines.cvm.classical import eval_condition
from pecos.error_models.class_errors_circuit import ErrorCircuits
from pecos.error_models.noise_impl_old.gate_groups import one_qubits, two_qubits
from pecos.error_models.noise_impl_old.init_noise import noise_init_bitflip
from pecos.error_models.noise_impl_old.meas_noise import noise_meas_bitflip
from pecos.error_models.noise_impl_old.memory_noise import noise_tq_mem
from pecos.error_models.noise_impl_old.sq_noise import noise_depolarizing_sq_gate
from pecos.error_models.noise_impl_old.tq_noise import (
    noise_two_qubit_gates_depolarizing_with_noiseless,
)
from pecos.error_models.parent_class_error_gen import ParentErrorModel

if TYPE_CHECKING:
    from pecos.typing import ErrorParams, GateParams, OutputDict


class DepolarizingErrorModel(ParentErrorModel):
    """Parameterized error model for Beta and ARC1."""

    def __init__(self) -> None:
        """Initialize a depolarizing error model.

        Sets up the error model with empty state for tracking qubits,
        error circuits, error parameters, and the quantum circuit.
        """
        super().__init__()

        self.qubit_set = set()

        self.error_circuits = None
        self.error_params = None
        self.circuit = None

    def scaling(self) -> None:
        """Apply scaling factors to error parameters."""
        if "p2_mem" not in self.error_params:
            self.error_params["p2_mem"] = None

        # conversion from average error to total error
        self.error_params["p1"] *= 3 / 2
        self.error_params["p2"] *= 5 / 4

        scale = self.error_params.get("scale", 1.0)
        self.error_params["p1"] *= scale
        self.error_params["p2"] *= scale

        if isinstance(self.error_params["p_meas"], tuple):
            self.error_params["p_meas"] = pc.mean(self.error_params["p_meas"])

        self.error_params["p_meas"] *= scale
        self.error_params["p_init"] *= scale

    def start(
        self,
        circuit: QuantumCircuit,
        error_params: ErrorParams,
        *,
        _reset_leakage: bool = True,  # Reserved for future leakage handling
    ) -> ErrorCircuits:
        """Initialize error generation for a quantum circuit.

        Args:
            circuit: Quantum circuit to generate errors for.
            error_params: Error parameters dictionary.
            _reset_leakage: Reserved for future leakage handling.

        Returns:
            Error circuits object containing generated errors.
        """
        self.qubit_set = set(range(circuit.metadata["num_qubits"]))

        self.error_circuits = ErrorCircuits()
        self.circuit = circuit

        # if error_params is None:

        if not error_params:
            msg = "Error params not set!"
            raise Exception(msg)

        self.error_params = {}
        self.error_params.update(error_params)

        self.scaling()  # scale self.error_params params

        if not self.error_params.get("noiseless_qubits"):
            self.error_params["noiseless_qubits"] = set()
        self.error_params["noiseless_qubits"] = set(
            self.error_params["noiseless_qubits"],
        )

        return self.error_circuits

    def reset(self) -> DepolarizingErrorModel:
        """Reset the error model to its initial state.

        Returns:
            New instance of DepolarizingErrorModel.
        """
        return DepolarizingErrorModel()

    def generate_tick_errors(
        self,
        tick_circuit: QuantumCircuit,
        time: int | tuple[int, ...],
        output: OutputDict | None = None,
        *,
        _reset_leakage: bool = False,  # Reserved for future leakage handling
        **_params: GateParams,  # Reserved for additional parameters
    ) -> ErrorCircuits:
        """The method that gets called each circuit tick to generate circuit noise for that tick."""
        # Get the tick
        tick_index = time[-1] if isinstance(time, tuple) else time

        circuit = tick_circuit.circuit

        before = QuantumCircuit()  # Faults that occur before the tick
        after = QuantumCircuit()  # Faults the occur after the tick
        remove_locations = set()  # gate locations to remove

        for symbol, locations, metadata in circuit.items(tick_index):
            cond = metadata.get("cond")

            if symbol == "cop":
                pass

            # ########################################
            # INITS WITH X NOISE
            elif symbol == "init |0>":
                noisy = set(locations) - self.error_params["noiseless_qubits"]
                noise_init_bitflip(noisy, after, "X", p=self.error_params["p_init"])

            # ########################################
            # ONE QUBIT GATES
            elif symbol in one_qubits:
                if eval_condition(cond, output):
                    noisy = set(locations) - self.error_params["noiseless_qubits"]
                    noise_depolarizing_sq_gate(noisy, after, p=self.error_params["p1"])

            # ########################################
            # TWO QUBIT GATES
            elif symbol in two_qubits:
                if eval_condition(cond, output):
                    noise_two_qubit_gates_depolarizing_with_noiseless(
                        locations,
                        after,
                        p=self.error_params["p2"],
                        noiseless_qubits=self.error_params["noiseless_qubits"],
                    )

                    if self.error_params["p2_mem"]:
                        noise_tq_mem(locations, after, p=self.error_params["p2_mem"])

            # ########################################
            # MEASURE X NOISE
            elif symbol == "measure Z":
                if eval_condition(cond, output):
                    noisy = set(locations) - self.error_params["noiseless_qubits"]
                    noise_meas_bitflip(
                        noisy,
                        metadata,
                        after,
                        p=self.error_params["p_meas"],
                    )

            elif symbol in {"repump", "leak", "unleak"}:
                pass

            else:
                msg = f"This error model doesn't handle gate: {symbol}!"
                raise Exception(msg)

        self.error_circuits.add_circuits(time, before, after, remove_locations)

        return self.error_circuits
