using Plots,LaTeXStrings
default(markersize=3,linewidth=1.5)
using LinearAlgebra,DifferentialEquations
include("FNC.jl");

x,Dx,Dxx = FNC.diffper(300,[-4,4])
f = (u,c,t) -> -c*(Dx*u);

u_init = @. 1 + exp(-3*x^2)
IVP = ODEProblem(f,u_init,(0.,3.),2.)
sol = solve(IVP,RK4());

an = @animate for t = LinRange(0,3,120) 
    plot(x,sol(t),
        xaxis=(L"x"),yaxis=([1,2],L"u(x,t)"),    
        title="Advection equation, t=$(round(t,digits=2))",leg=:none )
end
gif(an,"advection.gif")

rho_c = 1080;  rho_m = 380;  q_m = 10000;
Q0prime(rho) = q_m*4*rho_c^2*(rho_c-rho_m)*rho_m*(rho_m-rho)/(rho*(rho_c-2*rho_m) + rho_c*rho_m)^3;

x,Dx,Dxx = FNC.diffper(800,[0,4]);

ode = (rho,ep,t) -> -Q0prime.(rho).*(Dx*rho) + ep*(Dxx*rho);

rho_init = @. 400 + 10*exp(-20*(x-3)^2)
IVP = ODEProblem(ode,rho_init,(0.,1.),0.02)
sol = solve(IVP,alg_hints=[:stiff]);

plot(x,[sol(t) for t=0:.2:1],label=["t=$t" for t=0:.2:1],
    xaxis=(L"x"),yaxis=("car density"),title="Traffic flow")     

an = @animate for t = LinRange(0,0.9,91) 
    plot(x,sol(t),
        xaxis=(L"x"),yaxis=([400,410],"density"),    
        title="Traffic flow, t=$(round(t,digits=2))",leg=:none )
end
gif(an,"traffic1.gif")

rho_init = @. 400 + 80*exp(-16*(x-3)^2)
IVP = ODEProblem(ode,rho_init,(0.,0.5),0.02)
sol = solve(IVP,alg_hints=[:stiff]);

an = @animate for t = LinRange(0,0.5,101) 
    plot(x,sol(t),
        xaxis=(L"x"),yaxis=([400,480],"density"),    
        title="A traffic jam, t=$(round(t,digits=2))",leg=:none )
end
gif(an,"traffic2.gif")

x,Dx = FNC.diffper(400,[0 1])
uinit = @. exp(-80*(x-0.5)^2);

ode = (u,c,t) -> -c*(Dx*u);
IVP = ODEProblem(ode,uinit,(0.,2.),2.)
sol = solve(IVP,RK4());

t = 2*(0:80)/80
u = hcat([sol(t) for t in t]...)
contour(x,t,u',color=:viridis,
    xaxis=(L"x"),yaxis=(L"t"),title="Linear advection",leg=:none)

avgtau1 = sum(diff(sol.t))/(length(sol.t)-1)

x,Dx = FNC.diffper(800,[0 1])
uinit = @. exp(-80*(x-0.5)^2);
IVP = ODEProblem(ode,uinit,(0.,2.),2.)
sol = solve(IVP,RK4());

avgtau2 = sum(diff(sol.t))/(length(sol.t)-1)
@show ratio = avgtau1 / avgtau2

n = 100
x,Dx = FNC.diffmat2(n,[0 1])
uinit = @. exp(-80*(x-0.5)^2);

trunc = u -> u[1:n];  extend = v -> [v;0];
ode = (v,c,t) -> -c*trunc(Dx*extend(v));
IVP = ODEProblem(ode,uinit[1:n],(0.,1.),-1)
sol = solve(IVP,RK4());

plot(x[1:n],sol.t,sol[:,:]',
    xlabel=L"x",ylabel=L"t",title="Inflow boundary condition",leg=:none)

trunc = u -> u[2:n+1];  extend = v -> [0;v];
sol = solve(IVP,RK4());
plot(x[1:n],sol.t,sol[:,:]',
    xlabel=L"x",ylabel=L"t",title="Outflow boundary condition",leg=:none)

x,Dx = FNC.diffper(40,[0,1])
lambda = eigvals(Dx);

scatter(real(lambda),imag(lambda),xlim=[-40,40],ylim=[-40,40],
    title="Eigenvalues for pure advection",leg=:none)     

zc = @.exp(1im*2pi*(0:360)/360);    # points on |z|=1

z = zc .- 1;                        # shift left by 1
plot(Shape(real(z),imag(z)),color=RGB(.8,.8,1),layout=(1,2),subplot=1)
scatter!(real(0.1*lambda),imag(0.1*lambda),subplot=1,
    xlim=[-5,5],ylim=[-5,5],aspect_ratio=1,title="Euler",leg=:none)

z = zc .+ 1;                        # shift right by 1
plot!(Shape([-6,6,6,-6],[-6,-6,6,6]),color=RGB(.8,.8,1),subplot=2)
plot!(Shape(real(z),imag(z)),color=:white,subplot=2)
scatter!(real(0.1*lambda),imag(0.1*lambda),subplot=2,
    xlim=[-5,5],ylim=[-5,5],aspect_ratio=1,title="Backward Euler",leg=:none)

plot([],[],label="",leg=:left,aspect_ratio=1)
x,Dx,Dxx = FNC.diffper(40,[0,1]);
tau = 0.1
for ep = [0.001 0.01 0.05]
  lambda = eigvals(-Dx + ep*Dxx)
  scatter!(real(tau*lambda),imag(tau*lambda),label="\\epsilon=$ep")
end
title!("Eigenvalues for advection-diffusion")     

x,Dx,Dxx = FNC.diffcheb(40,[0,1]);
A = -Dx[2:end,2:end];    # leave out first row and column

lambda = eigvals(A)
scatter(real(lambda),imag(lambda),
    xlim=[-300,100],title="Eigenvalues of advection with zero inflow",leg=:none,aspect_ratio=1) 

maximum( real(lambda) )

m = 200
x,Dx = FNC.diffmat2(m,[-1,1]);

trunc = u -> u[2:m];
extend = v -> [0;v;0];

dwdt = function(w,c,t)
    u = extend(w[1:m-1])
    z = w[m:2*m]
    dudt = Dx*z
    dzdt = c^2*(Dx*u)
    return [ trunc(dudt); dzdt ]
end;

u_init = @.exp(-100*(x)^2)
z_init = -u_init
w_init = [ trunc(u_init); z_init ];  

IVP = ODEProblem(dwdt,w_init,(0.,2.),2)
sol = solve(IVP,RK4());

n = length(sol.t)-1
U = [ zeros(1,n+1); sol[1:m-1,:]; zeros(1,n+1) ];
Z = sol[m:2*m,:];

an = @animate for (i,t) = enumerate(sol.t) 
    plot(x,U[:,i],layout=(1,2),subplot=1,
        xaxis=(L"x"),yaxis=([-1,1],L"u(x,t)"),    
        title="Wave equation, t=$(round(t,digits=3))",leg=:none )
    plot!(x,sol.t,U',subplot=2,xlabel=L"x",ylabel=L"t",title="Space-time view",leg=:none)
end
gif(an,"wave.gif")

m = 120
x,Dx = FNC.diffcheb(m,[-1,1]);
c = @. 1 + (sign(x)+1)/2
trunc = u -> u[2:m];
extend = v -> [0;v;0];

dwdt = function(w,c,t)
    u = extend(w[1:m-1])
    z = w[m:2*m]
    dudt = Dx*z
    dzdt = c.^2 .* (Dx*u)
    return [ trunc(dudt); dzdt ]
end;

u_init = @.exp(-100*(x+0.5)^2);
z_init = -u_init;
w_init = [ trunc(u_init); z_init ];    
IVP = ODEProblem(dwdt,w_init,(0.,5.),c)
sol = solve(IVP,RK4());

an = @animate for t = LinRange(0,5,200)
    plot(x,extend(sol(t,idxs=1:m-1)),
        xaxis=(L"x"),yaxis=([-1,1],L"u(x,t)"),    
        title="Wave equation, variable speed",label="t=$(round(t,digits=2))" )
end
gif(an,"wavereflect.gif")