"""Logical instruction implementations for the 4.8.8 color code.

This module provides logical instruction implementations for the 4.8.8 color code,
including syndrome extraction, error correction procedures, and logical
measurement operations specific to the color code architecture.
"""

# Copyright 2018 The PECOS Developers
# Copyright 2018 National Technology & Engineering Solutions of Sandia, LLC (NTESS). Under the terms of Contract
# DE-NA0003525 with NTESS, the U.S. Government retains certain rights in this software.
#
# 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 typing import TYPE_CHECKING

from pecos.circuits.quantum_circuit import QuantumCircuit
from pecos.qeccs.default_logical_instruction import DefaultLogicalInstruction
from pecos.qeccs.helper_functions import pos2qudit
from pecos.typing import QECCInstrParams

if TYPE_CHECKING:
    from pecos.protocols import QECCProtocol


class InstrSynExtraction(DefaultLogicalInstruction):
    """Instruction for a round of syndrome extraction.

    Parent class sets self.qecc.
    """

    def __init__(
        self,
        qecc: "QECCProtocol",
        symbol: str,
        **params: QECCInstrParams,
    ) -> None:
        """Initialize the syndrome extraction instruction for Color 4.8.8 code.

        Args:
            qecc: The parent QECC instance.
            symbol: The instruction symbol identifier.
            **params: Additional instruction parameters.
        """
        super().__init__(qecc, symbol, **params)

        self.abstract_circuit = QuantumCircuit(**params)

        self.ancilla_x_check = set()
        self.ancilla_z_check = set()

        # Go through the ancillas and grab the data qubits that are on either side of it.

        self.pos2qudit = pos2qudit(qecc.layout)

        for q in sorted(self.ancilla_qudit_set):
            self._create_checks(q)

        # Determine the logical operations
        # --------------------------------
        z_qudits = set(qecc.sides["bottom"])
        x_qudits = set(qecc.sides["bottom"])

        logical_ops = [  # Each element in the list corresponds to a logical qubit
            # The keys label the type of logical operator
            {
                "X": QuantumCircuit([{"X": x_qudits}]),
                "Z": QuantumCircuit([{"Z": z_qudits}]),
            },
        ]

        self.initial_logical_ops = logical_ops
        self.final_logical_ops = logical_ops

        self.logical_signs = None
        self.logical_stabilizers = None

        # Must be called at the end of initiation.
        self._compile_circuit(self.abstract_circuit)

    def _create_checks(self, ancilla: int) -> None:
        self.ancilla_x_check.add(ancilla)
        self.ancilla_z_check.add(ancilla)

        # look up ancilla location
        x, y = self.qecc.layout[ancilla]

        square = []
        octagon = []

        square.extend(
            (
                self.pos2qudit.get((x - 1, y + 1)),
                self.pos2qudit.get((x + 1, y + 1)),
                self.pos2qudit.get((x - 1, y - 1)),
                self.pos2qudit.get((x + 1, y - 1)),
            ),
        )

        found_square = False

        for d in square:
            if d is not None:
                found_square = True
                break

        if found_square:
            self.abstract_circuit.append(
                "X check",
                polygon="square",
                locations={ancilla},
                datas=square,
            )
            self.abstract_circuit.append(
                "Z check",
                polygon="square",
                locations={ancilla},
                datas=square,
            )

        else:
            if y != 0:
                octagon.extend(
                    (
                        self.pos2qudit.get((x - 1, y + 3)),
                        self.pos2qudit.get((x + 1, y + 3)),
                        self.pos2qudit.get((x - 3, y + 1)),
                        self.pos2qudit.get((x + 3, y + 1)),
                        self.pos2qudit.get((x - 3, y - 1)),
                        self.pos2qudit.get((x + 3, y - 1)),
                        self.pos2qudit.get((x - 1, y - 3)),
                        self.pos2qudit.get((x + 1, y - 3)),
                    ),
                )
            else:
                octagon.extend(
                    (
                        self.pos2qudit.get((x - 1, y + 2)),
                        self.pos2qudit.get((x + 1, y + 2)),
                        self.pos2qudit.get((x - 3, y)),
                        self.pos2qudit.get((x + 3, y)),
                    ),
                )
                octagon.extend([None, None, None, None])

            self.abstract_circuit.append(
                "X check",
                polygon="octagon",
                locations={ancilla},
                datas=octagon,
            )
            self.abstract_circuit.append(
                "Z check",
                polygon="octagon",
                locations={ancilla},
                datas=octagon,
            )


class InstrInitZero(DefaultLogicalInstruction):
    """Instruction for initializing a logical zero.

    It is just like syndrome extraction except the data qubits are initialized in the zero state at tick = 0.

    `ideal_meas` == True will cause the measurements to be replace with ideal measurements.

    Parent class sets self.qecc.
    """

    def __init__(
        self,
        qecc: "QECCProtocol",
        symbol: str,
        **params: QECCInstrParams,
    ) -> None:
        """Initialize the logical zero state preparation instruction.

        Initializes all data qubits in the |0⟩ state followed by syndrome extraction.

        Args:
            qecc: The parent QECC instance.
            symbol: The instruction symbol identifier.
            **params: Additional instruction parameters including:
                - ideal_meas: If True, measurements are replaced with ideal measurements.
        """
        super().__init__(qecc, symbol, **params)

        self.symbol = "instr_init_zero"

        self.data_qudit_set = self.qecc.data_qudit_set
        self.ancilla_qudit_set = self.qecc.ancilla_qudit_set

        # This is basically syndrome extraction round where all the data qubits are initialized to zero.
        syn_ext = qecc.instruction("instr_syn_extract", **params)

        # Make a shallow copy of the abstract circuits.
        self.abstract_circuit = syn_ext.abstract_circuit.copy()
        self.abstract_circuit.params.update(params)

        self.ancilla_x_check = syn_ext.ancilla_x_check
        self.ancilla_z_check = syn_ext.ancilla_z_check

        data_qudits = syn_ext.data_qudit_set
        self.abstract_circuit.append("init |0>", locations=data_qudits, tick=0)

        self.initial_logical_ops = [  # Each element in the list corresponds to a logical qubit
            # The keys label the type of logical operator
            {"X": None, "Z": None},  # None => can be anything
        ]

        # Special for state initialization:
        # ---------------------------------
        # list of tuples of logical check and delogical stabilizer for each logical qudit.
        self.final_logical_ops = [
            {
                "X": QuantumCircuit([{"Z": set(qecc.sides["bottom"])}]),
                "Z": QuantumCircuit([{"X": set(qecc.sides["bottom"])}]),
            },
        ]

        # List of corresponding logical sign. (The logical sign if the instruction is preformed ideally.)
        self.logical_signs = [0]
        self.logical_stabilizers = ["Z"]
        # ---------------------------------

        # Must be called at the end of initiation.
        self._compile_circuit(self.abstract_circuit)


class InstrInitPlus(DefaultLogicalInstruction):
    """Instruction for initializing a logical plus.

    It is just like syndrome extraction except the data qubits are initialized in the plus state at tick = 0.

    `ideal_meas` == True will cause the measurements to be replace with ideal measurements.

    Parent class sets self.qecc.
    """

    def __init__(
        self,
        qecc: "QECCProtocol",
        symbol: str,
        **params: QECCInstrParams,
    ) -> None:
        """Initialize the logical plus state preparation instruction.

        Initializes all data qubits in the |+⟩ state followed by syndrome extraction.

        Args:
            qecc: The parent QECC instance.
            symbol: The instruction symbol identifier.
            **params: Additional instruction parameters including:
                - ideal_meas: If True, measurements are replaced with ideal measurements.
        """
        super().__init__(qecc, symbol, **params)

        self.symbol = "instr_init_plus"

        self.data_qudit_set = self.qecc.data_qudit_set
        self.ancilla_qudit_set = self.qecc.ancilla_qudit_set

        # This is basically syndrome extraction round where all the data qubits are initialized to plus.
        syn_ext = qecc.instruction("instr_syn_extract", **params)

        # Make a shallow copy of the abstract circuits.
        self.abstract_circuit = syn_ext.abstract_circuit.copy()
        self.abstract_circuit.params.update(params)

        self.ancilla_x_check = syn_ext.ancilla_x_check
        self.ancilla_z_check = syn_ext.ancilla_z_check

        data_qudits = syn_ext.data_qudit_set
        self.abstract_circuit.append("init |0>", locations=data_qudits, tick=0)
        self.abstract_circuit.append("H", locations=data_qudits, tick=1)

        self.initial_logical_ops = [  # Each element in the list corresponds to a logical qubit
            # The keys label the type of logical operator
            {"X": None, "Z": None},  # None => can be anything
        ]

        # Special for state initialization:
        # ---------------------------------
        # list of tuples of logical check and delogical stabilizer for each logical qudit.
        self.final_logical_ops = [
            {
                "X": QuantumCircuit([{"X": set(qecc.sides["bottom"])}]),
                "Z": QuantumCircuit([{"Z": set(qecc.sides["bottom"])}]),
            },
        ]

        # List of corresponding logical sign. (The logical sign if the instruction is preformed ideally.)
        self.logical_signs = [0]
        self.logical_stabilizers = ["X"]
        # ---------------------------------

        # Must be called at the end of initiation.
        self._compile_circuit(self.abstract_circuit)
