import numpy as np
import plotly.io as pio
pio.renderers.default='browser'
import os 
os.chdir(os.path.dirname(os.path.abspath(__file__)))
import Plot_Perco as pperco
from scipy import integrate
from math import pi,ceil,floor



def sqrt_g_star_function(k):
    #small useful function in GFF_Simulation
    New = (np.sin(k/2))**2
    return (np.sqrt((2/3)*np.sum(New,axis=1)))**(-1)



def GFF_Simulation(N):
    #We define the GFF using the method described in Sheffield "Gaussian free field for mathematicians", Section 4.4.
    vecs = np.mgrid[0:N, 0:N,0:N].reshape(3,-1).T    
    Lambda_Star = 2*(np.pi/N)*vecs
    Lambda_Star = Lambda_Star[1:] #(remove origin) 
    sqrt_g_star_k = sqrt_g_star_function(Lambda_Star)
    K = np.concatenate(([0], sqrt_g_star_k))

    Gaussians = np.random.normal(0,1, (N**3) )+ (1j)*np.random.normal(0,1, (N**3) )

    sdv_times_Gauss = K * Gaussians
    elements = sdv_times_Gauss.reshape((N,N,N))
    
    return np.real((N**(-1.5))* np.fft.fftn(elements))



def GFF_Cable_Perco(field,unif,N,h):
        #Field is the discrete GFF, unif are the uniform variables used to add Bernoulli random variables on the edges corresponding to the cable system
        #The output is the open edges for the GFF on the cable sytem above level h.
        
        #First compute the discrete level sets of the GFF
        GFF_temp=field>h

        #Then determine which edges are open using the uniform random variables
        GFF_cable_temp=[]
        para=np.exp(-(np.maximum(field-h,0)*np.maximum(np.roll(field,-1,axis=0)-h,0))/3)
        GFF_cable_temp.append(para<unif[0])
    
        para=np.exp(-(np.maximum(field-h,0)*np.maximum(np.roll(field,-1,axis=1)-h,0))/2)
        GFF_cable_temp.append(para<unif[1])
    
        para=np.exp(-(np.maximum(field-h,0)*np.maximum(np.roll(field,-1,axis=2)-h,0))/3)
        GFF_cable_temp.append(para<unif[2])
        

        #Then create an array of size 2N, such that the even vertices are the discrete level sets, and the vertices 2*x+e_i is equal to 1 if and only the edge between x and x+1 is open on the cable system.
        GFF_cable=np.zeros((2*N,2*N,2*N))
        GFF_cable[::2, ::2,::2]=GFF_temp
        GFF_cable[1::2, ::2,::2]=GFF_cable_temp[0]
        GFF_cable[::2, 1::2,::2]=GFF_cable_temp[1]
        GFF_cable[::2, ::2,1::2]=GFF_cable_temp[2]
        return GFF_cable
    

def GFF_Cable_Output(N,hrange):
    #This function generates a variable GFF_Cable_All_Levels, which takes real values for each vertex of the torus, and such that GFF_Cable_All_Levels>=h corresponds to all the points which have not been hit by the cable GFF above level h, for any h in hrange


    #We first sample the discrete GFF.
    GFF_Sample=GFF_Simulation(N)
    
    #We define some uniform random variables which will be useful for the cable system
    unif=np.random.uniform(size=(3,N,N,N))
    
    #We first compute the first level
    GFF_Cable_All_Levels=hrange[0]-1+GFF_Cable_Perco(GFF_Sample, unif, N, hrange[0])*(-hrange[0]+1+max(hrange[-1],1)*3/2)

    #And then iterate
    for k in np.arange(len(hrange)-1):
        #We remove the edges which are closed for the first time at this step, by assigning a value smaller than the current h, but bigger than the last step.
        GFF_Cable_Sample=GFF_Cable_Perco(GFF_Sample, unif, N, hrange[k+1])
        GFF_Cable_All_Levels[(GFF_Cable_All_Levels>=hrange[k]) & (GFF_Cable_Sample==0)]=(hrange[k+1]+hrange[k])/2
    return GFF_Cable_All_Levels
        

def GFF_Cable_Image(N,h,percent,boundary,HTML):
    #This function generates an image of the clusters for the cable system GFF above level h, both as png and html
    #Only the clusters with volume larger than (percent)/100 times the volume of the largest cluster are displayed
    #If percent is negative, we display only the (-percent) largest clusters instead.
    #If boundary=0 then the clusters have periodic boundary condition, and otherwise we use free boundary condition.
    #Put HTML=1 if you want to use the html file, otherwise put html=0 for better performance.
    name='GFF_Cable'    
    pperco.Percolation_Image(N,GFF_Cable_Output(N,[h])>=h,h,percent,boundary,0,name,'h',HTML)



def GFF_Cable_Video(N,hrange,percent,boundary):
    #This function generates a video of the clusters for the cable system GFF, with frames corresponding to levels in hrange
    #Only the clusters with volume larger than (percent)/100 times the volume of the largest cluster are displayed
    #If percent is negative, we display only the (-percent) largest clusters instead.
    #If boundary=0 then the clusters have periodic boundary condition, and otherwise we use free boundary condition.

    name='GFF_Cable'
    print(GFF_Cable_Output(N,hrange))
    pperco.Percolation_Video(N,GFF_Cable_Output(N,hrange),hrange,percent,boundary,0,name,'h')


       
    
def GFF_Cable_Slide(N,hrange,percent,boundary):
    #This function generates a slider of the clusters for the cable system GFF above level h for all h in hrange
    #Only the clusters with volume larger than (percent)/100 times the volume of the largest cluster are displayed
    #If percent is negative, we display only the (-percent) largest clusters instead.
    #If boundary=0 then the clusters have periodic boundary condition, and otherwise we use free boundary condition.

    name='GFF_Cable'
    pperco.Percolation_Slide(N,GFF_Cable_Output(N,hrange),hrange,percent,boundary,0,name,'h')
    
    
def GFF_Cable_Density(h):
    #Computes the probability that an edge is open for the GFF on the cable system at level h
    
    #This is a good approximation for the Green function at 0
    g=1.516386
    
    #First define the function whose integral will give the probability to be open for an edge of the cable system GFF at level h
    f=lambda t, s : (1-np.exp(-(t-h)*(s-h)/3))*np.exp(-(g*(t**2)+g*(s**2)+2*(1-g)*t*s)/(4*g-2))/(2*pi*np.sqrt((2*g-1)))
    return round(integrate.dblquad(f, h, np.inf, h, np.inf)[0],4)
                    

#%%
os.chdir(os.path.dirname(os.path.abspath(__file__)))
    

#Examples of parameters which produce nice outputs. N can be lowered to decrease computation time, and the time to display the html files
hc=0 #The critical parameter
#We display the density of the open edges at the critical parameter
print('Density of open edges at h_c='+str(GFF_Cable_Density(hc)))
size_interval=0.5
hmin=hc-size_interval
hmax=hc+size_interval #Maximal and minimal value we want to display
#We also display the size of the interval we chose in term of density, in order to be able to choose similar ones for different models
print('Interval we plot in terms of density of open edges=['+str(GFF_Cable_Density(hmax))+','+str(GFF_Cable_Density(hmin))+']')


#We first compute a few images as html and png. N needs to be small for the html to not be heavy
N=150
GFF_Cable_Image(N,hc,10,1,1)
hrange=np.arange(floor(hmin*10)/10,ceil(hmax*10)/10+0.1,0.1)
for h in hrange:
    GFF_Cable_Image(N,h,-2,1,1)


#We now make a video of the phase transition. As there is no html we can make N larger, but not too large due to plotly limitations
N=200
number_of_frames=250
h_between_frames=round((hmax-hmin)/number_of_frames,5)
hrange=np.arange(hmin,hmax+h_between_frames/2,h_between_frames)
GFF_Cable_Video(N,hrange,10,1)


#Finally, we plot a few points as a 3D interactive slider. N needs to be quite small for the total html file to not be too heavy
N=100
hrange=np.arange(floor(hmin*10)/10,ceil(hmax*10)/10+0.1,0.1)
GFF_Cable_Slide(N,hrange,10,1)


