from numpy import *
import matplotlib.pyplot as plt

# make sure warnings raise errors
seterr(all='raise')


def run_simulation(L, u0, dt, timesteps, nu=1):
    N = len(u0)
    # compute a vector storing the frequencies for the fourier transform, it already takes the interval [0, L]
    # into account, no need to divde with / L later in the computations
    k = fft.rfftfreq(N, d=L/(N-1))
    solution = zeros((timesteps+1, N))
    solution[0] = u0

    previous_u2k = fft.rfft(u0 ** 2)
    current_uk = fft.rfft(u0)
    # f_L (k), this is constant so we only compute it once.
    # note that we don't need to divide by /L because fftfreq already does it :)
    fL = (2*pi * k) ** 2 - nu * (2 * pi * k) ** 4

    for i in range(timesteps):
        # update solution (note that we don't divide by L!)
        current_u2k = fft.rfft(solution[i] ** 2)
        next_uk = (1 + dt/2 * fL) / (1 - dt/2 * fL) * current_uk - 1j * pi * k * (3 * current_u2k - previous_u2k) / (2 - dt * fL) * dt
        # store after inverse FFT, ignoring its imaginary part
        solution[i+1] = real(fft.irfft(next_uk))
        # update values for the next iteration
        current_uk = next_uk
        previous_u2k = current_u2k

    return solution


def wrap_simulation(L, N, dt=0.05, tmax=200, nu=1):
    x = linspace(0, L, N)
    time_values = linspace(0, tmax, int(tmax / dt)+1)
    s = run_simulation(L, cos(2*pi/L * x) + 0.1 * cos(4*pi/L * x), dt, len(time_values)-1, nu=nu)
    return x, time_values, s


# define a few constants
x, time_values, s = wrap_simulation(100, 1024, nu=1)
# pick a nice color for the plot
colormap = plt.get_cmap("jet")
# print the values
plt.pcolormesh(x, time_values, s, cmap=colormap)
# print the colorbar on the right
plt.colorbar()
# open a window with the plot :)
plt.show()

# ok now let's be real and try to find the critical value of L :)
Ls = linspace(1, 100, 100)


def test_multiple_Ls(nu=1):
    As = zeros(Ls.shape)
    critical_L = None
    for i, L in enumerate(Ls):
        # for every L, we compute the norm A
        #print("Computing for L = %f" % L)
        tmax = 200 # testing shows that a time of 200 s is clearly enough
        _, _, s = wrap_simulation(L, 1024, dt=0.2, tmax=tmax, nu=nu)
        As[i] = sqrt(mean(s[-1]**2)) # compute with the last iteration
        if As[i] > 1e-1 and critical_L is None:
            critical_L = L
    return Ls, As, critical_L


# Plot A as a function of L for nu = 1
Ls, As, critical_L = test_multiple_Ls(1)
plt.plot(Ls, As)
plt.xlabel("L")
plt.ylabel("A")
plt.show()

# Now find the critical value for multiple nus
nus_exp = linspace(-0.5, 1.75, 20)
# nu between e^-0.5 and e^1.75
nus = exp(nus_exp) # use a log scale
critical_Ls = zeros(nus.shape)
for i, nu in enumerate(nus):
    _, _, critical_L = test_multiple_Ls(nu)
    critical_Ls[i] = critical_L
    print("Found critical L %10f for nu %10f" % (critical_L, nu))
plt.plot(nus, critical_Ls, label="Simulation")

# Perform a linear regression to find the slope
from sklearn.linear_model import LinearRegression

reg = LinearRegression().fit(log(nus).reshape(-1, 1), log(critical_Ls)) # we fit the logs
# because we fitted the logs, we have the predict with the log, then take the exp()
plt.plot(nus, exp(reg.predict(log(nus).reshape(-1, 1))), label="Linear Regression (slope = %f)" % reg.coef_)
plt.xlabel("$\\nu$")
plt.ylabel("Critical L")
plt.xscale("log")
plt.yscale("log")
plt.legend()
plt.show()


# notes
# nu = 0.1 causes the solution to diverge!
# for nu above ~6 the critical L is 1
# [0.60653066, 0.68278248, 0.76862053, 0.86524996, 0.97402745, 1.09648023, 1.23432754, 1.38950472, 1.56419047, 1.76083737, 1.98220632, 2.23140533, 2.51193314, 2.82772834, 3.18322469, 3.58341333, 4.03391288, 4.54104833, 5.11193983, 5.75460268]
# [ 5.0,  6.0,  6.0,  6.0,  7.0,  7.0,  7.0,  8.0,  8.0,  9.0,  9.0, 10.0, 10.0, 11.0, 12.0, 12.0, 13.0, 14.0, 14.0, 15.0 ]