"""Guppy standard library for Quantinuum systems device operations."""

from typing import no_type_check

from guppylang_internals.decorator import custom_function, custom_type, hugr_op
from guppylang_internals.std._internal.compiler.qsystem import (
    ReadFutureBoolCompiler,
    future_bool_type,
)
from guppylang_internals.std._internal.compiler.quantum import (
    InoutMeasureResetCompiler,
)
from guppylang_internals.std._internal.compiler.tket_exts import (
    QSYSTEM_EXTENSION,
)
from guppylang_internals.std._internal.util import quantum_op

from guppylang import guppy
from guppylang.std.angles import angle, pi
from guppylang.std.array import array
from guppylang.std.builtins import owned
from guppylang.std.futures import Future
from guppylang.std.option import Option, nothing, some
from guppylang.std.quantum import qubit


@guppy
@no_type_check
def phased_x(q: qubit, angle1: angle, angle2: angle) -> None:
    r"""phased_x gate command.

    .. math::

        \mathrm{PhasedX}(\theta_1, \theta_2)=
          \mathrm{Rz(\theta_2)Rx(\theta_1)Rz(-\theta_2)}&=
          \begin{pmatrix}
          \cos(\frac{ \theta_1}{2}) &
            -i e^{-i \theta_2}\sin(\frac{\theta_1}{2})\\
          -i e^{i \theta_2}\sin(\frac{\theta_1}{2}) &
            \cos(\frac{\theta_1}{2})
           \end{pmatrix}
    """
    f1 = float(angle1)
    f2 = float(angle2)
    _phased_x(q, f1, f2)


@guppy
@no_type_check
def zz_max(q1: qubit, q2: qubit) -> None:
    r"""zz_max gate command. A maximally entangling zz_phase gate.

    This is a special case of the zz_phase command with :math:`\theta = \frac{\pi}{2}`.

    zz_max(q1, q2)

    Qubit ordering: [q1, q2]

    .. math::
        \mathrm{ZZMax}=
        \exp(\frac{- i\pi}{4}\big(Z \otimes Z \big))=
          \begin{pmatrix}
            e^{\frac{-i\pi}{4}} & 0 & 0 & 0 \\
            0 & e^{\frac{i\pi}{4}} & 0 & 0 \\
            0 & 0 & e^{\frac{i\pi}{4}} & 0 \\
            0 & 0 & 0 & e^{\frac{-i\pi}{4}}
        \end{pmatrix}
    """
    zz_phase(q1, q2, pi / 2)


@guppy
@no_type_check
def zz_phase(q1: qubit, q2: qubit, angle: angle) -> None:
    r"""zz_phase gate command.

    zz_phase(q1, q2, theta)

    Qubit ordering: [q1, q2]

    .. math::
        \mathrm{ZZPhase}(\theta)=
        \exp(\frac{- i \theta}{2}\big(Z \otimes Z \big))=
          \begin{pmatrix}
            e^{\frac{-i \theta}{2}} & 0 & 0 & 0 \\
            0 & e^{\frac{i \theta}{2}} & 0 & 0 \\
            0 & 0 & e^{\frac{i \theta}{2}} & 0 \\
            0 & 0 & 0 & e^{\frac{-i \theta}{2}}
        \end{pmatrix}

    >>> @guppy
    ... def qsystem_cx(q1: qubit, q2: qubit) -> None:
    ...     phased_x(angle(3/2), angle(-1/2), q2)
    ...     zz_phase(q1, q2, angle(1/2))
    ...     rz(angle(5/2), q1)
    ...     phasedx(angle(-3/2), q1)
    ...     rz(angle(3/2), q2)
    >>> qsystem_cx.compile()
    """
    f = float(angle)
    _zz_phase(q1, q2, f)


@guppy
@no_type_check
def rz(q: qubit, angle: angle) -> None:
    r"""rz gate command.

    .. math::
        \mathrm{Rz}(\theta)=
        \exp(\frac{- i \theta}{2} Z)=
          \begin{pmatrix}
            e^{\frac{-i \theta}{2}} & 0  \\
            0 & e^{\frac{i \theta}{2}}
        \end{pmatrix}
    """

    f1 = float(angle)
    _rz(q, f1)


@hugr_op(quantum_op("Measure", ext=QSYSTEM_EXTENSION))
@no_type_check
def measure(q: qubit @ owned) -> bool:
    """Measure a qubit destructively."""


@custom_function(InoutMeasureResetCompiler("MeasureReset", QSYSTEM_EXTENSION))
@no_type_check
def measure_and_reset(q: qubit) -> bool:
    """MeasureReset operation from the qsystem extension."""


@hugr_op(quantum_op("Reset", ext=QSYSTEM_EXTENSION))
@no_type_check
def reset(q: qubit) -> None:
    """Reset a qubit to the :math:`|0\\rangle` state."""


# TODO
# @hugr_op(quantum_op("TryQAlloc", ext=QSYSTEM_EXTENSION))
# @no_type_check
# def _try_qalloc() -> Option[qubit]:
#     ..


@hugr_op(quantum_op("QFree", ext=QSYSTEM_EXTENSION))
@no_type_check
def qfree(q: qubit @ owned) -> None: ...


@hugr_op(quantum_op("LazyMeasureLeaked", ext=QSYSTEM_EXTENSION))
@no_type_check
def _measure_leaked(q: qubit @ owned) -> Future[int]:
    """Measure the qubit or return 2 if it is leaked."""


@guppy
@no_type_check
def measure_leaked(q: qubit @ owned) -> "MaybeLeaked":
    """Measure the qubit and return a MaybeLeaked result."""
    fm = _measure_leaked(q)
    return MaybeLeaked(fm)


@guppy.struct
@no_type_check
class MaybeLeaked:
    """A class representing a measurement that may have leaked.

    This is used to represent the result of `measure_leaked`, which can either
    return a boolean measurement result or indicate that the qubit has leaked.
    """

    _measurement: Future[int]  # type: ignore[type-arg]

    @guppy
    @no_type_check
    def is_leaked(self: "MaybeLeaked") -> bool:
        """Check if the measurement indicates a leak."""
        return self._measurement.copy().read() == 2

    @guppy
    @no_type_check
    def to_result(self: "MaybeLeaked @ owned") -> Option[bool]:
        """Returns the measurement result or `nothing` if leaked."""
        int_value: int = self._measurement.read()
        if int_value == 2:
            return nothing()
        measurement = int_value == 1
        return some(measurement)

    @guppy
    @no_type_check
    def discard(self: "MaybeLeaked @ owned") -> None:
        self._measurement.discard()


@hugr_op(quantum_op("LazyMeasure", ext=QSYSTEM_EXTENSION))
@no_type_check
def lazy_measure(q: qubit @ owned) -> "Measurement":
    """Request a destructive lazy measurement of a qubit, returning a `Measurement`
    value. Call `.read()` on the value to block until the result is available.
    """


N = guppy.nat_var("N")


@guppy
@no_type_check
def lazy_measure_array(qubits: array[qubit, N] @ owned) -> array["Measurement", N]:
    """Request a destructive lazy measurement of an array of qubits, returning an array
    of `Measurement` values. Call `.read()` on each value to block until results are
    available.
    """
    return array(lazy_measure(q) for q in qubits)


@custom_type(
    future_bool_type(),
    copyable=False,
    droppable=False,
)
class Measurement:
    """Represents the result of a lazy measurement which needs to be explicitly read
    before being used."""

    @custom_function(compiler=ReadFutureBoolCompiler())
    @no_type_check
    def read(self: "Measurement" @ owned) -> bool:
        """Read the measurement result, consuming it. Blocks until the result is
        available if the measurement hasn't been performed yet since being requested.
        """

    @guppy
    @no_type_check
    def __consume_as_bool__(self: "Measurement" @ owned) -> bool:
        return self.read()


@guppy
@no_type_check
def collect_measurements(measurements: array[Measurement, N] @ owned) -> array[bool, N]:
    """Block on each measurement until it is available and collect results into an
    array of bools.
    """
    return array(m.read() for m in measurements)


# ------------------------------------------------------
# --------- Internal definitions -----------------------
# ------------------------------------------------------


@hugr_op(quantum_op("PhasedX", ext=QSYSTEM_EXTENSION))
@no_type_check
def _phased_x(q: qubit, angle1: float, angle2: float) -> None:
    """PhasedX operation from the qsystem extension.

    See ``guppylang.std.qsystem.phased_x`` for a public definition that
    accepts angle parameters.
    """


@hugr_op(quantum_op("ZZPhase", ext=QSYSTEM_EXTENSION))
@no_type_check
def _zz_phase(q1: qubit, q2: qubit, angle: float) -> None:
    """ZZPhase operation from the qsystem extension.

    See ``guppylang.std.qsystem.zz_phase`` for a public definition that
    accepts angle parameters.
    """


@hugr_op(quantum_op("Rz", ext=QSYSTEM_EXTENSION))
@no_type_check
def _rz(q: qubit, angle: float) -> None:
    """Rz operation from the qsystem extension.

    See ``guppylang.std.qsystem.rz`` for a public definition that
    accepts angle parameters.
    """
