import cmath
import pennylane as qml
from pennylane import numpy as np

eps = 1e-13


def random_state_unnormalized(n=4):
    x = np.random.rand(n)+1j*np.random.rand(n)
    return x


def random_state_normalized(n=4):
    x = random_state_unnormalized(n)
    x /= np.linalg.norm(x)
    return x


def get_vector(alpha, beta):
    """
    Function to compute 3D Cartesian coordinates
    from 2D qubit vector.
    """

    # get phases
    angle_alpha = cmath.phase(alpha)
    angle_beta = cmath.phase(beta)

    # avoiding wrong normalization due to rounding errors
    if cmath.isclose(angle_alpha, cmath.pi):
        angle_alpha = 0
    if cmath.isclose(angle_beta, cmath.pi):
        angle_beta = 0

    if (angle_beta < 0 and angle_alpha < angle_beta) or (angle_beta > 0 and angle_alpha > angle_beta):
        denominator = cmath.exp(1j*angle_beta)
    else:
        denominator = cmath.exp(1j*angle_alpha)

    # eliminate global phase
    alpha_new = alpha/denominator
    beta_new = beta/denominator

    # special case to avoid division by zero
    if abs(alpha) == 0 or abs(beta) == 0:
        if alpha == 0:
            return [0, 0, -1], [np.pi, 0]
        else:
            return [0, 0, 1], [0, 0]
    else:
        # compute theta and phi from alpha and beta
        theta = 2*cmath.acos(alpha_new)
        phi = -1j*cmath.log(beta_new/cmath.sin(theta/2))

        # compute the Cartesian coordinates
        x = cmath.sin(theta)*cmath.cos(phi)
        y = cmath.sin(theta)*cmath.sin(phi)
        z = cmath.cos(theta)

    return [x.real, y.real, z.real], [theta, phi]


def sx(i, n):
    assert i < n, "i has to be smaller than n"
    return np.kron(np.kron(np.eye(2**(i)), np.array([[0, 1], [1, 0]])), np.eye(2**(n-i-1)))


def sy(i, n):
    assert i < n, "i has to be smaller than n"
    return np.kron(np.kron(np.eye(2**(i)), np.array([[0, -1j], [1j, 0]])), np.eye(2**(n-i-1)))


def sz(i, n):
    assert i < n, "i has to be smaller than n"
    return np.kron(np.kron(np.eye(2**(i)), np.array([[1, 0], [0, -1]])), np.eye(2**(n-i-1)))
