from __future__ import annotations

from typing import Union

import numpy as np
import torch
from rich.console import Console
from rich.padding import Padding
from torch import Tensor, cdouble, tensor

from qadence.blocks import (
    AbstractBlock,
    PrimitiveBlock,
)
from qadence.blocks.primitive import ProjectorBlock
from qadence.blocks.utils import (
    chain,
    kron,
)
from qadence.logger import get_logger
from qadence.parameters import (
    Parameter,
)
from qadence.types import OpName, TNumber

logger = get_logger(__name__)


class X(PrimitiveBlock):
    """The X gate."""

    name = OpName.X

    def __init__(self, target: int):
        super().__init__((target,))

    @property
    def generator(self) -> AbstractBlock:
        return self

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

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


class Y(PrimitiveBlock):
    """The Y gate."""

    name = OpName.Y

    def __init__(self, target: int):
        super().__init__((target,))

    @property
    def generator(self) -> AbstractBlock:
        return self

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

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


class Z(PrimitiveBlock):
    """The Z gate."""

    name = OpName.Z

    def __init__(self, target: int):
        super().__init__((target,))

    @property
    def generator(self) -> AbstractBlock:
        return self

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

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


class I(PrimitiveBlock):
    """The identity gate."""

    name = OpName.I

    def __init__(self, target: int):
        super().__init__((target,))

    def __ixor__(self, other: AbstractBlock | int) -> AbstractBlock:
        if not isinstance(other, AbstractBlock):
            raise ValueError(
                f"Can only initialize a kron block with another block. Got {type(other)}."
            )
        return other

    def __imul__(self, other: AbstractBlock | TNumber | Parameter) -> AbstractBlock:
        if not isinstance(other, AbstractBlock):
            raise ValueError(
                "In-place multiplication is available " "only for AbstractBlock instances"
            )
        return other

    @property
    def generator(self) -> AbstractBlock:
        return I(*self.qubit_support)

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

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

    def __ascii__(self, console: Console) -> Padding:
        return Padding("──────", (1, 1, 1, 1))


class Projector(ProjectorBlock):
    """The projector operator."""

    name = OpName.PROJ

    def __init__(
        self,
        ket: str,
        bra: str,
        qubit_support: int | tuple[int, ...],
    ):
        super().__init__(ket=ket, bra=bra, qubit_support=qubit_support)

    @property
    def generator(self) -> None:
        raise ValueError("Property `generator` not available for non-unitary operator.")

    @property
    def eigenvalues_generator(self) -> None:
        raise ValueError("Property `eigenvalues_generator` not available for non-unitary operator.")


class N(Projector):
    """The N = (1/2)(I-Z) operator."""

    name = OpName.N

    def __init__(self, target: int, state: str = "1"):
        super().__init__(ket=state, bra=state, qubit_support=(target,))

    @property
    def generator(self) -> None:
        raise ValueError("Property `generator` not available for non-unitary operator.")

    @property
    def eigenvalues_generator(self) -> None:
        raise ValueError("Property `eigenvalues_generator` not available for non-unitary operator.")

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


class S(PrimitiveBlock):
    """The S / Phase gate."""

    name = OpName.S

    def __init__(self, target: int):
        self.generator = I(target) - Z(target)
        super().__init__((target,))

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

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

    def dagger(self) -> SDagger:
        return SDagger(*self.qubit_support)


class SDagger(PrimitiveBlock):
    """The Hermitian adjoint/conjugate transpose of the S / Phase gate."""

    name = OpName.SDAGGER

    def __init__(self, target: int):
        self.generator = I(target) - Z(target)
        super().__init__((target,))

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

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

    def dagger(self) -> S:
        return S(*self.qubit_support)


class H(PrimitiveBlock):
    """The Hadamard or H gate."""

    name = OpName.H

    def __init__(self, target: int):
        self.generator = (1 / np.sqrt(2)) * (X(target) + Z(target) - np.sqrt(2) * I(target))
        super().__init__((target,))

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

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


class Zero(PrimitiveBlock):
    name = OpName.ZERO

    def __init__(self) -> None:
        self.generator = 0 * I(0)
        super().__init__((0,))

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

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

    def __add__(self, other: AbstractBlock) -> AbstractBlock:
        return other

    def __iadd__(self, other: AbstractBlock) -> AbstractBlock:
        return other

    def __sub__(self, other: AbstractBlock) -> AbstractBlock:
        return -other

    def __isub__(self, other: AbstractBlock) -> AbstractBlock:
        return -other

    def __mul__(self, other: AbstractBlock | TNumber | Parameter) -> AbstractBlock:
        return self

    def __imul__(self, other: AbstractBlock | TNumber | Parameter) -> AbstractBlock:
        return self

    def __ixor__(self, other: AbstractBlock | TNumber | Parameter) -> AbstractBlock:
        return self

    def __pow__(self, power: int) -> AbstractBlock:
        return self


class T(PrimitiveBlock):
    """The T gate."""

    name = OpName.T

    def __init__(self, target: int):
        self.generator = I(target) - Z(target)
        super().__init__((target,))

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

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

    @property
    def n_qubits(self) -> int:
        return 1

    def dagger(self) -> TDagger:
        return TDagger(*self.qubit_support)


class TDagger(PrimitiveBlock):
    """The Hermitian adjoint/conjugate transpose of the T gate."""

    # FIXME: this gate is not support by any backend
    name = "T_dagger"

    def __init__(self, target: int):
        self.generator = I(target) - Z(target)
        super().__init__((target,))

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

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

    @property
    def n_qubits(self) -> int:
        return 1

    def dagger(self) -> T:
        return T(*self.qubit_support)


class SWAP(PrimitiveBlock):
    """The SWAP gate."""

    name = OpName.SWAP

    def __init__(self, control: int, target: int) -> None:
        a11 = 0.5 * (Z(control) - I(control))
        a22 = -0.5 * (Z(target) + I(target))
        a12 = 0.5 * (chain(X(control), Z(control)) + X(control))
        a21 = 0.5 * (chain(Z(target), X(target)) + X(target))
        self.generator = (
            kron(-1.0 * a22, a11) + kron(-1.0 * a11, a22) + kron(a12, a21) + kron(a21, a12)
        )
        super().__init__((control, 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)

    @property
    def n_qubits(self) -> int:
        return 2

    @property
    def _block_title(self) -> str:
        c, t = self.qubit_support
        s = f"{self.name}({c}, {t})"
        return s if self.tag is None else (s + rf" \[tag: {self.tag}]")


TPauliBlock = Union[X, Y, Z, I, N]
