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

f(n,α) = sum(g(n,m,α) for m in grid 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

grid = -N:N # Spatial grid

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.(grid))

    # Equivalent to Ode45 from Matlab
    prob = ODEProblem(RHS, initial_values, tspan, [h])
    sol = solve(prob, Tsit5(), reltol=1e-8, abstol=1e-8)  
    z= [abs(sol[i,j])^2  for j=1:size(sol)[2], i=1:size(sol)[1]] # sol[i,j] is the ith component at timestep j. 

    #plot_main(sol) # Surface and contour plot of |u|²
    return z
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(grid)], 2)...)

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

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

    x = -N:N
    y = sol.t
    z= [abs(sol[i,j])^2  for j=1:size(sol)[2], i=1:size(sol)[1]] # 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=-N:5:N))
    p2 =  contourf(fig[1,2], x, y, z, axis=(xlabel="Space", ylabel="Time"))
end

main()