# imports
using Distributions
using Plots
using Optim
using HDF5

# the number of increases in resolution
I = 5;

# output location
output = "$(@__DIR__)/../output/pbc"
processed = "$(@__DIR__)/../processed/pbc"

# get desired outputs
y  = []
dy = []
t  = []
x  = []
λ  = []
M  = []
N  = []
Δt = []
Δx = []

for n = 0:I
    h5open("$output/n=$n.h5", "r") do file
        push!(y , read(file["y"]))
        push!(dy, read(file["dy"]))
        push!(t , read(file["t"]))
        push!(x , read(file["x"]))
        push!(λ , read(attributes(file)["λ"]))
        push!(M , read(attributes(file)["M"]))
        push!(N , read(attributes(file)["N"]))
        push!(Δt, read(attributes(file)["Δt"]))
        push!(Δx, read(attributes(file)["Δx"]))
    end
end

# λ, L and T are constants across all simulations
λ = λ[1]
L = x[1][length(x[1])]
T = t[1][length(t[1])];

# analytical solution
function sol(t, x, T, L)
    # gaussian wave characteristics
    μ  = L/2
    σ  = L/10

    # this correction term comes from the fact that the pbc force me to have "ghost points"
    # this avoids issues at the boundariy where each wave would see a different step going from x=L to x=0
    L += Δx[1]

    # we want to wrap around the moving peak, μ(t) ≡ λt - μ, in the spatial box
    # this requires us to do the modulus of ct, where c = 1, with L
    # however, we want (ct)%L = L, so we need to have a custom modulus to ensure that
    if t != 0 && t%L == 0
        t = L
    else
        t = t%L
    end
    
    # compute the distance to the moving peak μ(t) which is now being wrapped in x-axis
    # use periodic boundary conditions to see which of the two distances is the shortest
    ds = abs(x - t - μ)
    if ds >= L/2
        ds = L - ds
    end

    return 1/(sqrt(2*π)*σ) * exp(-0.5*(ds/σ)^2)
end;

# compute the analytical solution with n=0 grid (for speed)
S = fill!(Array{Float64, 2}(undef, M[1], N[1]), NaN)
for i=1:M[1], j=1:N[1]
    S[i,j] = sol((i-1)*Δt[1], (j-1)*Δx[1], T, L)
end

# useful plot limits
yupper = 0.065
ylower = -0.005;

# plot the heatmap for the analytical solution
heatmap(x[1], t[1], S)
title!("Analytical solution")

# make an heatmap in t and x for the n-th solution
n = 0
heatmap(x[n+1], t[n+1], y[n+1])
xaxis!("x")
yaxis!("t")
title!("Solution for n = $n")

# pointwise diff between all n's and the analytical solution
DS = []
for n=0:I
    Dₙ = fill!(Array{Float64, 2}(undef, M[1], N[1]), NaN)
    for i=1:M[1]
        for j=1:N[1]
            Dₙ[i,j] = y[n+1][2^n*i-(2^n-1),2^n*j-(2^n-1)] - S[i,j]
        end
    end

    push!(DS, Dₙ)
end

# pointwise difference between n'th solution and analytical
n = 0
heatmap(x[n+1], t[n+1], DS[n+1])
xaxis!("x")
yaxis!("t")
title!("Pointwise diff between n = $n and analytical")

# compute the norn of a vector using simpson's 1/3 integration method
function simpson(V, h)
    # length of the vector
    len = length(V)

    # if vector is even then the left over points will be integrated by trapezoid method
    even = false
    if len%2 == 0
        even = true
    end

    # simpson's 1/3 method
    sum = 0
    for i=1:2:len-2
        sum += 2h/6 * (V[i]^2 + 4V[i+1]^2 + V[i+2]^2)
    end

    # integrate left over points
    if even
        sum += (V[len]^2 - V[len-1]^2)/h
    end

    return sqrt(sum)
end;

# compute the pointwise difference between n+1 and n for all n's
Q = []
D = []
for n=1:I
    Qₙ = fill!(Vector{Float64}(undef, M[n]), NaN)
    Dₙ = fill!(Array{Float64, 2}(undef, M[n], N[n]), NaN)

    for i=1:M[n]
        for j=1:N[n]
            Dₙ[i,j] = y[n+1][2*i-1,2*j-1] - y[n][i,j]
        end
        Qₙ[i] = simpson(Dₙ[i,:], Δx[n])
    end

    push!(Q, Qₙ)
    push!(D, Dₙ)
end

# plot Q(t)
p = plot()
for n=1:I
    plot!(p, t[n], Q[n], label="n=$(n) - n=$(n-1)")
end
xaxis!(p, "t")
yaxis!(p, "norm")

# look for the best fits to make the n-th + 1 curve fit the n-th curve
x₀ = [4.]  # initial guess
for n=1:I-1
    result = optimize(x -> sum((Q[n] .- x[1]*Q[n+1][1:2:length(Q[n+1])]).^2), x₀)
    println("Multiplier factor between n=$n and n=$(n-1): $(result.minimizer[1])")
end

# heatmap plot with the difference between the n'th solution and the n'th+1 solution
n = 0
heatmap(x[n+1], t[n+1], D[n+1])
xaxis!("x")
yaxis!("t")
title!("Pointwise diff between n=$(n+1) and n=$n")

# plot several n's compared to the analytical solution
anim = @animate for i=1:M[1]
    plot(x[1], S[i,:], label="Analytical", xlim=(0, L), ylim=(ylower, yupper), leg=:bottom, dpi=150)
    for n=0:I
        plot!(x[n+1], y[n+1][2^n*i-(2^n-1),:], label="Numerical (n = $n)")
    end
end
mp4(anim, "$processed/several.mp4", fps=32)
anim = nothing  # CPU and RAM usage is crazy if this is left hanging around for some reason... clear it!
GC.gc()         # .

# difference between the n-th solution and the analytical solution, for all n's
anim = @animate for i=1:M[1]
    plot()
    for n=0:I
        plot!(x[1], y[n+1][2^n*i-(2^n-1),1:2^n:N[n+1]] .- S[i,:], label="n = $n")
    end
end
mp4(anim, "$processed/diff-several.mp4", fps=32)
anim = nothing  # CPU and RAM usage is crazy if this is left hanging around for some reason... clear it!
GC.gc()         # .

p = plot(leg=:top)
for n=0:I
    plot!(p, t[1], DS[n+1][:,1], label="n = $n", leg=:top)
end
xaxis!(p, "t")
yaxis!(p, "u(t, 0)")
title!(p, "Difference from numerical to analytical at x = 0")
display(p)

p = plot(leg=:top)
for n=0:I
    plot!(p, t[1], DS[n+1][:,N[1]], label="n = $n", leg=:top)
end
xaxis!(p, "t")
yaxis!(p, "u(t, 0)")
title!(p, "Difference from numerical to analytical at x = L")
display(p)