Untitled

module fDNLS_Direct
using DifferentialEquations, GLMakie, LinearAlgebra, OffsetArrays

using SpecialFunctions: zeta
GLMakie.activate!()

#=
DifferentialEquations - Used for the ODE Solver
GLMakie               - 2D/3D plotting backend,
LinearAlgebra,
SpecialFunctions      - Used for importing the Riemann zeta function
=#

function main(N, end_time)
    #=
    Per the boundary conditions, solve for u = (u_{-N+1}, ..., u_{N-1})
    =#

    h      = 1
    α      = 0.2 
    tspan = (0.0,25.0)
    x = -N+1:N-1

    # Per u_(-N) = 0 and u_(N) = 0, then we solve the (2N+1) - 2 = 2N-1 total equations
    RHS(u,p,t) = im*(h*coefficientMatrix(α,x)*u - abs2.(u).*u)

   # initial_values = 0.9*Complex.(sech.(x)).* exp.(im*x)

    initial_values = onSite(30,α,x)# .* exp.(im*x) # onSite(ω, α)
    # Equivalent to Ode45 from Matlab
    prob = ODEProblem(RHS, initial_values, tspan)
    sol = solve(prob, Tsit5(), reltol=1e-8, abstol=1e-8)   

    plot_main(sol, x) # Surface and contour plot of |u|²
    return sol
end

function coefficientMatrix(α,x) 

    M = OffsetArray(zeros(length(x), length(x)), x,x) # Offset indices for M[i,j]. For example, M[-N+1, -N+1] corresponds to M[1,1]. 
    
    for i in firstindex(M,1) : lastindex(M,1)
        for j in axes(M,2)
            i == j ? M[i,i] = 1/abs(i - (-N))^(1+α) + sum(1/abs(i - m)^(1+α) for m in x if m!=i) + 1/abs(i - N)^(1+α) : M[i,j] = -1/abs(i - j)
        end
    end

    return M.parent   
end

function onSite(ω,α,x)
    # Assumption: qₙ = q₋ₙ, q₀ >> 1 >> q₁ >> ...
    Q = OffsetVector(zeros(Complex{Float64},length(x)),x)
    Q[0] = sqrt(ω + 2*zeta(1+α))

    # Obtain q₂, q₃, ... via equation (3.1)
    for n in 1:maximum(x) #drop first 1 since already initialized
        asymptotic_onsite = sum(Q[j]/(n+j)^(1+α) for j in 1-n:n-1)/(2*zeta(1+α) - 1/(2n)^(1+α) + ω)
        Q[-n] = asymptotic_onsite
        Q[n]  = asymptotic_onsite
    end

    return Q.parent # typeof(Q) = OffsetVector, typeof(Q.parent) = Vector
end

function plot_main(sol,x)
    fig = Figure(backgroundcolor=:snow2)

    y = sol.t
    z= [abs(sol[i,j])^2 for i in axes(sol,1), j in axes(sol,2)]  # sol[i,j] is the ith component at timestep j.

    p1 = surface(fig[1,1], x, y, z,
        axis=(type=Axis3, xlabel="Space", ylabel="Time", zlabel="|u|²", 
            title=L"\text{Intensity plot with } u(x,0) = \text{ asymptotic onsite sequence}"))

    #p2 =  contourf(fig[1,2], x, y, z, axis=(xlabel="Space", ylabel="Time"))
    #p2  =  scatter(fig[1,2], x, onSite(30,0.2,x), xlabel="Space", ylabel="qₙ", title="Asymptotic onsite sequence")

    normSol_inf = [norm(sol[i], Inf) for i in 1:length(sol.t)]

    p3 = lines(fig[2,:], y, normSol_inf, label=L"k=1", axis=(xlabel="Time", ylabel=L"L^{\infty}", xticks=0:0.5:25))
    axislegend()
    display(fig)
end

const N=18
const end_time=20

main(N, end_time)
end