%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
from clawpack import pyclaw
from clawpack import riemann
from ipywidgets import interact
from ipywidgets import widgets
from exact_solvers import traffic_ramps
from utils import riemann_tools

def c(rho, xi):
    return (1.-2*rho)

def make_plot_function(rho_l,rho_r,D):
    states, speeds, reval, wave_types = traffic_ramps.exact_riemann_solution(rho_l,rho_r,D)
    def plot_function(t):
        ax = riemann_tools.plot_riemann(states,speeds,reval,wave_types,t=t,t_pointer=0,
                                        extra_axes=True,variable_names=['Density']);
        # Characteristic plotting isn't working right for this problem
        riemann_tools.plot_characteristics(reval,c,None,ax[0])
        traffic_ramps.phase_plane_plot(rho_l,rho_r,D,axes=ax[2])
        ax[1].set_ylim(0,1)

        plt.show()
    return plot_function


def riemann_solution(rho_l, rho_r, D):
    plot_function = make_plot_function(rho_l,rho_r,D)
    interact(plot_function, t=widgets.FloatSlider(value=0.1,min=0,max=.9));

"""
Explanation: Traffic flow with an on-ramp
In this chapter we return to the LWR traffic model that we investigated in two earlier chapters.  The LWR model involves a single length of one-way road; in this chapter we will think of this road as a highway.  On a real highway, there are cars entering and leaving the highway from other roads.  In general, real traffic flow must be modeled on a network of roads. the development of continuum traffic models based on LWR and other simple models is an important and very active area of research; see for instance <cite data-cite="holden1995mathematical"><a href="riemann.html#holden1995mathematical">(Holden, 1995)</a></cite> for an investigation of the Riemann problem at a junction, and <cite data-cite="garavello2006traffic"><a href="riemann.html#garavello2006traffic">(Garavello, 2006)</a></cite> for an overview of the area.  Here we take a first step in that direction by considering the presence of a single on-ramp, where traffic enters the highway.
Let the flux of cars from the on-ramp be denoted by $D$; we assume that $D$ is constant in time but concentrated at a single point ($x=0$ in space).  Our model equation then becomes
\begin{align} \label{TFR:balance_law}
    \rho_t + \left(\rho(1-\rho)\right)_x & = D \delta(x),
\end{align}
where $\delta(x)$ is the Dirac delta function.  Equation \eqref{TFR:balance_law} is our first example of a balance law.  The term on the right hand side does not take the form of a flux, and the total mass of cars is not conserved.  We refer to the right-hand-side term as a source term -- quite appropriately in the present context, since it represents a source of cars entering the highway.  In a more realistic model, traffic on the on-ramp itself would also be modeled.  However, our goal here is primarily to illustrate the effect of a source term like that in \eqref{TFR:balance_law} on the solution of the Riemann problem.
Typically, source terms have only an infinitesimal effect on the Riemann solution over short times, since they are distributed in space.  The term considered here is an example of a singular source term; it has a non-negligible effect on the Riemann solution because it is concentrated at $x=0$.
Recall that the flux of cars in the LWR model is given by
$$f(\rho) = \rho(1-\rho)$$
where $0 \le \rho \le 1$.  Thus the maximum flux is $f_\text{max} = 1/4$, achieved when $\rho=1/2$.  We assume always that $D \le 1/4$, so that all the cars arriving from the on-ramp can enter the highway.
As discussed already in the chapter on traffic with a varying speed limit, the flux of cars must be continuous everywhere, and in particular at $x=0$.  Let $\rho^-, \rho^+$ denote the density $\rho$ in the limit as $\xi \to 0$ from the left and right, respectively. Then this condition means that
\begin{align} \label{TFR:source_balance} 
    f(\rho^-) + D = f(\rho^+).
\end{align}
For $D\ne0$, this condition implies that a stationary jump exists at $x=0$, similar to the stationary jump we found in the case of a varying speed limit.
One approach to solving the Riemann problem is to focus on finding $\rho^-$ and $\rho^+$; the wave structure on either side of $x=0$ can then be deduced in the same way we have done for problems without a source term -- connecting $\rho_l$ to $\rho_-$ and $\rho_r$ to $\rho_+$ by entropy-satisfying shock or rarefaction waves.  This approach was undertaken by Greenberg et. al. in <cite data-cite="greenberg1997analysis"><a href="riemann.html#greenberg1997analysis">(Greenberg, 1997)</a></cite> for Burgers' equation; the main results (Table 1 therein) can be transferred to the LWR model in a straightforward way.  As they noted, there is typically more than one choice of $(\rho^+, \rho^-)$ that leads to an entropy-satisfying weak solution; some additional admissibility condition is required in order to choose one.  Herein we will motivate the choice of $\rho^+, \rho^-$ based on physical considerations; the resulting values agree with those of Greenberg et. al. (see also <cite data-cite="isaacson1992nonlinear"><a href="riemann.html#isaacson1992nonlinear">(Isaacson, 1992)</a></cite> for yet another approach that yields the same admissibility conditions).
Spatially-varying fluxes and source terms
The similarity between the existence of an on-ramp at $x=0$ and a change in the speed limit at $x=0$ can be seen mathematically as follows.  For the varying speed limit, we studied a conservation law of the form
$$\rho_t + f(\rho,x)_x =0.$$
Using the chain rule, this is equivalent to
$$\rho_t + f_\rho(\rho,x) \rho_x = - f_x(\rho,x).$$
Hence the variable-coefficient system can also be viewed as a balance law.  If $f$ is discontinuous at $x=0$, then
$f_x$ is a delta function.  Notice that the presence of an on-ramp (positive source term) corresponds to a decrease in the speed limit.  This makes sense -- both of these have the effect of reducing the rate at which cars from upstream ($x<0$) can proceed downstream ($x>0$).  Thus the Riemann solutions we find in this chapter will be similar to those found in the presence of a decrease in speed limit.
In the remainder of the chapter, we investigate the solution of the Riemann problem for this balance law.
Conditions for existence of a solution
In our model, cars entering from the on-ramp are always given priority.  In a real-world scenario, traffic on the on-ramp could also back up and the flux $D$ from the ramp could be decreased.  However, a much more complicated model would be required in order to account for this; see <cite data-cite="delle2014pde"><a href="riemann.html#delle2014pde">(delle Monache, 1992)</a></cite> for an example of such a model.
The flux $D$ from the on-ramp cannot raise the density above $\rho_\text{max}=1$ (representing bumper-to-bumper traffic).  This leads to some restrictions on $D$ in order to guarantee existence of a solution to the Riemann problem:

$D \le 1/4$.  This condition is necessary since otherwise the flux from the on-ramp would exceed the maximum flux of the highway, even without any other oncoming traffic.
If $\rho_r > 1/2$, then  $D \le f(\rho_r)$.  The reason for this is as follows: if $\rho_r > 1/2$, then characteristics to the right of $x=0$ go to the left.  Hence there cannot be any right-going wave (a more detailed analysis shows that a right-going transonic shock is impossible), and it must be that $\rho^+ = \rho_r$.  Thus $D = f(\rho_r) - f(\rho^-) \le f(\rho_r)$.

It turns out that these two conditions are also sufficient for the existence of a solution to the Riemann problem.
End of explanation
"""


rho_l = 0.2
rho_r = 0.2
D = 0.05

riemann_solution(rho_l, rho_r, D)

traffic_ramps.plot_car_trajectories(rho_l,rho_r,D)

"""
Explanation: Light traffic, little inflow
What happens when the on-ramp has a relatively small flux of cars, and the highway around the ramp is not congested?  There will be a region of somewhat higher density starting at the ramp and propagating downstream.  This is demonstrated in the example below.
End of explanation
"""


rho_l = 0.2
rho_r = 0.8
D = 0.05

riemann_solution(rho_l, rho_r, D)

traffic_ramps.plot_car_trajectories(rho_l,rho_r,D,xmax=0.2)

"""
Explanation: In contrast to the LWR model without a ramp, here we see three constant states separated by two waves in the Riemann solution.  The first is a stationary wave where the traffic density abruptly increases due to the cars entering from the ramp, as predicted by \eqref{TFR:source_balance}.  Indeed that condition determines the middle state $\rho_m$ as the solution of
$$f(\rho_l) + D = f(\rho_m)$$
For given values of $\rho_l$ and $D$, this is a quadratic equation for $\rho_m$, with solution
\begin{align} \label{TFR:qm1}
    \rho_m = \frac{1 \pm \sqrt{1-4(f(\rho_l)+D)}}{2}.
\end{align}
As in the case of varying speed limit, we can choose the physically relevant solution by applying the condition that the characteristic speed not change sign at $x=0$.  This dictates that we choose the minus sign, so that $\rho_m<1/2$, since $\rho_l < 1/2$.
Downstream, there is a rarefaction as these cars accelerate and again spread out.
The solution just proposed will break down if either of the following occur:

If the downstream density $\rho_r$ is greater than $\rho_m$, then a shock wave will form rather than a rarefaction.
If the combined flux from upstream and from the ramp exceeds $f_\text{max}$, there will be a shock wave moving upstream due to congestion at the mouth of the on-ramp.  This happens if $f(\rho_l) + D > 1/4$; notice that this is precisely the condition for the value of $\rho_m$ in \eqref{TFR:qm1} to become complex.

We consider each of these scenarios in the following sections.
Uncongested upstream, congested downstream: transonic shock
What if upstream traffic and flux from the on-ramp are light, but traffic is significantly heavier just after the on-ramp?  In this case a shock wave will form, since if $\rho_r > \rho_l$, characteristics from the left and right regions must cross.  The shock may move to the left or right, depending on how congested the downstream segment is.  In either case, there will again be a stationary jump at $x=0$ due to the cars entering from the on-ramp.
End of explanation
"""


rho_l = 0.2
rho_r = 0.2
D = 0.25

riemann_solution(rho_l, rho_r, D)

traffic_ramps.plot_car_trajectories(rho_l,rho_r,D)

"""
Explanation: The influx of cars from the ramp here causes a traffic jam that moves upstream.  As we discuss further below, some real highways limit the influx in order to prevent this.
Experiment with the value of $\rho_r$ and in the example above.  Can you give a precise condition that determines whether the shock will move left or right?
Light traffic, heavy inflow
Now we come to the most interesting case.  Since the maximum flux is $1/4$, it follows that if $f(\rho_l) + D = 1/4$, then the oncoming traffic from the highway and the on-ramp can just fit onto the road at $x=0$.  The smaller value of $\rho_l$ for which this equation holds is $\rho^* = 1/4 - \sqrt{D}$.  If $\rho_l$ exceeds this value, then not all the cars arriving at $x=0$ can  fit on the road there; since our model gives priority to the cars coming from the on-ramp, the road to the left of $x=0$ must suffer a traffic jam -- a shock wave moving to the left.
As long as $\rho_r < 1/2$, the value of $\rho^+$ will be exactly $1/2$, so as to maximize the flux through $x=0$.  Downstream, a rarefaction will form as cars accelerate into the less-congested highway.
End of explanation
"""


rho_l = 0.6
rho_r = 0.2
D = 0.12

riemann_solution(rho_l, rho_r, D)

traffic_ramps.plot_car_trajectories(rho_l,rho_r,D)

"""
Explanation: Notice that in the extreme case that $D=1/4$, the cars from the on-ramp completely block the cars coming from the left; those cars come to a complete stop and never pass $x=0$.  This may seem surprising, since the density of cars to the right of $x=0$ is just $1/2$.  However, since the flux must increase by $1/4$ at $x=0$, it follows that the flux just to the left of $x=0$ must be zero.
Counterintuitively, when two roads merge, limiting the influx of traffic from one or both of them can significantly increase the overall rate of traffic flow.  Contrary to our model, the usual approach is to prioritize the cars already on the highway and restrict the influx of cars from an on-ramp.  This is done in practice nowadays on many highway on-ramps in congested areas.
Congested upstream, uncongested downstream
End of explanation
"""


rho_l = 0.6
rho_r = 0.8
D = 0.12

riemann_solution(rho_l, rho_r, D)

traffic_ramps.plot_car_trajectories(rho_l,rho_r,D)

"""
Explanation: Congested on both sides
Next let us consider what happens if the incoming traffic from the upstream highway and the on-ramp exceeds the maximum flux, but the road is also congested for $x>0$ (i.e., $\rho_r>1/2$).  Then no waves can travel to the right, and a left-going shock will form.  If downstream congestion is too high, then the traffic from the on-ramp will not all be able to enter the highway and no solution is possible in this model (see the second condition for existence, above).
End of explanation
"""


rho_l = 0.1
rho_r = 0.6
D = 0.08
riemann_solution(rho_l, rho_r, D)

traffic_ramps.plot_car_trajectories(rho_l,rho_r,D)

rho_l = 1.0
rho_r = 0.7
D = 0.1

riemann_solution(rho_l, rho_r, D)

traffic_ramps.plot_car_trajectories(rho_l,rho_r,D)

"""
Explanation: Further examples
End of explanation
"""


f = lambda q: q*(1-q)

def plot_all(rho_l,rho_r,D):
    states, speeds, reval, wave_types = traffic_ramps.exact_riemann_solution(rho_l,rho_r,D)
    ax = riemann_tools.plot_riemann(states,speeds,reval,wave_types,t=0.5,extra_axes=2);
    riemann_tools.plot_characteristics(reval,c,None,ax[0])
    traffic_ramps.phase_plane_plot(rho_l,rho_r,D,axes=ax[2],show=False)
    traffic_ramps.plot_car_trajectories(rho_l,rho_r,D,axes=ax[3])

    plt.show()
    
interact(plot_all,
            rho_l = widgets.FloatSlider(min=0.,max=1.,step=0.01,value=0.4,description=r'$\rho_l$'),
            rho_r = widgets.FloatSlider(min=0.,max=1.,step=0.01,value=0.7,description=r'$\rho_r$'),
            D = widgets.FloatSlider(min=0.,max=0.25,step=0.01,value=0.1),
        );

"""
Explanation: Interactive solver
End of explanation
"""


