Untitled

using DifferentialEquations, GLMakie, LinearAlgebra, SpecialFunctions
GLMakie.activate!()

#=
The fDNLS is written as iuₙ' = ∑(uₙ - uₘ)/|n-m|^(1+α) - |uₙ|^2 uₙ, where the sum is evaluated such that m \in 1:N and m != n. 
Dirichlet boundary values may be imposed: u₁ = 0, uₙ = 0.

Standing wave solutions are of the form u(x,t) = q(x)exp[i ω t], where q(x) here is the "Ground State Solution" obtained by solution by the 
corresponding system of nonlinear equations after substituting the ansatz. To solve the above system of ODEs, take u(x,0) = q(x). Ref. https://royalsocietypublishing.org/doi/epdf/10.1098/rspa.2014.0364
=#

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

const N = 31    

centerGrid  = (N-1)/2;
x = -centerGrid:centerGrid;

h(j,n,α) = 1/(n-j)^(1+α) + 1/(n+j)^(1+α) # Used in the onSite() function
Q₊  = zeros(Int(centerGrid+1))

function main()

    p   = [1] # ODEProblem() requries a parameter list: `p`. Here, p = [h]
    α   = 0.2
    h   = 1
   
    tspan = (0.0,10.0)
    A = Complex.(diagm(ones(length(x))))
    A[1,1] = 0
    A[length(x),length(x)] = 0 

    RHS(u,p,t) = 1*im*(p[1]*coefficientMatrix(α)*u - A*@.(((abs(u))^2)*u) )

    #initial_values = Complex.(sech.(x))
    initial_values = Complex.(onSite(30,α)) # onSite(ω, α)
    # Equivalent to Ode45 from Matlab
    prob = ODEProblem(RHS, initial_values, tspan, [h])
    sol = solve(prob, Tsit5(), reltol=1e-8, abstol=1e-8)   

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

function diagonal_coefficients(n,α)
    
    if n == 1 || n == N
        diag_value = 0
    else
        diag_value = 1/abs(n-1)^(1+α) + sum(1/abs(n-m)^(1+α) for m in 2:N-1 if m!=n) + 1/abs(n-N)^(1+α) # Entry for M[n,n]. row n, col n
    end
    
    return diag_value
end

function off_diagonal_coefficients(n,m,α)
    
    if n == 1 || n == N
        off_diag_value = 0
    else
        off_diag_value = 1/abs(n-m)^(1+α) # Entry for M[n,m]. row n, col m
    end
    
    return off_diag_value
end

function coefficientMatrix(α) 
    # M is the coefficient matrix corresponding to the 1-dimensional nonlinear discrete operator of the fDNLS
    M = Matrix{Float64}(undef, length(x), length(x))
    for i in 1:size(M,1) # Iterate through the rows
        for j in 1:size(M,2) # Iterate through the columns
            if i==j
                M[i,i] = diagonal_coefficients(i, α)
            elseif i!=j
                M[i,j] = off_diagonal_coefficients(i,j,α)
            end
        end
    end

    return M
end

function onSite(ωᵢ,α)
    # Assumption: qₙ = q₋ₙ, q₀ >> 1 >> q₁ >> ...
    q₀ = sqrt(ωᵢ + 2*zeta(1+α))
    q₁ = q₀/(2*zeta(1+α) - 1/2^(1+α) + ωᵢ)

    Q₊[1] = q₀
    Q₊[2] = q₁

    # Obtain q₂, q₃, ... via equation (3.1)
    for n in 3:size(Q₊)[1]
        Q₊[n] = (q₀/(n)^(1+α) + sum(h(j,n,α)*Q₊[j+1] for j in 1:n-1))/(2*zeta(1+α) - 1/(2*n)^(1+α) + ωᵢ)
    end

    Q₋ = reverse(Q₊[2:size(Q₊)[1]]) # [q₋ₙ, ..., q₂, q₁]. Does not include q₀. 
    Q = append!(Q₋, Q₊)             # [q₋ₙ, ..., q₂, q₁, q₀, q₁, ..., qₙ] 

    return Q
end

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

    y = sol.t
    z= [abs(sol[i,j])^2  for i=1:size(sol)[1], j=1:size(sol)[2]]  # sol[i,j] is the ith component at timestep j. # 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|²", xticks=-centerGrid:2:centerGrid, title="Intensity plot with u(x,0) = q(x)"))
    p2 =  contourf(fig[1,2], x, y, z, axis=(xlabel="Space", ylabel="Time"))
    display(fig)
end

main()