from m3t2 import rho_red_light, flux
# help(rho_red_light)

import numpy 
from matplotlib import pyplot
from matplotlib import rcParams
rcParams['font.family'] = 'serif'
rcParams['font.size'] = 16
from matplotlib import animation

nx = 100
nt = 30
dx = 4.0/nx

rho_in = 5.
rho_max = 10.

V_max = 1.

x = numpy.linspace(0+.5*dx, 4-.5*dx, nx)

rho = rho_red_light(nx, rho_max, rho_in)

def animate(data):
    x = numpy.linspace(0+0.5*dx, 4-0.5*dx, nx)
    y = data
    line.set_data(x,y)
    return line,

def godunov(rho, nt, dt, dx, rho_max, V_max):
	""" Computes the solution with the Godunov scheme using the Lax-Friedrichs flux.

	Parameters
	----------
	rho 	: array of floats
			  Density at current time-loop
	nt 		: int
			  Number of time steps
	dt 		: float
			  Time-step size
	dx 		: float
			  Mesh spacing
	rho_max : float
			  Maximum allowed car density
	V_max 	: Speed limit

	Returns
	-------
	rho_n 	: array of floats
			  Density after nt time steps at every point x
	"""

	rho_n = numpy.zeros((nt, len(rho)))
	rho_n[0,:] = rho.copy()

	rho_plus = numpy.zeros_like(rho)
	rho_minus = numpy.zeros_like(rho)
	f = numpy.zeros_like(rho)
	
	for t in range(1,nt):
		rho_plus[:-1] = rho[1:]
		rho_minus = rho.copy()
		f = 0.5 * (flux(rho_minus, rho_max, V_max) + \
			flux(rho_plus, rho_max, V_max) - dx/dt*(rho_plus-rho_minus))
		rho_n[t,1:-1] = rho[1:-1] + dt/dx*(f[:-2]-f[1:-1])
		rho_n[t,0] = rho[0]
		rho_n[t,-1] = rho[-1]
		rho = rho_n[t].copy()

	return rho_n
	
# sigma = 1.0
# dt = sigma*dx/V_max

# rho = rho_red_light(nx, rho_max,rho_in)
# rho_n = godunov(rho, nt, dt, dx, rho_max, V_max)

# fig = pyplot.figure();
# ax = pyplot.axes(xlim=(0,4),ylim=(4.5,11),xlabel=('Distance'),ylabel=('Traffic density'));
# line, = ax.plot([],[],color='#003366', lw=2); 

# anim = animation.FuncAnimation(fig, animate, frames=rho_n, interval=50)

# pyplot.show()

def minmod(e,dx):
	"""
	Compute the minmod approximation to the slope

	Parameters
	----------
	e 	: array of float
		  input data
	dx 	: float
		  space step

	Returns
	-------
	sigma : array of float
			minmod slope
	"""

	sigma = numpy.zeros_like(e)
	de_minus = numpy.ones_like(e)
	de_plus = numpy.ones_like(e)

	de_minus[1:] = (e[1:]-e[:-1])/dx
	de_plus[:-1] = (e[1:]-e[:-1])/dx

	for i in range(1, len(e)-1):
		if (de_minus[i] * de_plus[i] < 0.):
			sigma[i] = 0.
		elif (numpy.abs(de_minus[i]) < numpy.abs(de_plus[i])):
			sigma[i] = de_minus[i]
		else:
			sigma[i] = de_plus[i]
	return sigma

def muscl(rho, nt, dt, dx, rho_max, V_max):
	""" 
	Computes the solution with MUSCL scheme using the Lax-Friedrichs flux,
	RK2 in time and minmod slope limiting.

	Parameters
	----------
	rho 	: array of floats
			  Density at current time-loop
	nt 		: int
			  Number of time steps
	dt 		: float
			  Time-step size
	dx 		: float
			  Mesh spacing
	rho_max : float
			  Maximum allowed car density
	V_max   : float
			  Speed limit

	Returns
	-------
	rho_n 	: array of floats
			  Density after nt time steps in every point x

	"""
	rho_n = numpy.zeros((nt, len(rho)))
	rho_n[0,:] = rho.copy()
	
	rho_plus = numpy.zeros_like(rho)
	rho_minus = numpy.zeros_like(rho)
	f = numpy.zeros_like(rho)
	rho_star = numpy.zeros_like(rho)

	for t in range(1,nt):
		sigma = minmod(rho,dx)

		rho_left = rho + sigma*dx/2
		rho_right = rho - sigma*dx/2

		f_left = flux(rho_left, rho_max, V_max)
		f_right = flux(rho_right, rho_max, V_max)
		f[:-1] = 0.5*(f_left[:-1]+f_right[1:]-dx/dt*(rho_right[1:]-rho_left[:-1]))

		rho_star[1:-1] = rho[1:-1]+dt/dx*(f[:-2]-f[1:-1])
		rho_star[0] = rho[0]
		rho_star[-1] = rho[-1]

		sigma = minmod(rho_star,dx)

		rho_left = rho_star + sigma*dx/2
		rho_right = rho_star - sigma*dx/2
		f_left = flux(rho_left, rho_max, V_max)
		f_right = flux(rho_right, rho_max, V_max)
		f[:-1] = 0.5*(f_left[:-1]+f_right[1:]-dx/dt*(rho_right[1:]-rho_left[:-1]))

		rho_n[t,1:-1] = .5*rho[1:-1]+.5*(rho_star[1:-1]+dt/dx*\
			(f[:-2]-f[1:-1]))
		rho_n[t,0] = rho[0]
		rho_n[t,-1] = rho[-1]
		rho = rho_n[t].copy()

	return rho_n

sigma = 1.0
dt = sigma*dx/V_max

rho = rho_red_light(nx, rho_max,rho_in)
rho_n1 = muscl(rho, nt, dt, dx, rho_max, V_max)
rho_n2 = godunov(rho, nt, dt, dx, rho_max, V_max)

fig = pyplot.figure();
ax = pyplot.axes(xlim=(0,4),ylim=(4.5,11),xlabel=('Distance'),ylabel=('Traffic density'));
line, = ax.plot([],[],color='b', lw=2); 

anim1 = animation.FuncAnimation(fig, animate, frames=rho_n1, interval=50)
# anim2 = animation.FuncAnimation(fig, animate, frames=rho_n2, interval=50)
pyplot.show()









