import math
import numpy as np
from scipy import signal
from astropy import units as u, constants as ac

DY_POLARIZABILITY = 184.4 # in a.u, most precise measured value of Dy polarizability
DY_MASS = 164*u.u
DY_DIPOLE_MOMENT = 9.93 * ac.muB

#####################################################################
# HELPER FUNCTIONS                                                  #
#####################################################################

def orderOfMagnitude(number):
    return math.floor(math.log(number, 10))

def rotation_matrix(axis, theta):
    """
    Return the rotation matrix associated with counterclockwise rotation about
    the given axis by theta radians.
    In 2-D it is just,
        thetaInRadians = np.radians(theta)
        c, s = np.cos(thetaInRadians), np.sin(thetaInRadians)
        R = np.array(((c, -s), (s, c)))
    In 3-D, one way to do it is use the Euler-Rodrigues Formula as is done here   
    """
    axis = np.asarray(axis)
    axis = axis / math.sqrt(np.dot(axis, axis))
    a = math.cos(theta / 2.0)
    b, c, d = -axis * math.sin(theta / 2.0)
    aa, bb, cc, dd = a * a, b * b, c * c, d * d
    bc, ad, ac, ab, bd, cd = b * c, a * d, a * c, a * b, b * d, c * d
    return np.array([[aa + bb - cc - dd, 2 * (bc + ad), 2 * (bd - ac)],
                     [2 * (bc - ad), aa + cc - bb - dd, 2 * (cd + ab)],
                     [2 * (bd + ac), 2 * (cd - ab), aa + dd - bb - cc]])

def find_nearest(array, value):
    array = np.asarray(array)
    idx = (np.abs(array - value)).argmin()
    return idx

def modulation_function(mod_amp, n_points, func = 'arccos'):
    
    if func == 'sin':
        phi         = np.linspace(0, 2*np.pi, n_points)
        mod_func    = mod_amp * np.sin(phi)
    elif func == 'arccos':
        # phi        = np.linspace(0, 2*np.pi, n_points)
        # mod_func   = mod_amp * (2/np.pi * np.arccos(phi/np.pi-1) - 1)
        phi      = np.linspace(0, 2*np.pi, int(n_points/2))
        tmp_1    = 2/np.pi * np.arccos(phi/np.pi-1) - 1
        tmp_2    = np.flip(tmp_1)
        mod_func = mod_amp * np.concatenate((tmp_1, tmp_2))
    elif func == 'triangle':
        phi         = np.linspace(0, 2*np.pi, n_points)
        mod_func    = mod_amp * signal.sawtooth(phi, width = 0.5) # width of 0.5 gives symmetric rising triangle ramp
    elif func == 'square':
        phi         = np.linspace(0, 1.99*np.pi, n_points)
        mod_func    = mod_amp * signal.square(phi, duty = 0.5)
    else:
        mod_func = None    
    
    if mod_func is not None:
        dx          = (max(mod_func) - min(mod_func))/(2*n_points)

    return dx, mod_func

#####################################################################
# BEAM PARAMETERS                                                   #
#####################################################################

# Rayleigh length
def z_R(w_0, lamb)->np.ndarray:
    return np.pi*w_0**2/lamb

# Beam Radius
def w(pos, w_0, lamb):
    return w_0*np.sqrt(1+(pos / z_R(w_0, lamb))**2)