using Kinetic, Plots, LinearAlgebra, JLD2
using ProgressMeter: @showprogress

cd(@__DIR__)
D = Dict{Symbol,Any}()
begin
    D[:matter] = "gas"
    D[:case] = "cavity"
    D[:space] = "2d1f2v"
    D[:flux] = "kfvs"
    D[:collision] = "bgk"
    D[:nSpecies] = 1
    D[:interpOrder] = 2
    D[:limiter] = "vanleer"
    D[:boundary] = "maxwell"
    D[:cfl] = 0.5
    D[:maxTime] = 1.0

    D[:x0] = -1.5
    D[:x1] = 1.5
    D[:nx] = 40
    D[:y0] = -1.5
    D[:y1] = 1.5
    D[:ny] = 40
    D[:pMeshType] = "uniform"
    D[:nxg] = 0
    D[:nyg] = 0

    D[:umin] = -5.0
    D[:umax] = 5.0
    D[:nu] = 24
    D[:vmin] = -5.0
    D[:vmax] = 5.0
    D[:nv] = 24
    D[:vMeshType] = "rectangle"
    D[:nug] = 0
    D[:nvg] = 0

    D[:knudsen] = 0.5
    D[:mach] = 0.0
    D[:prandtl] = 1.0
    D[:inK] = 0.0
    D[:omega] = 0.81
    D[:alphaRef] = 1.0
    D[:omegaRef] = 0.5

    D[:uLid] = 0.15
    D[:vLid] = 0.0
    D[:tLid] = 1.0
end

ks = SolverSet(D)
ctr, a1face, a2face = init_fvm(ks, ks.ps, :dynamic_array; structarray = true)

s2 = 0.03^2
s1 = 0.1#4π * s2
flr = 0.0001
init_density(x, y) = max(flr, 1.0 / s1 * exp(-(x^2 + y^2) / 4.0 / s2))

function init_field!(ks, ctr, a1face, a2face)
    for i = 1:ks.ps.nx, j = 1:ks.ps.ny
        rho0 = init_density(ks.ps.x[i, j], ks.ps.y[i, j])

        ctr[i, j].prim .= [rho0, 0.0, 0.0, 1.0]
        ctr[i, j].w .= prim_conserve(ctr[i, j].prim, ks.gas.γ)
        ctr[i, j].f .= maxwellian(ks.vs.u, ks.vs.v, ctr[i, j].prim)
    end
    for i in eachindex(a1face)
        a1face[i].fw .= 0.0
        a1face[i].ff .= 0.0
    end
    for i in eachindex(a2face)
        a2face[i].fw .= 0.0
        a2face[i].ff .= 0.0
    end
end

init_field!(ks, ctr, a1face, a2face)

t = 0.0
dt = timestep(ks, ctr, t)
nt = Int(ks.set.maxTime ÷ dt) + 1
res = zero(ks.ib.wL)

@showprogress for iter = 1:20#nt
    reconstruct!(ks, ctr)
    #evolve!(ks, ctr, a1face, a2face, dt; mode = Symbol(ks.set.flux), bc = Symbol(ks.set.boundary))
    
    # horizontal flux
    @inbounds Threads.@threads for j = 1:ks.pSpace.ny
        for i = 2:ks.pSpace.nx
            flux_kfvs!(
                a1face[i, j].fw,
                a1face[i, j].ff,
                ctr[i-1, j].f,
                ctr[i, j].f,
                ks.vSpace.u,
                ks.vSpace.v,
                ks.vSpace.weights,
                dt,
                a1face[i, j].len,
            )
        end
    end
    
    # vertical flux
    vn = ks.vSpace.v
    vt = -ks.vSpace.u
    @inbounds Threads.@threads for j = 2:ks.pSpace.ny
        for i = 1:ks.pSpace.nx
            wL = KitBase.local_frame(ctr[i, j-1].w, 0., 1.)
            wR = KitBase.local_frame(ctr[i, j].w, 0., 1.)
            swL = KitBase.local_frame(ctr[i, j-1].sw[:, 2], 0., 1.)
            swR = KitBase.local_frame(ctr[i, j].sw[:, 2], 0., 1.)

            KitBase.flux_kfvs!(
                a2face[i, j].fw,
                a2face[i, j].ff,
                ctr[i, j-1].f,
                ctr[i, j].f,
                vn,
                vt,
                ks.vSpace.weights,
                dt,
                a2face[i, j].len,
            )

            a2face[i, j].fw .= KitBase.global_frame(a2face[i, j].fw, 0., 1.)
        end
    end
    
    # boundary flux
    @inbounds Threads.@threads for j = 1:ks.pSpace.ny
        KitBase.flux_boundary_maxwell!(
            a1face[1, j].fw,
            a1face[1, j].ff,
            ks.ib.bcL,
            ctr[1, j].f,
            ks.vSpace.u,
            ks.vSpace.v,
            ks.vSpace.weights,
            dt,
            ctr[1, j].dy,
            1.,
        )

        KitBase.flux_boundary_maxwell!(
            a1face[ks.pSpace.nx+1, j].fw,
            a1face[ks.pSpace.nx+1, j].ff,
            ks.ib.bcL,
            ctr[ks.pSpace.nx, j].f,
            ks.vSpace.u,
            ks.vSpace.v,
            ks.vSpace.weights,
            dt,
            ctr[ks.pSpace.nx, j].dy,
            -1.,
        )
    end
    
    @inbounds Threads.@threads for i = 1:ks.pSpace.nx
        KitBase.flux_boundary_maxwell!(
            a2face[i, 1].fw,
            a2face[i, 1].ff,
            ks.ib.bcL,
            ctr[i, 1].f,
            vn,
            vt,
            ks.vSpace.weights,
            dt,
            ctr[i, 1].dx,
            1,
        )
        a2face[i, 1].fw .= KitBase.global_frame(a2face[i, 1].fw, 0., 1.)
        
        KitBase.flux_boundary_maxwell!(
            a2face[i, ks.pSpace.ny+1].fw,
            a2face[i, ks.pSpace.ny+1].ff,
            ks.ib.bcL,
            ctr[i, ks.pSpace.ny].f,
            vn,
            vt,
            ks.vSpace.weights,
            dt,
            ctr[i, ks.pSpace.ny].dy,
            -1,
        )
        a2face[i, ks.pSpace.ny+1].fw .= KitBase.global_frame(
            a2face[i, ks.pSpace.ny+1].fw,
            0.,
            1.,
        )
    end

    update!(ks, ctr, a1face, a2face, dt, res; coll = Symbol(ks.set.collision), bc = Symbol(ks.set.boundary))

    t += dt
end

plot_contour(ks, ctr)
