from __future__ import annotations

import sympy
import torch
from rich.console import Console, RenderableType
from rich.panel import Panel
from rich.tree import Tree
from torch import Tensor, cdouble

from qadence.blocks import (
    ControlBlock,
    ParametricControlBlock,
)
from qadence.blocks.utils import (
    add,  # noqa
    chain,
    kron,
)
from qadence.logger import get_logger
from qadence.parameters import (
    Parameter,
    evaluate,
)
from qadence.types import OpName, TNumber, TParameter

from .parametric import PHASE, RX, RY, RZ
from .primitive import SWAP, I, N, X, Y, Z

logger = get_logger(__name__)


class CNOT(ControlBlock):
    """The CNot, or CX, gate."""

    name = OpName.CNOT

    def __init__(self, control: int, target: int) -> None:
        self.generator = kron(N(control), X(target) - I(target))
        super().__init__((control,), X(target))

    @property
    def eigenvalues_generator(self) -> Tensor:
        return torch.tensor([-2, 0, 0, 0], dtype=cdouble)

    @property
    def eigenvalues(self) -> Tensor:
        return torch.tensor([-1, 1, 1, 1], dtype=cdouble)

    def __ascii__(self, console: Console) -> RenderableType:
        (target, control) = self.qubit_support
        h = abs(target - control) + 1
        return Panel(self._block_title, expand=False, height=3 * h)

    def __rich_tree__(self, tree: Tree = None) -> Tree:
        if tree is None:
            return Tree(self._block_title)
        else:
            tree.add(self._block_title)
        return tree


class MCZ(ControlBlock):
    name = OpName.MCZ

    def __init__(self, control: tuple[int, ...], target: int) -> None:
        self.generator = kron(*[N(qubit) for qubit in control], Z(target) - I(target))
        super().__init__(control, Z(target))

    @property
    def eigenvalues_generator(self) -> Tensor:
        return torch.cat((torch.tensor(-2, dtype=cdouble), torch.zeros(2**self.n_qubits - 1)))

    @property
    def eigenvalues(self) -> Tensor:
        torch.cat((torch.tensor(-1, dtype=cdouble), torch.ones(2**self.n_qubits - 1)))

    def __ascii__(self, console: Console) -> RenderableType:
        (target, control) = self.qubit_support
        h = abs(target - control) + 1
        return Panel(self._block_title, expand=False, height=3 * h)

    def __rich_tree__(self, tree: Tree = None) -> Tree:
        if tree is None:
            return Tree(self._block_title)
        else:
            tree.add(self._block_title)
        return tree


class CZ(MCZ):
    """The CZ gate."""

    name = OpName.CZ

    def __init__(self, control: int, target: int) -> None:
        super().__init__((control,), target)


class MCRX(ParametricControlBlock):
    name = OpName.MCRX

    def __init__(
        self,
        control: tuple[int, ...],
        target: int,
        parameter: Parameter | TNumber | sympy.Expr | str,
    ) -> None:
        self.generator = kron(*[N(qubit) for qubit in control], X(target))
        super().__init__(control, RX(target, parameter))

    @classmethod
    def num_parameters(cls) -> int:
        return 1

    @property
    def eigenvalues_generator(self) -> Tensor:
        return torch.cat(
            (torch.zeros(2**self.n_qubits - 2), torch.tensor([1, -1], dtype=cdouble))
        )

    @property
    def eigenvalues(self) -> Tensor:
        val = evaluate(self.parameters.parameter, as_torch=True)
        lmbd = torch.cos(val / 2.0) - 1j * torch.sin(val / 2.0)
        return torch.cat((torch.ones(2**self.n_qubits - 2), lmbd, lmbd.conj()))


class CRX(MCRX):
    """The CRX gate."""

    name = OpName.CRX

    def __init__(
        self,
        control: int,
        target: int,
        parameter: Parameter | TNumber | sympy.Expr | str,
    ):
        super().__init__((control,), target, parameter)


class MCRY(ParametricControlBlock):
    name = OpName.MCRY

    def __init__(
        self,
        control: tuple[int, ...],
        target: int,
        parameter: Parameter | TNumber | sympy.Expr | str,
    ) -> None:
        self.generator = kron(*[N(qubit) for qubit in control], Y(target))
        super().__init__(control, RY(target, parameter))

    @classmethod
    def num_parameters(cls) -> int:
        return 1

    @property
    def eigenvalues_generator(self) -> Tensor:
        return torch.cat(
            (torch.zeros(2**self.n_qubits - 2), torch.tensor([1, -1], dtype=cdouble))
        )

    @property
    def eigenvalues(self) -> Tensor:
        val = evaluate(self.parameters.parameter, as_torch=True)
        lmbd = torch.cos(val / 2.0) - 1j * torch.sin(val / 2.0)
        return torch.cat((torch.ones(2**self.n_qubits - 2), lmbd, lmbd.conj()))


class CRY(MCRY):
    """The CRY gate."""

    name = OpName.CRY

    def __init__(
        self,
        control: int,
        target: int,
        parameter: TParameter,
    ):
        super().__init__((control,), target, parameter)


class MCRZ(ParametricControlBlock):
    name = OpName.MCRZ

    def __init__(
        self,
        control: tuple[int, ...],
        target: int,
        parameter: Parameter | TNumber | sympy.Expr | str,
    ) -> None:
        self.generator = kron(*[N(qubit) for qubit in control], Z(target))
        super().__init__(control, RZ(target, parameter))

    @classmethod
    def num_parameters(cls) -> int:
        return 1

    @property
    def eigenvalues_generator(self) -> Tensor:
        return torch.cat(
            (torch.zeros(2**self.n_qubits - 2), torch.tensor([1, -1], dtype=cdouble))
        )

    @property
    def eigenvalues(self) -> Tensor:
        val = evaluate(self.parameters.parameter, as_torch=True)
        lmbd = torch.cos(val / 2.0) - 1j * torch.sin(val / 2.0)
        return torch.cat((torch.ones(2**self.n_qubits - 2), lmbd, lmbd.conj()))


class CRZ(MCRZ):
    """The CRZ gate."""

    name = OpName.CRZ

    def __init__(
        self,
        control: int,
        target: int,
        parameter: Parameter | TNumber | sympy.Expr | str,
    ):
        super().__init__((control,), target, parameter)


class MCPHASE(ParametricControlBlock):
    name = OpName.MCPHASE

    def __init__(
        self,
        control: tuple[int, ...],
        target: int,
        parameter: Parameter | TNumber | sympy.Expr | str,
    ) -> None:
        self.generator = kron(*[N(qubit) for qubit in control], Z(target) - I(target))
        super().__init__(control, PHASE(target, parameter))

    @classmethod
    def num_parameters(cls) -> int:
        return 1

    @property
    def eigenvalues_generator(self) -> Tensor:
        return torch.cat(
            (torch.tensor([-2, 0], dtype=cdouble), (torch.zeros(2**self.n_qubits - 2)))
        )

    @property
    def eigenvalues(self) -> Tensor:
        v = evaluate(self.parameters.parameter, as_torch=True)
        return torch.cat((torch.ones(2**self.n_qubits - 1), torch.exp(1j * v)))

    def __rich_tree__(self, tree: Tree = None) -> Tree:
        if tree is None:
            return Tree(self._block_title)
        else:
            tree.add(self._block_title)
        return tree

    def __ascii__(self, console: Console) -> RenderableType:
        (target, control) = self.qubit_support
        h = abs(target - control) + 1
        return Panel(self._block_title, expand=False, height=3 * h)


class CPHASE(MCPHASE):
    """The CPHASE gate."""

    name = OpName.CPHASE

    def __init__(
        self,
        control: int,
        target: int,
        parameter: Parameter | TNumber | sympy.Expr | str,
    ):
        super().__init__((control,), target, parameter)


class CSWAP(ControlBlock):
    """The CSWAP (Control-SWAP) gate."""

    name = OpName.CSWAP

    def __init__(self, control: int | tuple[int, ...], target1: int, target2: int) -> None:
        if isinstance(control, tuple):
            control = control[0]

        a00m = -N(target=control)
        a00p = -N(target=control, state="0")
        a11 = -N(target=target1)
        a22 = -N(target=target2, state="0")
        a12 = 0.5 * (chain(X(target1), Z(target1)) + X(target1))
        a21 = 0.5 * (chain(Z(target2), X(target2)) + X(target2))
        no_effect = kron(a00m, I(target1), I(target2))
        swap_effect = (
            kron(a00p, -1.0 * a22, a11)
            + kron(a00p, -1.0 * a11, a22)
            + kron(a00p, a12, a21)
            + kron(a00p, a21, a12)
        )
        self.generator = no_effect + swap_effect
        super().__init__((control,), SWAP(target1, target2))

    @property
    def eigenvalues_generator(self) -> Tensor:
        return torch.tensor((1, -1, 1, 1, 1, 1, 1, 1), dtype=torch.cdouble)

    @property
    def eigenvalues(self) -> Tensor:
        return torch.tensor((1, -1, 1, 1, 1, 1, 1, 1), dtype=torch.cdouble)

    @property
    def nqubits(self) -> int:
        return 3


class Toffoli(ControlBlock):
    name = OpName.TOFFOLI

    def __init__(self, control: tuple[int, ...], target: int) -> None:
        self.generator = kron(*[N(qubit) for qubit in control], X(target) - I(target))
        super().__init__(control, X(target))

    @property
    def n_qubits(self) -> int:
        return len(self.qubit_support)

    @property
    def eigenvalues_generator(self) -> Tensor:
        return torch.tensor(
            [-2, *[0 for _ in range(2 ** len(self.qubit_support) - 1)]], dtype=cdouble
        )

    @property
    def eigenvalues(self) -> Tensor:
        return torch.tensor(
            [-1, *[1 for _ in range(2 ** len(self.qubit_support) - 1)]], dtype=cdouble
        )
