Untitled

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

#=
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;

f(n,α) = sum(g(n,m,α) for m in 1:N if m!=n) # Diagonal elements for the coefficient matrix M
g(n,m,α) = 1/abs(n-m)^(1+α) # Off-diagonal elements for the coefficient matrix M

function main()

    p   = [1] # ODEProblem() requries a parameter list: `p`. Here, p = [h]
    α   = 1
    h   = 1
   
    tspan = (0.0,200)

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

    initial_values = Complex.(sech.(x))

    # 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|²
end

function coefficientMatrix(α) 
    # M is the coefficient matrix corresponding to the 1-dimensional nonlinear discrete operator of the fDNLS
    M = Matrix{Float64}(undef, repeat([length(x)], 2)...)

    for i in 1:size(M,1)
        M[i,i] = f(i,α) # Coefficient for diagonal element uᵢ
        for j in i+1:size(M,2)
        M[i, j] = -g(i, j, α) # Coefficient for off-diagonal elements
        end
    end

        
    M = Symmetric(M)

    return M
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, yticks=0:100:200))
    p2 =  contourf(fig[1,2], x, y, z, axis=(xlabel="Space", ylabel="Time"))
    display(fig)
end

main()