import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import gamma,skew,pearsonr

#%%

####
####          SYSTEM SETUP. DOES NOT CHANGE EVENT-BY-EVENT. Choose colliding nuclei,
####          write nuclear parameters. Then create grid. Choose grid size and step.
####


# 2-parameter Fermi with deformation parameters beta2, gamma, beta3, beta4 (not used).
# Radius expansion in spherical harmonics
def rho_ws(x, y, z, R, a, b2, g, b3, b4):
    r = np.sqrt(x**2+y**2+z**2) + 1e-15
    B2 = np.sqrt(5/16/np.pi)*b2*((2*z**2-x**2-y**2)*np.cos(g) + np.sqrt(3)*(x**2-y**2)*np.sin(g))/r**2
    B3 = 1/4*np.sqrt(7/np.pi)*b3*(5*z**3-3*z*r**2)/r**3
    B4 = 3/16*np.sqrt(1/np.pi)*b4*(35*z**4-30*z**2*r**2+3*r**4)/r**4
    rho_local = 1/(1+np.exp((r-R*(1+B2+B3+B4))/a))
    return rho_local

##
##Nuclear parameters.
##

#16O
R_A, a_A, A_A, beta2_A, gamma_A, beta3_A, beta4_A = 2.61, 0.51, 16, 0, 0, 0, 0 #fm
R_B, a_B, A_B, beta2_B, gamma_B, beta3_B, beta4_B = 2.61, 0.51, 16, 0, 0, 0, 0 #fm

#76Ge
#R_A, a_A, A_A, beta2_A, gamma_A, beta3_A, beta4_A = 4.60, 0.517, 76, 0.26, 0.47, 0, 0.006 # fm
#R_B, a_B, A_B, beta2_B, gamma_B, beta3_B, beta4_B = 4.60, 0.517, 76, 0.26, 0.47, 0, 0.006 # fm

#129Xe
#R_A, a_A, A_A, beta2_A, gamma_A, beta3_A, beta4_A = 5.60, 0.492, 129, 0.2, 0.0, 0.0, 0.0 # fm
#R_B, a_B, A_B, beta2_B, gamma_B, beta3_B, beta4_B = 5.60, 0.49, 129, 0.2, 0.0, 0.0, 0.0 # fm

##
##Construction of grid for the event.
##

#grid
#Following TRENTo parameters, size -10 to 10 fm, with step 0.2 fm is standard.
size=50
dim=10 #fm
grid_step=2*dim/(size) #fm
area=grid_step**2
xx=np.linspace(-dim,dim,size+1)
xx=(xx[1:]+xx[:-1])/2

#grid coordinates
x_grid,y_grid=np.meshgrid(xx,xx,indexing='ij')

#choose number of events to generate
nev=int(4e3)

#matrix where we store observables
obs=np.zeros((nev,13)) #with deformation

# matrices for final 1- ,2-point functions
ave_S1 = np.zeros((size, size))
ave_S2 = np.zeros((size, size, size, size))

# arrays of nucleon positions
coords = np.zeros((nev, A_A, 3))

ev=0
for ev in range(nev):

####
####         GENERATE COORDINATES according to a matter density that fluctuates event-by-event.
####         The fluctuating quantities are the deformation parameters, beta and gamma, and the
####         orientation of the two nuclei in space. For each nucleus one has to generate angles
####         and deformation parameters.
####


    # Now generate nucleon coordinates according to WS distribution.
    # We perform independent sampling from the deformed density.
    # No short-range correlations are implemented.

    # Nucleus A
    lim = R_A+10*a_A  # fm  ## 10*a prescription from TRENTO
    x_A = np.zeros(A_A)
    y_A = np.zeros(A_A)
    z_A = np.zeros(A_A)

    # Optimized algorithm
    gen = int(10*A_A)  # factor 10 works well
    cont = 0
    while cont >= 0:
        # print cont,
        loop = cont
        cx = np.random.uniform(-lim, lim, size=gen)
        cy = np.random.uniform(-lim, lim, size=gen)
        cz = np.random.uniform(-lim, lim, size=gen)
        rand = np.random.random(gen,)
        n_eval = rho_ws(cx, cy, cz, R_A, a_A,
                        beta2_A, gamma_A, beta3_A, beta4_A)
        n_eval /= rand
        coords_ev = np.where(n_eval > 1)
        cont = cont+coords_ev[0].size
        x_ev = cx[coords_ev[0]]
        y_ev = cy[coords_ev[0]]
        z_ev = cz[coords_ev[0]]
        if cont <= A_A:
            x_A[loop:cont] = x_ev
            y_A[loop:cont] = y_ev
            z_A[loop:cont] = z_ev
        else:
            mmax = A_A-loop
            x_A[loop:] = x_ev[0:mmax]
            y_A[loop:] = y_ev[0:mmax]
            z_A[loop:] = z_ev[0:mmax]
            break

    # Now we randomly rotate the coordinates of the sampled nucleons in 3D.
    # 3-1-3 convention
    # first rotation in (x,y), u is uniformly distributed
    # second rotation in (y,z), cos(v) is uniformly distributed
    # third rotation in (x,y), w is uniformly distributed
    u = np.random.uniform(0, 2*np.pi)
    dummy = np.random.uniform(-1, 1)
    v = np.arccos(dummy)
    w = np.random.uniform(0, 2*np.pi)

    # rotation in (x,y) plane by angle u
    x_u = x_A*np.cos(u)-y_A*np.sin(u)
    y_u = x_A*np.sin(u)+y_A*np.cos(u)
    z_u = z_A

    # rotation in (y,z) plane by angle v
    x_v = x_u
    y_v = y_u*np.cos(v)-z_u*np.sin(v)
    z_v = y_u*np.sin(v)+z_u*np.cos(v)

    # rotation in (x,y) plane by angle w
    x_w = x_v*np.cos(w)-y_v*np.sin(w)
    y_w = x_v*np.sin(w)+y_v*np.cos(w)
    z_w = z_v

    # final coords for nucleus A
    x_A = x_w
    y_A = y_w
    z_A = z_w

    # Nucleus B
    lim = R_B+10*a_B  # fm  ## 10*a prescription from TRENTO
    x_B = np.zeros(A_B)
    y_B = np.zeros(A_B)
    z_B = np.zeros(A_B)

    # Optimized algorithm
    gen = int(10*A_B)  # factor 10 works well
    cont = 0
    while cont >= 0:
        # print cont,
        loop = cont
        cx = np.random.uniform(-lim, lim, size=gen)
        cy = np.random.uniform(-lim, lim, size=gen)
        cz = np.random.uniform(-lim, lim, size=gen)
        rand = np.random.random(gen,)
        n_eval = rho_ws(cx, cy, cz, R_B, a_B,
                        beta2_B, gamma_B, beta3_B, beta4_B)
        n_eval /= rand
        coords_ev = np.where(n_eval > 1)
        cont = cont+coords_ev[0].size
        x_ev = cx[coords_ev[0]]
        y_ev = cy[coords_ev[0]]
        z_ev = cz[coords_ev[0]]
        if cont <= A_B:
            x_B[loop:cont] = x_ev
            y_B[loop:cont] = y_ev
            z_B[loop:cont] = z_ev
        else:
            mmax = A_B-loop
            x_B[loop:] = x_ev[0:mmax]
            y_B[loop:] = y_ev[0:mmax]
            z_B[loop:] = z_ev[0:mmax]
            break

    # Now we randomly rotate the coordinates of the sampled nucleons in 3D.
    # 3-1-3 convention
    # first rotation in (x,y), u is uniformly distributed
    # second rotation in (y,z), cos(v) is uniformly distributed
    # third rotation in (x,y), w is uniformly distributed
    u = np.random.uniform(0, 2*np.pi)
    dummy = np.random.uniform(-1, 1)
    v = np.arccos(dummy)
    w = np.random.uniform(0, 2*np.pi)

    # rotation in (x,y) plane by angle u
    x_u = x_B*np.cos(u)-y_B*np.sin(u)
    y_u = x_B*np.sin(u)+y_B*np.cos(u)
    z_u = z_B

    # rotation in (y,z) plane by angle v
    x_v = x_u
    y_v = y_u*np.cos(v)-z_u*np.sin(v)
    z_v = y_u*np.sin(v)+z_u*np.cos(v)

    # rotation in (x,y) plane by angle w
    x_w = x_v*np.cos(w)-y_v*np.sin(w)
    y_w = x_v*np.sin(w)+y_v*np.cos(w)
    z_w = z_v

    # final coords for nucleus B
    x_B = x_w
    y_B = y_w
    z_B = z_w
    
    # store coords for nucleus A
    coords[ev, :, 0] = x_A
    coords[ev, :, 1] = y_A
    coords[ev, :, 2] = z_A

###
### GENERATE AN IMPACT PARAMETER.
###
    b=0
    bmax=24 #fm
    #Generate b until at least one collision takes place.
    check=0
    while check==0:
        #b=np.random.random()**0.5*bmax #fm
        b=0
        #Shift coordinates along X to include impact parameter.
        #Pay attention to the choice of + and - for A or B.
        x_Ab=x_A-b/2
        x_Bb=x_B+b/2

####
####         DEFINE PARTICIPANT NUCLEONS. Follow PHOBOS Glauber (black disks).
####         Two nucleons interact if they are separated by a distance smaller
####         that the 'ball diameter' D=sqrt(sigma/pi), where sigma is the
####         nucleon-nucleon cross section.
####

        sigma = 7    #fm^2 for 5.02 TeV
        #sigma = 6.4  #fm^2 for 2.76 TeV
        #sigma = 4.2  #fm^2 for 200 GeV
        #sigma = 3.32 #fm^2 for 27 GeV
        #sigma = 3.12 #fm^2 for 7.7 GeV
        D=np.sqrt(sigma/np.pi) #fm

        ## Optimized algorithm to select participant nucleons.

        #np.repeat A and np.tile B. Symmetric is also possible.
        array_xA=np.repeat(x_Ab,A_B)
        array_xB=np.tile(x_Bb,A_B)
        array_yA=np.repeat(y_A,A_B)
        array_yB=np.tile(y_B,A_B)

        #Compute distances between A_A*A_B nucleons.
        #Save index where distance-D is negative.
        dist=np.sqrt((array_xA-array_xB)**2+(array_yA-array_yB)**2)-D
        index=np.where(dist<0)[0]

        #Extract index for participants A and B.
        ind_A=np.unique(np.floor(index/A_B)).astype('int')
        ind_B=np.unique(index%A_B)

        #Extract participant coordinates.
        xp_A=x_Ab[ind_A]
        yp_A=y_A[ind_A]
        xp_B=x_Bb[ind_B]
        yp_B=y_B[ind_B]

#        coords_px=np.concatenate((xp_A,xp_B),axis=0)
#        coords_py=np.concatenate((yp_A,yp_B),axis=0)

        #Extract spectator coordinates.
#        ind_sA=np.delete(np.arange(A_A),ind_A)
#        ind_sB=np.delete(np.arange(A_B),ind_B)

#        xs_A=x_Ab[ind_sA]
#        ys_A=y_A[ind_sA]
#        xs_B=x_Bb[ind_sB]
#        ys_B=y_B[ind_sB]

#        coords_sx=np.concatenate((xs_A,xs_B),axis=0)
#        coords_sy=np.concatenate((ys_A,ys_B),axis=0)

        #number of participants
        np_A=xp_A.size
        np_B=xp_B.size

#        #number of spectators
#        ns_A=xs_A.size
#        bs_B=xs_B.size

        #exit loop if there is at least one participant
        if np_A!=0:
            check=check+1

    #store participant nucleons and b
    obs[ev,0]=b
    obs[ev,1]=np_A+np_B
    
    ########
    # do not include participants
    xp_A=x_A
    xp_B=x_B
    yp_A=y_A
    yp_B=y_B
    np_A=A_A
    np_B=A_A

####
####          GENERATE PARTICIPANT T_A and T_B, AND ENTROPY PROFILE.
####          Each participant is smeared with a Gaussian. The sum of
####          all Gaussians for nucleus A or B gives the corresponding
####          nuclear thickness. The total entropy profile of the collision
####          is then given by a function of TA and TB, following TRENTo.
####

##We need to evaluate the thickness functions over some grid.
##Following TRENTo parameters, size -10 to 10 fm, with step 0.2 fm is standard.

    #normalization for each Gaussian follows gamma distribution of params k, 1/k
    #k=2.0
    w_A=np.ones(np_A) #gamma.rvs(k,loc=0,scale=1/k,size=np_A) #
    w_B=np.ones(np_B) #gamma.rvs(k,loc=0,scale=1/k,size=np_B) #

    #"thickness" functions of nuclei A and B, as a sum of Gaussians.
    T_A=np.zeros((size,size))
    T_B=np.zeros((size,size))
    
    #sigma of a nucleon
    w=0.5 #fm

    #Optimized algorithm to sample Gaussians on the grid.
    #Loop over participants in A. Evaluate Gaussian up to 5*sigma.
    wmax=5 ### as implemented in Trento
    for j in range(np_A):
        x_loop=(x_grid-xp_A[j])**2
        y_loop=(y_grid-yp_A[j])**2
        radius=np.sqrt(x_loop+y_loop)
        indic=np.where(radius<w*wmax)
        indic_x=indic[0]
        indic_y=indic[1]
        T_A[indic_x,indic_y]+=w_A[j]*1/(2*np.pi*w**2)*np.exp(-(x_loop[indic_x,indic_y]+y_loop[indic_x,indic_y])/(2*w**2)) #fm^-2

    #Loop over participants in B. Evaluate Gaussian up to 5*sigma.
    for j in range(np_B):
        x_loop=(x_grid-xp_B[j])**2
        y_loop=(y_grid-yp_B[j])**2
        radius=np.sqrt(x_loop+y_loop)
        indic=np.where(radius<w*wmax)
        indic_x=indic[0]
        indic_y=indic[1]
        T_B[indic_x,indic_y]+=w_B[j]*1/(2*np.pi*w**2)*np.exp(-(x_loop[indic_x,indic_y]+y_loop[indic_x,indic_y])/(2*w**2)) #fm^-2

    #Compute entropy per rapidity profile. We follow TRENTo.
    #Define a avariable p, then use generalized mean.
    p=1
    if p==0:
        s = np.sqrt(T_A*T_B)  #fm^-2

    if p==1:
        s = (T_A + T_B ) / 2  #fm^-2

####
#### STEP 5 : CALCULATE "OBSERVABLES". With the knowledge of s, we can move on
####          to the calculation of eccentricities, rms size, total entropy.
####          Just use algebra of 2D grids.
####

    #Total entropy.
    s_tot=np.sum(s)
    obs[ev,2]=s_tot*area
    
    #Total energy
    ene = s**(4/3)
    obs[ev,3]=np.sum(ene)*area

    #Recentering correction
    center_x=np.sum(s*x_grid)/s_tot
    center_y=np.sum(s*y_grid)/s_tot
    x_cen=x_grid-center_x
    y_cen=y_grid-center_y

    #RMS radius, eccentricity, triangularity
    z=x_cen+y_cen*1.j
    r2=x_cen**2+y_cen**2
    r3=r2**(3/2)
    e2_num=s*z**2
    e2_den=s*r2
    rms_rad=np.sqrt(np.sum(e2_den)/s_tot) #no area factor
    obs[ev,4]=rms_rad
    e3_num=s*z**3
    e3_den=s*r3
    E2=-np.sum(e2_num)/np.sum(e2_den)
    obs[ev,6]=np.abs(E2)
    E3=-np.sum(e3_num)/np.sum(e3_den)
    obs[ev,7]=np.abs(E3)
    E1=-np.sum(r2*s*z)/np.sum(e3_den)
    obs[ev,5]=np.abs(E1)

    # calculation of n-point functions
    # store <s(x)>
    ave_S1 += s

    # store <s(x)s(y)>
    ave_S2_ev = np.multiply.outer(s, s) 
    ave_S2 += ave_S2_ev 

    if ev%400==0:
        print(ev, end=' ')

#%%

#################################
## compute perturbative quantities
#################################


# calculate the 1-point function
C1x = ave_S1/nev  # average density

# calculate the connected 2-point function
C1xC1y = np.multiply.outer(C1x, C1x)  # product of averages
C2xy = ave_S2/nev - C1xC1y

# These arrays of indices permit one to generate all possible combinations of 4 coordinates used in the computations involving the connected 2-point function, C2(x,y), to be seen below.
# 4 arrays of indices (all possible size^4 combination)
nu = np.arange(size)

i1_2 = np.tile(nu, size**3)  # size^4 different
i2_2 = np.repeat(i1_2[0:size**3], size**1)  # size^3 different
i3_2 = np.repeat(i1_2[0:size**2], size**2)  # size^2 different
i4_2 = np.repeat(i1_2[0:size**1], size**3)  # size different


## eccentricity
E2E2_nume = np.real(np.sum((xx[i1_2]+xx[i2_2]*1.j)**2*(xx[i3_2]-xx[i4_2]*1.j)**2*C2xy[i1_2, i2_2, i3_2, i4_2])*area**2)
E2E2_deno = np.sum((x_grid**2+y_grid**2)*C1x*area)**2
E2E2 = E2E2_nume / E2E2_deno

## triangularity
E3E3_nume = np.real(np.sum((xx[i1_2]+xx[i2_2]*1.j)**3*(xx[i3_2]-xx[i4_2]*1.j)**3*C2xy[i1_2, i2_2, i3_2, i4_2])*area**2)
E3E3_deno = np.sum((x_grid**2+y_grid**2)**(3/2)*C1x*area)**2
E3E3 = E3E3_nume / E3E3_deno

print("")
print("exact e2{2} for beta2 = ",beta2_A, "is ", np.mean(obs[:,6]**2)**0.5)
print("exact e3{2} for beta2 = ",beta2_A, "is ", np.mean(obs[:,7]**2)**0.5)
print("perturbative e2{2} = ", E2E2**0.5)
print("perturbative e3{2} = ", E3E3**0.5)
print("")

#%%

#######################

r2=np.zeros(nev)
r3=np.zeros(nev)
r4=np.zeros(nev)
r6=np.zeros(nev)
r2perp=np.zeros(nev)
r3perp=np.zeros(nev)
r4perp=np.zeros(nev)
r6perp=np.zeros(nev)

e2=np.zeros(nev)
e3=np.zeros(nev)

print("")
for ev in range(nev):

    x=coords[ev][:,0]
    y=coords[ev][:,1]
    z=coords[ev][:,2]

    r2[ev]=np.sum((x**2+y**2+z**2)**(1))
    r3[ev]=np.sum((x**2+y**2+z**2)**(3/2))
    r4[ev]=np.sum((x**2+y**2+z**2)**(2))
    r6[ev]=np.sum((x**2+y**2+z**2)**(3))

    r2perp[ev]=np.sum((x**2+y**2)**(1))
    r3perp[ev]=np.sum((x**2+y**2)**(3/2))
    r4perp[ev]=np.sum((x**2+y**2)**(2))
    r6perp[ev]=np.sum((x**2+y**2)**(3))

    e2[ev]=np.abs(np.sum((x+y*1.j)**(2))) # numerator only
    e3[ev]=np.abs(np.sum((x+y*1.j)**(3))) # numerator only

    if ev % 50000 == 0:
        print(ev, end=" "),

A=A_A

# exact results directly from point-like nucleons
r2_psi = np.mean(r2)/A
r4_psi = np.mean(r4)/A
r6_psi = np.mean(r6)/A

r2perp_psi = np.mean(r2perp)/A
r3perp_psi = np.mean(r3perp)/A
r4perp_psi = np.mean(r4perp)/A
r6perp_psi = np.mean(r6perp)/A

E2_psi = 1/(A**2-A)*(np.mean(e2**2)-np.mean(r4perp))
E3_psi = 1/(A**2-A)*(np.mean(e3**2)-np.mean(r6perp))

print("")
print('<r^2>_psi=', r2_psi)
print('<r^4>_psi=', r4_psi)
print('<r^6>_psi=', r6_psi)
print('<r_perp^2>_psi=', r2perp_psi)
print('<r_perp^3>_psi=', r3perp_psi)
print('<r_perp^4>_psi=', r4perp_psi)
print('<r_perp^6>_psi=', r6perp_psi)
print('<E_2>_psi=', E2_psi)
print('<E_3>_psi=', E3_psi)
print('')
print("exact e2{2} for beta2 = ",beta2_A, "is ", np.mean(obs[:,6]**2)**0.5)
print("exact e3{2} for beta2 = ",beta2_A, "is ", np.mean(obs[:,7]**2)**0.5)
print("perturbative e2{2} = ", E2E2**0.5)
print("perturbative e3{2} = ", E3E3**0.5)
print('Point-like <e_2^2>^1/2 [spherical] =', (1/2 * (1/A)*r4perp_psi/r2perp_psi**2)**0.5 )
print('Point-like <e_3^2>^1/2 [spherical] =', (1/2 * (1/A)*r6perp_psi/r3perp_psi**2)**0.5 )
print('Point-like <e_2^2>^1/2 [deformed] =', (1/2 * ((1/A)*r4perp_psi/r2perp_psi**2 + (A-1)/A*E2_psi/r2perp_psi**2 ))**0.5)
print('Point-like <e_3^2>^1/2 [deformed] =', (1/2 * ((1/A)*r6perp_psi/r3perp_psi**2 + (A-1)/A*E3_psi/r3perp_psi**2 ))**0.5)
print('')

#%%