"""Functions to test and benchmark unitary compilation."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../src/benchmark/bench_compilation.ipynb.

# %% auto #0
__all__ = ['sigma_x', 'sigma_y', 'sigma_z', 'SpecialUnitaries', 'qubit_tensor_product', 'BaseHamiltonian', 'IsingHamiltonian',
           'XXZHamiltonian']

# %% ../../src/benchmark/bench_compilation.ipynb #a593dee3-594e-4760-a02d-db5559f5f25f
from ..imports import *

# %% ../../src/benchmark/bench_compilation.ipynb #e9337291-4d9e-4524-8fe8-9d198a2abb24
class SpecialUnitaries:
    """Special unitary matrices to benchmark compilation."""
    
    @staticmethod
    def QFT(num_qubits: int) -> torch.Tensor:
        """The Quantum Fourier transform (QFT) unitary for `num_qubits`-qubits."""
        
        N  = 2**num_qubits
        wN = np.exp(2.0j*np.pi/N)

        U = torch.zeros((N, N), dtype=torch.complex128)      
        for x in range(N):
            U[:, x] = torch.tensor([np.power(wN, x*k, dtype=complex) for k in range(N)])

        U *= 1.0/np.sqrt(N)   
        return U

# %% ../../src/benchmark/bench_compilation.ipynb #86c32afa-32ad-4a52-9aa6-144d8f76f1f2
sigma_x = torch.tensor([[0, 1],
                        [1, 0]],
                       dtype=torch.complex128)

sigma_y = torch.tensor([[ 0, -1j],
                        [1j,   0]],
                       dtype=torch.complex128)

sigma_z = torch.tensor([[1,  0],
                        [0, -1]],
                       dtype=torch.complex128)

# %% ../../src/benchmark/bench_compilation.ipynb #485317ad-e46c-4652-bd2b-89c0c03afea2
def qubit_tensor_product(num_qubits: int, *ops: torch.Tensor, pos: int | Sequence[int]) -> torch.Tensor:
    """
    Make tensor product with identities, assumes `ops` placed at `pos` in the tensor product ordering.
    """

    _ops = [torch.eye(2) for i in range(num_qubits)]

    if isinstance(pos, int):
        pos = [pos]
    elif isinstance(pos, Sequence):
        assert len(pos) == len(ops)
    else:
        raise NotImplementedError()

    for pos_i, ops_i in zip(pos, ops):
        _ops[pos_i] = ops_i
        
    mat = _ops[0]
    for op in _ops[1:]:
        mat = torch.kron(mat, op)

    return mat

# %% ../../src/benchmark/bench_compilation.ipynb #ecb265cc-521a-44fd-bf8d-c97b46938932
class BaseHamiltonian(abc.ABC):
    """Base implementation of a Hamiltonian."""

    def __init__(self, device: Optional[str | torch.device] = None) -> None:
        self.device = default(device, "cpu")
        self._generate_matrix()
        
        if not torch.allclose(self.data.adjoint(), self.data):
            raise RuntimeError("Generated Hamiltonian matrix is not self-adjoint!")
    
    @abc.abstractmethod
    def _generate_matrix(self) -> torch.Tensor:
        """Generates the Hamiltonian matrix into `self.data`."""
        raise NotImplementedError()

    def get_evolution(self, t: float | torch.Tensor, split_complex_channel: bool = False, dtype: Optional[torch.dtype] = None) -> torch.Tensor:
        """
        Assuming `h_bar=1`. Returns the unitary evolution in marix form.
        """
        U = torch.linalg.matrix_exp(-1j * t * self.data)

        if split_complex_channel:
            U = torch.stack([torch.real(U), torch.imag(U)])

        if exists(dtype):
            U = U.to(dtype)
        
        return U

# %% ../../src/benchmark/bench_compilation.ipynb #8e1da512-6e2f-4eb4-aade-5cb822b6e6f9
class IsingHamiltonian(BaseHamiltonian):
    """Implementation of the Ising Hamiltonian on a qubit chain."""
    
    def __init__(self, 
                 h: float, 
                 J: float, 
                 num_qubits: int, 
                 periodic_boundary: bool = True,
                 device: Optional[str | torch.device] = None) -> None:
        """
        h:     Magnetic field 
        J:     Coupling constant 
        """
        self.h = h
        self.J = J      
        self.num_qubits = num_qubits
        self.periodic_boundary = periodic_boundary
        super().__init__(device)
        
    def _generate_matrix(self) -> torch.Tensor:
        """
        Note:  We take big endian convention in placing the `i,j`-sigmas in tensor product ordering.
               For little endian we need to use `pos = self.num_qubits-i`.
        """
        
        N   = 2**self.num_qubits
        ham = torch.zeros((N, N), dtype=torch.complex128)

        pairs = [(i, i+1) for i in range(self.num_qubits-1)]
        
        if self.periodic_boundary:
            pairs.append((self.num_qubits-1, 0))

        for (i, j) in pairs:
            Z_term = qubit_tensor_product(self.num_qubits, sigma_z, sigma_z, pos=[i, j])

            # Coupling + Perturbation
            ham += -self.J * Z_term

        # Magnetic
        for i in range(self.num_qubits):
            ham += -self.h * qubit_tensor_product(self.num_qubits, sigma_x, pos=i)

        self.data = ham.to(self.device)

# %% ../../src/benchmark/bench_compilation.ipynb #bcff4597-de9d-436e-908a-ff683e868478
class XXZHamiltonian(BaseHamiltonian):
    """Implementation of the XXZ Hamiltonian on a qubit chain."""
    
    def __init__(self, 
                 h: float, 
                 J: float, 
                 delta: float, 
                 num_qubits: int, 
                 periodic_boundary: bool = True,
                 device: Optional[str | torch.device] = None) -> None:
        """
        h:     Magnetic field 
        J:     Coupling constant 
        delta: Perturbation
        """
        self.h = h
        self.J = J      
        self.delta = delta
        self.num_qubits = num_qubits
        self.periodic_boundary = periodic_boundary
        super().__init__(device)
        
    def _generate_matrix(self) -> torch.Tensor:
        """
        Note:  We take big endian convention in placing the `i,j`-sigmas in tensor product ordering.
               For little endian we need to use `pos = self.num_qubits-i`.
        """
        
        N   = 2**self.num_qubits
        ham = torch.zeros((N, N), dtype=torch.complex128)

        pairs = [(i, i+1) for i in range(self.num_qubits-1)]
        
        if self.periodic_boundary:
            pairs.append((self.num_qubits-1, 0))

        for (i, j) in pairs:
            X_term = qubit_tensor_product(self.num_qubits, sigma_x, sigma_x, pos=[i, j])
            Y_term = qubit_tensor_product(self.num_qubits, sigma_y, sigma_y, pos=[i, j])
            Z_term = qubit_tensor_product(self.num_qubits, sigma_z, sigma_z, pos=[i, j])

            # Coupling + Perturbation
            ham += -self.J * (X_term + Y_term + self.delta * Z_term)

        # Magnetic
        for i in range(self.num_qubits):
            ham += -self.h * qubit_tensor_product(self.num_qubits, sigma_x, pos=i)

        self.data = ham.to(self.device)
