"""This module implements the DiscreteLayerGenerator class."""
from __future__ import annotations

import logging
from typing import Callable
from typing import Sequence

from bqskit.compiler.passdata import PassData
from bqskit.ir.circuit import Circuit
from bqskit.ir.gate import Gate
from bqskit.ir.gates.constant.cx import CNOTGate
from bqskit.ir.gates.constant.h import HGate
from bqskit.ir.gates.constant.t import TGate
from bqskit.ir.gates.parameterized.pauliz import PauliZGate
from bqskit.ir.operation import Operation
from bqskit.passes.search.generator import LayerGenerator
from bqskit.qis.state.state import StateVector
from bqskit.qis.state.system import StateSystem
from bqskit.qis.unitary.unitarymatrix import UnitaryMatrix
from bqskit.utils.typing import is_sequence


_logger = logging.getLogger(__name__)


class DiscreteLayerGenerator(LayerGenerator):
    """
    The DiscreteLayerGenerator class.

    Expands circuits using only discrete gates. This is a non-reinforcement
    learning version of diagonalizing in
    https://arxiv.org/abs/2409.00433.
    """

    def __init__(
        self,
        gateset: Sequence[Gate] = [HGate(), TGate(), CNOTGate()],
        double_headed: bool = False,
        dividing_gate_type: Callable[[int], Gate] = PauliZGate,
    ) -> None:
        """
        Construct a DiscreteLayerGenerator.

        Args:
            gateset (Sequence[Gate]): A sequence of gates that can be used
                in the output circuit. These must be non-parameterized gates.
                (Default: [HGate, TGate, CNOTGate])

            double_headed (bool): If True, successors will be generated by
                both appending and prepending gates. This lets unitaries be
                diagonalized instead of inverted. (Default: False)

            dividing_gate (Callable[[int], Gate]): A gate that goes between
                the two heads of the discrete searches. If double_headed is
                False, this gate simply goes at the beggining of the circuit.
                (Default: PauliZGate)

        Raises:
            ValueError: If the gateset is not a sequence.

            ValueError: If the gateset contains a parameterized gate.

            ValueError: If the radices of gates are different.

        TODO:
            Check universality of gateset.
        """
        if not is_sequence(gateset):
            m = f'Expected sequence of gates, got {type(gateset)}.'
            raise ValueError(m)

        radix = gateset[0].radixes[0]
        for gate in gateset:
            if gate.num_params > 0:
                m = 'Expected gate for constant gates, got parameterized'
                m += f' {gate} gate.'
                raise ValueError(m)
            for rad in gate.radixes:
                if rad != radix:
                    m = f'Radix mismatch on gate: {gate}. '
                    m += f'Expected {radix}, got {rad}.'
                    raise ValueError(m)
        self.gateset = gateset
        self.double_headed = double_headed
        self.dividing_gate_type = dividing_gate_type

    def gen_initial_layer(
        self,
        target: UnitaryMatrix | StateVector | StateSystem,
        data: PassData,
    ) -> Circuit:
        """
        Generate the initial layer, see LayerGenerator for more.

        Raises:
            ValueError: If `target` has a radix mismatch with
                `self.initial_layer_gate`.
        """

        if not isinstance(target, (UnitaryMatrix, StateVector, StateSystem)):
            m = f'Expected unitary or state, got {type(target)}.'
            raise TypeError(m)

        for radix in target.radixes:
            if radix != self.gateset[0].radixes[0]:
                m = 'Radix mismatch between target and gateset.'
                raise ValueError(m)

        init_circuit = Circuit(target.num_qudits, target.radixes)

        if self.double_headed:
            n = target.num_qudits
            span = list(range(n))
            init_circuit.append_gate(self.dividing_gate_type(n), span)

        return init_circuit

    def cancels_something(
        self,
        circuit: Circuit,
        gate: Gate,
        location: tuple[int, ...],
    ) -> bool:
        """Ensure applying gate at location does not cancel a previous gate."""
        last_cycle = circuit.num_cycles - 1
        try:
            op = circuit.get_operation((last_cycle, location[0]))
            op_gate, op_location = op.gate, op.location
            if op_location == location and op_gate.get_inverse() == gate:
                return True
            return False
        except IndexError:
            return False

    def count_repeats(
        self,
        circuit: Circuit,
        gate: Gate,
        qudit: int,
    ) -> int:
        """Count the number of times the last gate is repeated on qudit."""
        count = 0
        for cycle in reversed(range(circuit.num_cycles)):
            try:
                op = circuit.get_operation((cycle, qudit))
                if op.gate == gate:
                    count += 1
                else:
                    return count
            except IndexError:
                continue
        return count

    def gen_successors(self, circuit: Circuit, data: PassData) -> list[Circuit]:
        """
        Generate the successors of a circuit node.

        Raises:
            ValueError: If circuit is a single-qudit circuit.
        """
        if not isinstance(circuit, Circuit):
            raise TypeError(f'Expected circuit, got {type(circuit)}.')

        if circuit.num_qudits < 2:
            raise ValueError('Cannot expand a single-qudit circuit.')

        # Get the coupling graph
        coupling_graph = data.connectivity

        # Generate successors
        successors = []
        hashes = set()
        singles = [gate for gate in self.gateset if gate.num_qudits == 1]
        multis = [gate for gate in self.gateset if gate.num_qudits > 1]

        def add_to_successors(circuit: Circuit) -> None:
            h = self.hash_circuit_structure(circuit)
            if h not in hashes:
                successors.append(circuit)
                hashes.add(h)

        for gate in singles:
            for qudit in range(circuit.num_qudits):
                if gate.radixes[0] != circuit.radixes[qudit]:
                    continue
                if self.cancels_something(circuit, gate, (qudit,)):
                    continue
                if isinstance(gate, TGate):
                    if self.count_repeats(circuit, TGate(), qudit) >= 7:
                        continue
                successor = circuit.copy()
                successor.append_gate(gate, [qudit])

                add_to_successors(successor)

                if self.double_headed:
                    successor = circuit.copy()
                    op = Operation(gate, [qudit])
                    successor.insert(0, op)
                    add_to_successors(successor)

        for gate in multis:
            for edge in coupling_graph:
                if self.cancels_something(circuit, gate, edge):
                    continue
                qudit_radixes = [circuit.radixes[q] for q in edge]
                if gate.radixes != qudit_radixes:
                    continue
                successor = circuit.copy()
                successor.append_gate(gate, edge)
                add_to_successors(successor)

                if self.double_headed:
                    successor = circuit.copy()
                    op = Operation(gate, edge)
                    successor.insert(0, op)
                    add_to_successors(successor)

        return successors

    def hash_circuit_structure(self, circuit: Circuit) -> int:
        hashes = []
        for op in circuit:
            hashes.append(hash(op))
        return hash(tuple(hashes))
