"""
Generate a 2D Gaussian random field from white noise filtered in Fourier space.
The method is described in the following paper:
    Hu, Y.Z. and Tonder, K., 1992. Simulation of 3-D random rough surface by 2-D digital filter and Fourier analysis. International journal of machine tools and manufacture, 32(1-2), pp.83-90.
    DOI: 10.1016/0890-6955(92)90064-N
This part of the code comes from the following GitHub repository (C++/Python): https://github.com/vyastreb/rfgen

The remaining part of the code is about analysis of the generated random field.

Author: Vladislav Yastrebov, MINES Paris - PSL, Centre des matériaux, CNRS UMR 7633, BP 87, 91003 Evry, France (with help of GPT4 and Copilot)
Date: Jan 2024 - Feb 2025
License: BSD 3 clause
"""

import numpy as np
import matplotlib.pyplot as plt
from numpy.fft import fft2, ifft2, fftfreq, fftshift, ifftshift, fftn, ifftn
from scipy.optimize import curve_fit
import time
import rfgen

"""
  Analysis of the generated random field
"""

N0 = 1024           # Size of the random field
k_low_1 = 1 / N0    # Lower cutoff of the power spectrum
k_low_2 = 10 / N0   # Upper cutoff of the power spectrum
k_high = 256 / N0   # Upper cutoff of the power spectrum
Hurst = 0.6         # Hurst exponent
dim = 2             # Dimension of the random field
seed = 1509          # Seed for the random number generator
np.random.seed(seed)

# Generate the random field
random_field_kl_1 = rfgen.selfaffine_field(dim = dim, N = N0, Hurst = Hurst, k_low = k_low_1, k_high = k_high)
random_field_kl_2 = rfgen.selfaffine_field(dim = dim, N = N0, Hurst = Hurst, k_low = k_low_2, k_high = k_high)

# Normalize
random_field_kl_1 -= np.mean(random_field_kl_1)
random_field_kl_1 /= np.std(random_field_kl_1)

random_field_kl_2 -= np.mean(random_field_kl_2)
random_field_kl_2 /= np.std(random_field_kl_2)

# Plot
plt.rcParams['figure.figsize'] = [10,5]
fig,ax = plt.subplots(1,2)
ax[0].axis('off')
ax[1].axis('off')

a = ax[0].imshow(random_field_kl_1, cmap='RdYlBu_r', interpolation='bicubic', vmin = -3, vmax = 3)
b = ax[1].imshow(random_field_kl_2, cmap='RdYlBu_r', interpolation='bicubic', vmin = -3, vmax = 3)
ax[0].set_title("k_low = 1")
ax[1].set_title("k_low = 10")
fig.colorbar(a, ax=ax[0], shrink=0.82, pad=0.02, orientation='horizontal')
fig.colorbar(b, ax=ax[1], shrink=0.82, pad=0.02, orientation='horizontal')
plt.tight_layout()
fig.savefig("RandomFields.png", dpi=300)
plt.show()

# construct PDF
Nbins = 50
bins = np.linspace(-4, 4, Nbins)
hist1, bins = np.histogram(random_field_kl_1, bins=bins, density=True)
hist2, bins = np.histogram(random_field_kl_2, bins=bins, density=True)

# plot PDF
x = np.linspace(-4, 4, Nbins)
plt.rcParams['figure.figsize'] = [10,5]
plt.rcParams.update({'font.size': 14})
fig,ax = plt.subplots(1,2)
plt.title("PDF of the random field")
ax[0].title.set_text("k_low = 1")
ax[0].grid()
ax[0].set_xlabel("Random field value")
ax[0].set_ylabel("PDF")
ax[0].set_xlim(-4,4)
ax[0].set_ylim(0,0.5)
ax[0].bar(bins[:-1], hist1, width=bins[1]-bins[0], align="edge", color="r", alpha=0.5)
ax[0].plot(x, 1/np.sqrt(2*np.pi)*np.exp(-x**2/2),"-",color="k")

ax[1].grid()
ax[1].title.set_text("k_low = 10")
ax[1].set_xlabel("Random field value")
ax[1].set_ylabel("PDF")
ax[1].set_xlim(-4,4)
ax[1].set_ylim(0,0.5)
# ax[1].plot(bins[:-1], hist2, 'o-')
ax[1].bar(bins[:-1], hist2, width=bins[1]-bins[0], align="edge", color="g", alpha=0.5)
ax[1].plot(x, 1/np.sqrt(2*np.pi)*np.exp(-x**2/2),"-",color="k")
fig.tight_layout()
fig.savefig("PDF.pdf")
plt.show()



# Extract random profile
x = np.random.randint(0,N0)
profile_1 = random_field_kl_1[:,x]
profile_2 = random_field_kl_2[:,x]
x = np.linspace(0,1,N0)
fig,ax = plt.subplots()
plt.xlabel("x")
plt.ylabel("profile")
plt.xlim(0,1)
plt.grid()
plt.plot(x,profile_1,"-",color="r",label="k_low = 1")
plt.plot(x,profile_2,"-",color="g",label="k_low = 10")
plt.legend()
plt.show()

# Classical way to evaluate the correlation function $R(Delta x)$
t0 = time.time()
R_1 = np.zeros(N0-1)
R_2 = np.zeros(N0-1)
for i,shift in enumerate(range(0,N0-1)):
    # Shift profiles
    profile1_shifted = np.roll(profile_1,shift)
    profile2_shifted = np.roll(profile_2,shift)
    # Evaluate product and integrate
    product_1 = np.sum(profile_1*profile1_shifted) / N0
    product_2 = np.sum(profile_2*profile2_shifted) / N0
    # Assign to correlation function
    R_1[i] = product_1
    R_2[i] = product_2

# Normalize correlation function
R_1 /= R_1[0]
R_2 /= R_2[0]
TimeSpent = time.time()-t0
print("\n// Find autocorrelation function using direct integration: CPU time ", TimeSpent)

# FFT 1D transform
t0 = time.time()
zhat = np.fft.fft(profile_1)
PSD = zhat * np.conj(zhat) / N0**2
Rfft_1 = np.fft.ifft(PSD)
Rfft_1 /= Rfft_1[0]
zhat = np.fft.fft(profile_2)
PSD = zhat * np.conj(zhat) / N0**2
Rfft_2 = np.fft.ifft(PSD)
Rfft_2 /= Rfft_2[0]
t1 = time.time()
print("// Find autocorrelation function using FFT transform: CPU time ", t1-t0)
print("// It is faster by a factor of ", TimeSpent/(t1-t0), "\n")

dx = np.linspace(0,1,N0-1)
fig,ax = plt.subplots()
plt.xlabel("Delta x/L")
plt.ylabel("Auto-correlation function")
plt.grid()
plt.xlim(0,0.5)
# find the first zero crossing
cutoff = 0
for i,rr in enumerate(R_1):
    if rr < 0:
       cutoff = i
       break 
plt.plot(dx,R_1[:int(N0)],"-",label="k_low = 1")
plt.plot(dx[::10],Rfft_1[::10],"o",color="b",label="FFT method: k_low = 1")
plt.axhline(y=0,color="k",ls="--")
plt.axvline(x=dx[cutoff],color="b",ls="--")
plt.text(dx[cutoff]+0.005,0.04,"~ Correlation length, $l^*$",fontsize=14,rotation=90)
# 
cutoff = 0
for i,rr in enumerate(R_2):
    if rr < 0:
       cutoff = i
       break 
plt.plot(dx,R_2[:int(N0)],"-",color="r",label="k_low = 10")
plt.plot(dx[::10],Rfft_2[::10],"o",color="r",label="FFT method: k_low = 10")
plt.axvline(x=dx[cutoff],color="r",ls=":")
plt.text(dx[cutoff]+0.005,0.03,"~ Correlation length, $l^*$",fontsize=14,rotation=90)

plt.legend()
fig.savefig("AutoCorrelationFunction.pdf")
plt.show()


# # 2D correlation function and PSD
zhat = np.fft.fft2(random_field_kl_1)
PSD_1 = zhat * np.conj(zhat) / N0**2
R2D_1 = np.fft.ifft2(PSD_1)
R2D_1 = np.real(R2D_1)
R2D_1 /= R2D_1[0,0]

# 2D correlation function and PSD
zhat = np.fft.fft2(random_field_kl_2)
PSD_2 = zhat * np.conj(zhat) / N0**2
R2D_2 = np.fft.ifft2(PSD_2)
R2D_2 = np.real(R2D_2)
R2D_2 /= R2D_2[0,0]


# Plot 2D correlation function
plt.rcParams['figure.figsize'] = [10,5]
plt.rcParams.update({'font.size': 10})
fig,ax = plt.subplots(1,2)
ax[0].set_title("k_low = 1")
ax[1].set_title("k_low = 10")
bar = ax[0].imshow(np.real(R2D_1), cmap='RdYlBu_r', interpolation='bicubic', extent=[0,1,0,1])
cbar = fig.colorbar(bar, ax=ax[0], shrink=0.82, pad=0.1, orientation='horizontal')
# Add contour lines and levels with labels
mylevels = np.linspace(-0.1,1,12)
CS = ax[0].contour(np.real(R2D_1),levels=mylevels,colors=("k","r","k","k","k","k","k","k","k"),extent=[0,1,0,1],linewidths=1,alpha=0.5)
ax[0].clabel(CS, inline=True, fontsize=6)

cx = 0.1
bar2 = ax[1].imshow(np.real(R2D_2[:int(cx*N0),:int(cx*N0)]), cmap='RdYlBu_r', interpolation='bicubic', extent=[0,cx,0,cx], origin="lower")
cbar2 = fig.colorbar(bar2, ax=ax[1], shrink=0.82, pad=0.1, orientation='horizontal')

mylevels = np.linspace(-0.1,1,12)
CS2 = ax[1].contour(np.real(R2D_2[:int(cx*N0),:int(cx*N0)]),levels=mylevels,colors=("k","r","k","k","k","k","k","k","k","k","k"),extent=[0,cx,0,cx],linewidths=1,alpha=0.5)
ax[1].clabel(CS2, inline=True, fontsize=6)
fig.savefig("2DCorrelationFunction.png", dpi=300)
plt.show()

# Plot PSD
plt.rcParams['figure.figsize'] = [10,5]
fig,ax = plt.subplots(1,2)
ax[0].set_title("k_low = 1")
ax[1].set_title("k_low = 10")

im = ax[0].imshow(np.log(np.real(PSD_1[:int(k_high*N0),:int(k_high*N0)])).T, 
                  vmin=-7, vmax=10, interpolation='bicubic', cmap="inferno", 
                  extent=[0,k_high,0,k_high], origin='lower')
fig.colorbar(im, ax=ax[0], orientation='horizontal', shrink=0.82, pad=0.06, label="log(PSD)")   

im1 = ax[1].imshow(np.log(np.real(PSD_2[:int(k_high*N0),:int(k_high*N0)])).T, 
                   vmin=-7, vmax=10, interpolation='bicubic', cmap="inferno", 
                   extent=[0,k_high,0,k_high], origin='lower')
fig.colorbar(im1, ax=ax[1], orientation='horizontal', shrink=0.82, pad=0.06, label="log(PSD)")
# Add contour lines an

plt.tight_layout()
fig.savefig("PSD.png", dpi=300)
plt.show()

# PSD of profile
PSD_profile_1x = PSD_1[0,:]
PSD_profile_1y = PSD_1[:,0]
PSD_profile_2x = PSD_2[0,:]
PSD_profile_2y = PSD_2[:,0]
kl = int(k_low_1*N0)
ks = int(k_high*N0)
k = np.linspace(kl,ks,ks-kl)

plt.rcParams['figure.figsize'] = [10,5]
fig,ax = plt.subplots(1,2)
ax[0].set_title("k_low = 1")
ax[1].set_title("k_low = 10")

ax[0].set_xlabel("k")
ax[0].set_xlim(0,1.1*np.log(ks+5))
ax[0].set_ylim(-10,12)
ax[0].set_ylabel("PSD")
ax[0].grid()
ax[0].plot(np.log(k),np.log(PSD_profile_1x[kl:ks]),"-",color="r",label="ky=0")
ax[0].plot(np.log(k),np.log(PSD_profile_1y[kl:ks]),"-",color="b",label="kx=0")
# fit power law
def func(x, a, b):
    return a - b * x
popt, pcov = curve_fit(func, np.log(k), np.log(0.5*(PSD_profile_1x[kl:ks] + PSD_profile_1y[kl:ks])))
print(popt)
ax[0].plot(np.log(k),func(np.log(k),*popt),"--",color="k",lw=3,label="Fit $k^{-"+str(round(popt[1],2))+"}$")
# ax[0].text(np.log(kl)+1,np.log(PSD_profile_1x[kl])-1,"$k^{-"+str(round(popt[1],2))+"}$",fontsize=14)
ax[0].legend(loc="lower left")

kl = int(k_low_2*N0)
k = np.linspace(kl,ks,ks-kl)
ax[1].set_xlabel("k")
ax[1].set_ylabel("PSD")
ax[1].set_xlim(0,1.1*np.log(ks))
ax[1].set_ylim(-10,12)
ax[1].grid()
ax[1].plot(np.log(k),np.log(PSD_profile_2x[kl:ks]),"-",color="r",label="ky=0")
ax[1].plot(np.log(k),np.log(PSD_profile_2y[kl:ks]),"-",color="b",label="kx=0")
# fit power law
def func(x, a, b):
    return a - b * x
popt, pcov = curve_fit(func, np.log(k), np.log(0.5*(PSD_profile_2x[kl:ks] + PSD_profile_2y[kl:ks])))
print(popt)
ax[1].plot(np.log(k),func(np.log(k),*popt),"--",color="k",lw=3,label="Fit $k^{-"+str(round(popt[1],2))+"}$")
# ax[1].text(np.log(kl)+1,np.log(PSD_profile_2x[kl])-1,"$k^{-"+str(round(popt[1],2))+"}$",fontsize=14)
ax[1].legend(loc="lower left")
fig.tight_layout()
fig.savefig("PSD_profile.pdf")
plt.show()