using Sbplib.Grids, Sbplib.SbpOperators, Sbplib.LazyTensors, Sbplib.RegionIndices
using OrdinaryDiffEq, Plots, Printf, Base.Threads

function apply_tm!(f,u,tm,ind)
    for I in ind
        @inbounds f[I] = (tm*u)[I]
    end
end

function apply_tm_all_regions!(f,u,tm,rinds)
    apply_tm!(f,u,tm,rinds[1])
    apply_tm!(f,u,tm,rinds[2])
    apply_tm!(f,u,tm,rinds[3])
end

region_indices(L, N, ::Lower) = map(x->Index{Lower}(x),1:closure_size(L))
region_indices(L, N, ::Interior) = map(x->Index{Interior}(x),closure_size(L)+1:N-closure_size(L))
region_indices(L, N, ::Upper) = map(x->Index{Upper}(x),N-closure_size(L)+1:N)

function get_region_indices(L,N)
    ind_lower = region_indices(L, N, Lower())
    ind_interior = region_indices(L, N, Interior())
    ind_upper = region_indices(L, N, Upper())
    return (ind_lower, ind_interior, ind_upper)
end

function wave_eq_sim(alg,T,CFL)
    # Domain
    N = 101
    g = EquidistantGrid(N,0.,1.)
    dx = min(spacing(g)...)

    # Spatial discretization
    Δ = Laplace(g,sbp_operators_path()*"standard_diagonal.toml"; order=4)
    (id_l, id_r) = boundary_identifiers(g)
    SAT_l = boundary_condition(Δ,id_l,NeumannBC())
    SAT_r = boundary_condition(Δ,id_r,NeumannBC())
    tm = (Δ + SAT_l + SAT_r)

    # RHS function
    rinds = get_region_indices(Δ,N)
    function f(du,u,p,t)
        du[1:N] .= u[N+1:end]
        apply_tm_all_regions!(view(du,N+1:2*N), view(u,1:N), tm, rinds)
    end
    # Initial condition
    sigma = 0.1
    ic_u = x->1/(sigma*sqrt(2*pi))*exp(-1/2*((x-0.5)^2/sigma^2))
    ic_u_t = x->0
    w0 = [evalOn(g,ic_u);
          evalOn(g,ic_u_t)]
    # Setup ODE and solve
    tspan = (0.,T)
    prob = ODEProblem(f,w0,tspan)
    sol = solve(prob, alg, dt=CFL*dx, saveat=0.05)

    # Plotting
    x = [x[1] for x in points(g)]
    anim = @animate for i ∈ eachindex(sol.t)
        u_i = sol.u[i]
        plot(x, u_i[1:N], ylims = (0,4), lw=3,ls=:dash,label="",title=@sprintf("u at t = %.3f", sol.t[i]))
    end
    gif(anim, "wave.gif", fps = 15)
end

wave_eq_sim(CarpenterKennedy2N54(),1.,0.25)