"""Bitflip noise with leakage for qubit initialization.

This module provides noise models for qubit initialization operations that
include both bitflip errors and leakage to states outside the computational
subspace, providing more realistic error modeling.
"""

# Copyright 2023 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

from pecos.error_models.noise_impl.noise_initz_bitflip import noise_initz_bitflip
from pecos.reps.pyphir.op_types import QOp

if TYPE_CHECKING:
    from pecos.protocols import MachineProtocol


def noise_initz_bitflip_leakage(op: QOp, p: float, machine: MachineProtocol) -> None:
    """The noise model for qubit (re)initialization, including leakage support.

    Args:
    ----
        op: Ideal quantum operation.
        p: Probability of bitflip.
        machine: Machine protocol instance containing leakage state information.
    """
    args = set(op.args)
    leaked = machine.leaked_qubits & args

    noise = []
    noisy_ops = machine.unleak(leaked)
    if noise:
        noise.extend(noisy_ops)

    not_leaked = args - leaked
    if not_leaked:
        remaining_inits = QOp(
            name=op.name,
            args=list(not_leaked),
            metadata=dict(op.metadata),
        )
        noisy_ops = noise_initz_bitflip(remaining_inits, p)
        if noisy_ops:
            noise.extend(noisy_ops)

    return noise
