using Plots
gr()
const g = 1.0
const dx = 0.02

function v0(x)
    if x >= 0.0
        return [1.0, 0.0, 0.0]
    else
        return [2.0, 0.0, 0.0]
    end
end


# V[x, t, q] where q=0 is h, q=1 is hu, q=2 is hv
function initialize_v()
    V = zeros(Int(round(2.0/dx)), 60, 3)

    for i in 1:size(V)[1]
        V[i, 1, :] = v0(-1.0 + (i-1)*dx)
    end
    return V
end


# numerical flux
function F(Vin)
    return [Vin[2], Vin[2]^2/Vin[1] + g * Vin[1]^2/2, Vin[2]*Vin[3]/Vin[1]]
end


function lambda_max(V, i, n)
    if i+1 <= size(V)[1]
        l = abs(0.5*(V[i, n, 2]/V[i, n, 1] + (V[i+1, n, 2]/V[i+1, n, 1])))
        l += sqrt(g * 0.5 * (V[i, n, 1] + V[i+1, n, 1]))
    else
        l = abs(V[i, n, 2]/V[i, n, 1]) + sqrt(g * V[i, n, 1])
    end
    return l
end


# flux at i + 1/2 , n
function Fin(V, i, n)
    lambda = lambda_max(V, i, n)
    if i+1 <= size(V)[1]
        return 0.5*(F(V[i, n, :]) + F(V[i+1, n, :])) - 0.5 * lambda * (V[i+1, n, :] - V[i, n, :])
    else
        return F(V[i, n, :])
    end
end


function main(dx)
    V = initialize_v()
    dts = zeros(60)
    for n in 1:60
        lambdas = [lambda_max(V, i, n) for i = 1:100]
        dt = 0.8 * minimum(dx ./ lambdas)
        dts[n] = dt
        V[1, n+1, :] = V[1, n, :]
        for i in 2:100
            V[i, n+1, :] = V[i, n, :] - dt/dx * (Fin(V, i, n) - Fin(V, i-1, n))
        end
        if sum(dts) > 0.5
            return V[:, 1:n, :], dts[1:n]
        end
    end
    return V, dts
end