using LinearAlgebra, PyPlot

nz = 101
z = range(0, stop=1, length=nz)
dz = z[2]-z[1]

ue = Diagonal(z)
be = Diagonal([2/dz; zeros(nz-2); -2/dz]) # for j=1,nz equations
Inz = Matrix(Diagonal(ones(nz)))
Inz2 = Matrix(Diagonal([1/2; ones(nz-2); 1/2]))
diag, sidediag = -2ones(nz), ones(nz-1)
D2z = Matrix(Tridiagonal(sidediag, diag, sidediag))
D2z = D2z/dz^2

D2zmod = D2z
D2zmod[1, 1:2] = [-2, 2]/dz^2       # for j=1  equation
D2zmod[nz, nz-1:nz] = [2, -2]/dz^2  # for j=nz equation

function create_operators(k)
    Laplacian = D2zmod - k^2*Inz
    invLaplacian = inv(Laplacian)
    # the operator A
    A =invLaplacian*(-im*k*ue*Laplacian + im*k*be)
    # the metric M
    M = Matrix(Tridiagonal(-sidediag, -diag, -sidediag))    
    M[1, 1:2], M[nz, nz-1:nz] = [1, -1], [-1, 1]
    M = M/dz + k^2*dz*Inz2
    Msqrt = sqrt(M)
    invMsqrt = inv(Msqrt)
    # the operator A_M
    AM = Msqrt*A*invMsqrt
    return A, M, AM, Msqrt, invMsqrt
end;

k = range(0, stop=3, length=51)

# the analytic growth rates
mu = k;
growthrate_analytic = real.( k ./ mu .* sqrt.(Complex.(1 ./ tanh.(mu/2)-mu/2).*(mu/2-tanh.(mu/2))));
growthrate_analytic[1] = 0

# the numerical growth rates
growthrate = zeros(length(k))
for j in 1:length(k)
    if k[j] == 0
        growthrate[j] = 0
    else
        A, M, AM = create_operators(k[j])
        growthrate[j] = maximum(real(eigvals(A)))        
    end
end

plot(k, growthrate_analytic, label="analytic")
plot(k, growthrate, "*", label="numerical")
xlabel(L"$k$", fontsize=18)
ylabel(L"$\mathrm{Re}(\sigma)$", fontsize=18)
title("Eady growth rates", fontsize=18)
legend();

function eigenanalysis(A)
    S, U = eigen(A);
    mostunstable = findall(real(S) .== maximum(real(S)))
    return S[mostunstable], U[:, mostunstable]
end

function find_phi0s(k, tstar)
    # finds the eigenvector of A, A^\dagger, A+A^\dagger,
    # and e^{A^dagger t}e^{A t} for t=tstar
    
    A, M, AM, Msqrt, invMsqrt = create_operators(k)
    maxgrowthrate = maximum(real(eigvals(A)))

    s1, u1 = eigenanalysis(AM);
    s1adj, v1 = eigenanalysis(Matrix(adjoint(AM)));
    s1inst, w1 = eigenanalysis(AM + Matrix(adjoint(AM)));
    s1opt, y1 = eigenanalysis(Matrix(adjoint(exp(AM*tstar))*exp(AM*tstar)))
    return s1, u1, s1adj, v1, s1inst, w1, s1opt, y1
end;

function plot_phi0s(;k=1.6, tstar=1)
    # plots the eigenvector of A, A^\dagger, A+A^\dagger,
    # and e^{A^dagger t}e^{A t} for t=tstar
    
    A, M, AM, Msqrt, invMsqrt = create_operators(k)
    s1, u1, s1adj, v1, s1inst, w1, s1opt, y1 = find_phi0s(k, tstar)
    
    nx = 100
    x = range(0, stop=2π, length=nx)
    eikx = Matrix(transpose(exp.(im*k*x)))
    
    # go back to 'normal' coordinates by multiplying with invMsqrt
    u1real = real.((invMsqrt*u1)*eikx)
    v1real = real.((invMsqrt*v1)*eikx)
    w1real = real.((invMsqrt*w1)*eikx)
    y1real = real.((invMsqrt*y1)*eikx)

    fig, axs = subplots(ncols=2, nrows=2, figsize=(6.5, 6.5))
    sca(axs[1]); contour(x, z, u1real)
    title(L"eigvector of $A$", fontsize=14)
    ylabel(L"$z$", fontsize=14)
    sca(axs[2]); contour(x, z, v1real)
    title(L"eigvector of $A^\dagger$", fontsize=14)
    xlabel(L"$x$", fontsize=14); ylabel(L"$z$", fontsize=14)
    sca(axs[3]); contour(x, z, w1real)
    title(L"eigvector of $A+A^\dagger$", fontsize=14)
    sca(axs[4]); contour(x, z, y1real)
    title("eigvector of \$e^{A^\\dagger t_*}e^{A t_*}\$, \$t_*= $tstar\$", fontsize=14)
    xlabel(L"$x$", fontsize=14)
end;


plot_phi0s(k=1.6, tstar=5);

function compute_energy(phi0, tfinal, AM)
    t = range(0, stop=tfinal, length=101)
    energy = zeros(length(t))
    dt = t[2]-t[1]
    expAMdt = exp(AM*dt)
    phit = phi0/sqrt(dot(phi0, phi0))
    for j in 1:length(t)
        energy[j] = dot(phit, phit)
        phit .= expAMdt*phit
    end
    return t, energy
end;

k=1.6
A, M, AM, Msqrt, invMsqrt = create_operators(k)

tstar, tfinal = 5, 10
s1, u1, s1adj, v1, s1inst, w1, s1opt, y1 = find_phi0s(k, tstar)

figure(figsize=(6, 2.8))
t, energy = compute_energy(u1, tfinal, AM)
semilogy(t, energy , label="eigenvector")
t, energy = compute_energy(v1, tfinal, AM)
semilogy(t, energy , label="adjoint")
t, energy = compute_energy(w1, tfinal, AM)
semilogy(t, energy , label="instantaneous")
t, energy = compute_energy(y1, tfinal, AM)
semilogy(t, energy , label="optimal \$t_*= $tstar\$")
semilogy(t, exp.(2*real(s1[1])*t), "--k")
xlabel(L"$t$", fontsize=16); ylabel(L"$E(t)$", fontsize=16)
legend();

plot_phi0s(k=2.6, tstar=5);

k=2.6
A, M, AM, Msqrt, invMsqrt = create_operators(k)

tstar, tfinal = 5, 10
s1, u1, s1adj, v1, s1inst, w1, s1opt, y1 = find_phi0s(k, tstar)

figure(figsize=(6, 2.8))
t, energy = compute_energy(u1, tfinal, AM); semilogy(t, energy , label="eigenvector")
t, energy = compute_energy(v1, tfinal, AM); semilogy(t, energy , label="adjoint")
t, energy = compute_energy(w1, tfinal, AM); semilogy(t, energy , label="instantaneous")
t, energy = compute_energy(y1, tfinal, AM); semilogy(t, energy , label="optimal \$t_*=10\$")
xlabel(L"$t$", fontsize=16); ylabel(L"$E(t)$", fontsize=16)
legend();

k=2.6
A, M, AM, Msqrt, invMsqrt = create_operators(k)

tstar, tfinal = 5, 100
s1, u1, s1adj, v1, s1inst, w1, s1opt, y1 = find_phi0s(k, tstar)

figure(figsize=(6, 2.8))
t, energy = compute_energy(u1, tfinal, AM); semilogy(t, energy , label="eigenvector")
t, energy = compute_energy(v1, tfinal, AM); semilogy(t, energy , label="adjoint")
t, energy = compute_energy(w1, tfinal, AM); semilogy(t, energy , label="instantaneous")
t, energy = compute_energy(y1, tfinal, AM); semilogy(t, energy , label="optimal \$t_*=10\$")
xlabel(L"$t$", fontsize=16); ylabel(L"$E(t)$", fontsize=16)
legend();

tstar, tfinal = 10, 50
k1 = 2.38 # this is <k_crit
A1, M1, AM1, M1sqrt, invM1sqrt = create_operators(k1)
s1, u1, s1adj, v1, s1inst, w1, s1opt, y1 = find_phi0s(k1, tstar)
t1, energy1 = compute_energy(y1, tfinal, AM1)
k2 = 2.42 # this is >k_crit
A2, M2, AM2, M2sqrt, invM2sqrt = create_operators(k2)
s2, u2, s2adj, v2, s2inst, w2, s2opt, y2 = find_phi0s(k2, tstar);
t2, energy2 = compute_energy(y2, tfinal, AM2)
gr1, gr2 = round(real(s1[1]), digits=3), round(real(s2[1]), digits=3)

figure(figsize=(6, 2.8))
semilogy(t1, energy1, label="\$k_1= $k1 \$, \$\\sigma_1 = $gr1\$")
semilogy(t2, energy2, label="\$k_2= $k2 \$, \$\\sigma_2 = $gr2\$")
semilogy(t1, 2e3*exp.(2*real(s1[1])*t1), "--k", linewidth=0.8, label="\$\\propto e^{2\\mathrm{Re}(\\sigma_1)t}\$")
semilogy(t1, 1e3*exp.(2*real(s2[1])*t1), "-.k", linewidth=0.8, label="\$\\propto e^{2\\mathrm{Re}(\\sigma_2)t}\$")
xlabel(L"$t$", fontsize=16); ylabel(L"$E(t)/E(0)$", fontsize=16); legend();

k, tstar = 2.6, 5
A, M, AM, Msqrt, invMsqrt = create_operators(k)
s1, u1, s1adj, v1, s1inst, w1, s1opt, y1 = find_phi0s(k, tstar)

phi0 = y1
phit = phi0 / sqrt(dot(phi0, phi0))

tfinal = 10
t = range(0, stop=tfinal, length = 51)
dt = t[2]-t[1]
energy = zeros(length(t))
expAMdt = exp(AM*dt)

nx = 100
x = range(0, stop=2π, length=nx)
eikx = Matrix(transpose(exp.(im*k*x)))

fig, axs = subplots(ncols=2, nrows=1, figsize=(12, 2.8))

for j in 1:length(t)
    energy[j] = dot(phit, phit)
    phit = expAMdt*phit
    
    phitreal = real.((invMsqrt*phit)*eikx)
    
    sca(axs[1]); cla();
    contour(x, z, phitreal)
    xlabel(L"$x$", fontsize=14)
    ylabel(L"$z$", fontsize=14)
    sca(axs[2]); cla();
    semilogy(t[1:j], energy[1:j])
    xlabel(L"$t$", fontsize=14)
    ylabel(L"$E(t)$", fontsize=14)
    sleep(0.01)
    IJulia.clear_output(true)
    display(fig)
end