import numpy as np
import numpy.random as random
from math import pi, sqrt


def random_orientation(samples=150, planar=True):
    a = random.rand(samples, 2)

    alpha = 2*pi*a[:, 0]
    if planar:
        beta = pi*a[:, 1]
        wt = np.sin(beta)
    else:
        beta = np.arccos(1-2*a[:, 1])
        wt = np.ones_like(beta)

    return alpha, beta, wt


def random_planar(samples=150):
    return random_orientation(samples=samples, planar=True)


def random_spherical(samples=150):
    return random_orientation(samples=samples, planar=False)


def grid(a=144, b=None, planar=True):
    if b:
        a_axis = range(a)
        b_axis = range(1, b+1)
    else:
        a = b = int(sqrt(a))
        a_axis = range(a)
        b_axis = range(1, b+1)
    alpha = 2*pi*np.array(a_axis*b)/a
    beta = []
    for i in b_axis:
        beta += [i]*a
    if planar:
        beta = pi*np.array(beta)/(b+1)
        wt = np.sin(beta)
    else:
        beta = np.arccos(1 - 2*np.array(beta)/(b+1))
        wt = np.ones_like(beta)

    return alpha, beta, wt


def grid_planar(a=12, b=None):
    return grid(a=a, b=b, planar=True)


def grid_spherical(a=12, b=None):
    return grid(a=a, b=b, planar=False)


def zcw(n=150, planar=True):
    fib0, fib1, fib2 = [0, 1, 1]
    while True:
        fib0, fib1, fib2 = fib1, fib0 + fib1, fib1+fib2
        if fib1+fib2 >= n:
            break

    arr = np.arange(fib2)
    alpha = 2*pi*np.mod(arr/fib2, 1.0)
    if planar:
        beta = pi*np.mod(arr*fib0/fib2, 1.0)
        wt = np.sin(beta)
    else:
        beta = np.arccos(1 - 2*np.mod(arr*fib0/fib2, 1.0))
        wt = np.ones_like(beta)

    return alpha, beta, wt


def zcw_planar(n=150):
    return zcw(n=n, planar=True)


def zcw_spherical(n=150):
    return zcw(n=n, planar=False)


def alderman(samples=144):
    n = int(0.5*sqrt(samples))
    alpha = []
    beta = []
    wt = []
    for i in range(-n, n):
        i_abs = abs(i)
        for j in range(-n+i_abs, n-i_abs+1):
            j_abs = abs(j)
            k = n-i_abs-j_abs
            inv_r_ij = 1/sqrt(i**2 + j**2 + k**2)
            a_ij = (np.arctan2(j, i)+2*pi) % (2*pi)
            alpha.append(a_ij)
            alpha.append(a_ij)
            b_ij = np.arccos(k*inv_r_ij)
            b_ij_neg = np.arccos(-k*inv_r_ij)
            beta.append(b_ij)
            beta.append(b_ij_neg)
            wt.append((n*inv_r_ij)**3)
            wt.append((n*inv_r_ij)**3)

    return np.array(alpha), np.array(beta), np.array(wt)


def sophe(samples=146):
    n = int(0.5*sqrt(2*samples))
    alpha = []
    beta = []
    for quad in [0, 1, 2, 3]:
        for i in range(1, n):
            for j in range(1, i+1):
                # 0,90 / 0,90
                alpha.append(0.5*pi*(i-j+1)/i + 0.5*quad*pi)
                beta.append(0.5*pi*i/(n-1))

                # 0,90 / 90,180
                if i != n-1:
                    alpha.append(0.5*pi*(i-j+1)/i + 0.5*quad*pi)
                    beta.append(0.5*pi*(n-i-1)/(n-1)+0.5*pi)

    return np.array(alpha), np.array(beta), np.ones_like(alpha)


def spherical_to_xyz(phi, theta, r=1.):
    x = r*np.cos(phi)*np.sin(theta)
    y = r*np.sin(phi)*np.sin(theta)
    z = r*np.cos(theta)

    return x, y, z


def xyz_to_spherical(x, y, z):
    r = np.sqrt(x**2 + y**2 + z**2)
    th = np.arccos(z/r)
    ph = np.arctan2(y, x)

    return r, th, ph
