Untitled

module fDNLS_Direct

using DifferentialEquations, GLMakie, LinearAlgebra, OffsetArrays, NonlinearSolve, LaTeXStrings, MathTeXEngine
NonlinearSolve.ForwardDiff.can_dual(::Type{ComplexF64}) = true
using SpecialFunctions: zeta, gamma
GLMakie.activate!()

@time begin

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

N      = 50 # 2N-1
α      = 0.9
ω      = 50
γ      = 1
ϵ      = 1
h      = 1
#ϵ = ( 2^α * gamma((1+α)/2) ) / (pi^(0.5) * abs(gamma(-α/2)) * h^(α))
#p      = [ϵ]
p = [ϵ, ω] # Parameter list

tspan = (0.0,300)

x=-N+1:N-1

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

    RHS(u,p,t) = -im*(p[1]*coefficientMatrix(α,x)*u - γ*abs2.(u).*u)

    # Define the initial value for the problem with q(x) and the problem with q(x)exp(ix)
    initial_values = Complex.(ground_state(ω, α, x, ϵ))
   # initial_value_with_exp = Complex.(ground_state(ω, α, x, ϵ)).*exp.(im*pi*x)

    # Define the two problems: one for initial value q(x) and the other for q(x)exp(ix)
    prob            = ODEProblem(RHS, initial_values, tspan, p)
    solution_direct = solve(prob, Tsit5(), reltol=1e-8, abstol=1e-8)   

    #=
    prob_with_exp = ODEProblem(RHS, initial_value_with_exp, tspan, p)
    solution_direct_with_exp = solve(prob_with_exp, Tsit5(), reltol=1e-8, abstol=1e-8) 
    =#
    plot_main(solution_direct,x,p,ϵ) # Surface and contour plot of |u|²
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)^(1+α))
        end
    end

    return M.parent   
end

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

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

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

function ground_state(ω, α, x, ϵ)
    #  0 = ω*qₙ + ϵ(Lq)ₙ - qₙ³
   
    RHS(u,p)      = p[2]*u + p[1]*coefficientMatrix(α,x)*u - u.^3

    initialValues = onSite(p[2],α,x) # retrieves the asymptoic onsite sequence Q = [q₋ₙ, ..., q₋₂, q₋₁, q₀, q₁, ..., qₙ] 
    probN         = NonlinearProblem(RHS, initialValues, p)
    sol           = solve(probN, NewtonRaphson(), reltol = 1e-9)

    # norm of RHS sol,p
   # err = norm(RHS(sol, p),2)
   # println(err)

    return sol
end


function plot_main(solution_direct,x,p,ϵ)
    fig = Figure(backgroundcolor=:snow2)

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

    #=
    y_with_exp = solution_direct_with_exp.t
    z_with_exp = [abs(solution_direct_with_exp[i,j])^2 for i in axes(solution_direct_with_exp,1), j in axes(solution_direct_with_exp,2)]
    =#

    #=
    surface_axis=Axis3(fig[1,1],xlabel="Space", ylabel="Time", zlabel="|u|²", title="Initial condition: qₙ = fSolve with onsite initial")
    p1 = surface!(surface_axis, x, y, z)

    surface_axis_with_exp=Axis3(fig[1:2,2],xlabel="Space", ylabel="Time", zlabel="|u|²", title="Initial condition: qₙexp(in)")
    p1_with_exp = surface!(surface_axis_with_exp, x, y_with_exp, z_with_exp)
    =#
    axis_contour = Axis(fig[1,1], xticks = -N+1:2:N-1, xminorticksvisible = true, 
    xlabel=L"h\mathbb{Z}", title="Intensity", subtitle="Initial condition: Onsite asymptotic sequence")
    contourf!(axis_contour, x, y, z)

    axis_initial_condition = Axis(fig[2,1], xticks = -N+1:2:N-1, xminorticksvisible = true, xlabel="space", title=L"-\omega q_{n} = \epsilon\sum_{m \in \mathbb{Z}, m\neq n}{\frac{q_n - q_m}{|n-m|^{1+\alpha}}} - q_n^3",
    subtitle="Paramters: α=$α, ω=$ω, ϵ=$ϵ ")
    #p2  =  scatter!(axis_initial_condition, x, ground_state(ω,α,x,ϵ))
    p2Ground  =  scatter!(axis_initial_condition, x,ground_state(ω,α,x,ϵ), label="fSolve with onsite initial")
    p2onSite  =  scatter!(axis_initial_condition, x, onSite(ω,α,x), label="Onsite asymptotic sequence")

    axislegend()

    #=
    norm_of_solution = [norm(solution_direct.u[i],2) for i in 1:size(solution_direct.u)[1]]
    axis_of_norm = Axis(fig[2,2],xlabel="Time",title="Conservation of 2-norm of solution")
    p3 = scatter!(axis_of_norm, solution_direct.t, norm_of_solution)
    =#

    display(fig)
end

main(N)
end

end