# This script implements a partial differential equation describing changes in 
# concentration with respect to longitudinal distance along a stream channel
# due to advection. I try using both the symbolic PDE method (which works)
# and the finite-difference operators method (which doesn't work yet)

##### ----- Define Parameters, Initial Conditions, Boundary Conditions ----- #####
v = 1.0 # Velocity
p = [v] # parameter vector
tspan = (0.0, 10.0) # simulation timespan
xₙ = 10 # physical length of simulation
Δx = 0.01 # grid spacing along x dimension (longitudinal)
N = convert(Int64, xₙ/Δx - 1) # number of interior points
x = range(0, step = Δx, length = (N+2)) # x-points of solution, including boundaries
h(p,t) = zeros(N) # history function describing solute concentrations before the start of the simulation
C₀ = zeros(N) # initial condition
for i in 1:N
  u₀[i] = 1.0
end

##### ----- Automated Finite-Difference for Symbolic PDES Method ----- #####
using ModelingToolkit
@parameters t x v
@variables(u(..))
@derivatives Dt'~t
@derivatives Dx'~x

#1D PDE and boundary conditions
eq = Dt(u(t,x)) ~ -v*Dx(u(t,x))
bcs = [u(0,x) ~ 0,
        u(t,0) ~ 1.0,
        Dx(u(t,10.0)) ~ 0.0]

# Space and time domains
domains = [t ∈ IntervalDomain(0.0,convert(Float64, xₙ)),
           x ∈ IntervalDomain(tspan[1],tspan[2])]

# PDE system
pdesys = PDESystem(eq,bcs,domains,[t,x],[u])

# Method of lines discretization
discretization = MOLFiniteDifference(Δx,3) #Third order finite difference approximation with nodes spaced by 0.01

# Convert the PDE problem into an ODE problem
prob = discretize(pdesys,discretization)

# Solve ODE problem
sol = solve(prob,Tsit5(),saveat=1.0)

# Plot results and compare with exact solution
t = sol.t
x = prob.space[2]

plt = plot()
for i in 1:length(t)
    plot!(x,Array(prob.extrapolation[1](t[i])*sol.u[i]),label="t=$(t[i])")
end
display(plt)

##### ----- Finite Difference Operator Method ----- #####
ΔUpwind = DiffEqOperators.UpwindDifference(1,4,Δx,xₙ) # first derivative operator, fourth-order approximation
Dirichletbc = DiffEqOperators.DirichletBC(1.0,0.0) # boundary condition operator

@advect(ΔUpwind,Dirichletbc)
prob = ODEProblem(advect!, u₀, tspan, p)
alg = Rosenbrock23(autodiff=false)
maxdt = Δx/v
sol = solve(prob, alg, saveat = 1.0, dtmax = 0.8*maxdt)
t = sol.t
u = sol.u

plt = plot()
for i in 1:length(t)
    plot!(x, bc*u[[i]][1] ,label="t=$(t[i])")
end
display(plt)
