using JuMP, PyPlot, Gurobi

Tmax = 160                 # max time of simulation (seconds)
dt = 0.25                   # discretization interval (seconds)

T = round(Int, Tmax/dt)    # total number of time steps
vmax = 100000/3600         # max speed, 100 km/h, converted to m/sec
amax = 2                   # max acceleration, 2 m/sec^2
xlight = 2000              # distance of the light, 2 km, converted to m.
xend = 4000                # distance of the finish, 4 km, converted to m.
ty = round(Int,5/dt)       # duration of the yellow light in timesteps
tr = round(Int,20/dt)      # duration of the red light in timesteps

m = Model(solver = GurobiSolver(OutputFlag=0))

# primary equations: assuming all goes right
@variable(m, 0 <= v[1:T] <= vmax)
@variable(m, -amax <= a[1:T] <= amax)
@variable(m, x[1:T] )

@constraint(m, x[1] == 0)
@constraint(m, v[1] == vmax)
@constraint(m, dveqn[t=1:T-1], v[t+1]-v[t] == a[t]*dt )
@constraint(m, dxeqn[t=1:T-1], x[t+1]-x[t] == v[t]*dt + 0.5*a[t]*dt^2 )

# contingency equations: assuming light turns yellow
@variable(m, 0 <= vhat[1:T,1:ty+tr] <= vmax)
@variable(m, -amax <= ahat[1:T,1:ty+tr] <= amax)
@variable(m, xhat[1:T,1:ty+tr])

@constraint(m, cxhat[t=1:T], xhat[t,1] == x[t])
@constraint(m, cvhat[t=1:T], vhat[t,1] == v[t])
@constraint(m, cdveqn[t=1:T,τ=1:ty+tr-1], vhat[t,τ+1]-vhat[t,τ] == ahat[t,τ]*dt )
@constraint(m, cdxeqn[t=1:T,τ=1:ty+tr-1], xhat[t,τ+1]-xhat[t,τ] == vhat[t,τ]*dt + 0.5*ahat[t,τ]*dt^2 )

# running a yellow light
@variable(m, zgo[t=1:T], Bin)
@constraint(m, xhat[:,ty] .>= xlight*zgo )

# stopping for a red light
@variable(m, zstop[t=1:T], Bin)
@constraint(m, xhat[:,ty+tr] .<= xlight + (1-zstop)*xend )

# logic constraint; one of them must happen!
@constraint(m, zgo + zstop .>= 1)

# termination criterion: are we there yet?
@variable(m, zfinish[1:T], Bin)
@constraint(m, x .>= xend*zfinish )

# objective function: reach the end as soon as possible assuming no ill-timed yellow lights.
@objective(m, Max, sum(zfinish) )
;

@time solve(m)

t_val = 0:dt:Tmax; t_val = t_val[1:end-1]
a_val = getvalue(a)
v_val = getvalue(v)*3.6
x_val = getvalue(x)/1000
z_fin = getvalue(zfinish)
z_go = getvalue(zgo)
z_stop = getvalue(zstop)
vhat_val = getvalue(vhat)*3.6
txlight = t_val[findmin(abs(x_val-xlight/1000))[2]]
txend = t_val[findmin(abs(x_val-xend/1000))[2]];

figure(figsize=(8,3))
plot(t_val, v_val, ".")
plot(txlight*[1,1], [0,120], "r--")
plot(txend*[1,1], [0,120], "m--")
xlim([60,80]); ylim([60,110])
grid()
xlabel("Time (sec)"); ylabel("Speed (km/h)")
title("Optimal speed as a function of time")
legend(["optimal speed","traffic light"],loc="lower right")
tight_layout()
#savefig("traffic_light_plot.png")