using JuMP, Clp, StochDynamicProgramming, PyPlot

# Number of timesteps (as we manage the dams over a year, it is equal to the number of weeks):
TF = 52

# Capacity of dams: 
VOLUME_MAX = 100
VOLUME_MIN = 0

# Specify the maximum flow that could be turnined: 
CONTROL_MAX = round(Int, .4/7. * VOLUME_MAX) + 1
CONTROL_MIN = 0

# Some statistics about aleas (water inflow):
W_MAX = round(Int, .5/7. * VOLUME_MAX)
W_MIN = 0
DW = 1

T0 = 1

# Define aleas' space:
N_ALEAS = Int(round(Int, (W_MAX - W_MIN) / DW + 1))
ALEAS = linspace(W_MIN, W_MAX, N_ALEAS);


COST = -66*2.7*(1 + .5*(rand(TF) - .5));

plot(COST, color="k", lw=2)
xlim(0, 53)
xlabel("Week")
ylabel("Cost")

# Define dynamic of the dam:
function dynamic(t, x, u, w)
    return [x[1] - u[1] - u[3] + w[1], x[2] - u[2] - u[4] + u[1] + u[3]]
end


# Define cost corresponding to each timestep:
function cost_t(t, x, u, w)
    return COST[t] * (u[1] + u[2])
end

"""Build aleas probabilities for each month."""
function build_aleas()
    aleas = zeros(N_ALEAS, TF)

    # take into account seasonality effects:
    unorm_prob = linspace(1, N_ALEAS, N_ALEAS)
    proba1 = unorm_prob / sum(unorm_prob)
    proba2 = proba1[N_ALEAS:-1:1]

    for t in 1:TF
        aleas[:, t] = (1 - sin(pi*t/TF)) * proba1 + sin(pi*t/TF) * proba2
    end
    return aleas
end


"""Build an admissible scenario for water inflow."""
function build_scenarios(n_scenarios::Int64, probabilities)
    scenarios = zeros(n_scenarios, TF)

    for scen in 1:n_scenarios
        for t in 1:TF
            Pcum = cumsum(probabilities[:, t])

            n_random = rand()
            prob = findfirst(x -> x > n_random, Pcum)
            scenarios[scen, t] = prob
        end
    end
    return scenarios
end

scenario = build_scenarios(1, build_aleas());

plot(scenario', lw=2, color="blue")
xlabel("Week")
ylabel("Water inflow")
xlim(0, 53)

"""Build probability distribution at each timestep.

Return a Vector{NoiseLaw}"""
function generate_probability_laws(n_scenarios)
    aleas = build_scenarios(n_scenarios, build_aleas())

    laws = Vector{NoiseLaw}(TF-1)

    # uniform probabilities:
    proba = 1/n_scenarios*ones(n_scenarios)

    for t=1:TF-1
        laws[t] = NoiseLaw(aleas[:, t], proba)
    end

    return laws
end

N_SCENARIO = 10
aleas = generate_probability_laws(10);

x_bounds = [(VOLUME_MIN, VOLUME_MAX), (VOLUME_MIN, VOLUME_MAX)];
u_bounds = [(CONTROL_MIN, CONTROL_MAX), (CONTROL_MIN, CONTROL_MAX), (0, 100), (0, 100)];

x0 = [50, 50]

model = LinearSPModel(TF, # number of timestep
                    u_bounds, # control bounds
                    x0, # initial state
                    cost_t, # cost function
                    dynamic, # dynamic function 
                    aleas);


set_state_bounds(model, x_bounds)

# LP solver: 
solver = ClpSolver()
# Set tolerance at 0.1% for convergence: 
EPSILON = 0.001
# Maximum iterations: 
MAX_ITER = 20


params = SDDPparameters(solver,
                        passnumber=10,
                        gap=EPSILON,
                        compute_ub=5, 
                        montecarlo_in_iter=100, 
                        max_iterations=MAX_ITER);

# this function build a SDDPInterface object, and run SDDP onto it. Then, it renders `sddp`
sddp = @time solve_SDDP(model, params, 1);

# We plot evolution of lower-bound: 

plot(sddp.stats.upper_bounds, lw=2, c="darkred", lw=2)

plot(sddp.stats.lower_bounds, lw=2, c="darkblue", lw=2)
grid()
xlabel("Number of iterations")
ylabel("Lower bound");

# and evolution of execution time: 
plot(sddp.stats.exectime, lw=2, c="k")
grid()

xlabel("Number of iterations")
ylabel("Time (s)");

alea_year = Array([7.0 7.0 8.0 3.0 1.0 1.0 3.0 4.0 3.0 2.0 6.0 5.0 2.0 6.0 4.0 7.0 3.0 4.0 1.0 1.0 6.0 2.0 2.0 8.0 3.0 7.0 3.0 1.0 4.0 2.0 4.0 1.0 3.0 2.0 8.0 1.0 5.0 5.0 2.0 1.0 6.0 7.0 5.0 1.0 7.0 7.0 7.0 4.0 3.0 2.0 8.0 7.0])


aleas = zeros(52, 1, 1)
aleas[:, 1, 1] = alea_year;


costs, stocks = StochDynamicProgramming.simulate(sddp, 100);

SDDP_COST = costs[1]
println("SDDP cost: ", SDDP_COST)

plot(stocks[:, :, 1])
plot(stocks[:, :, 2])
xlabel("Week")
ylabel("Stocks")
grid()
ylim(0, 100)
xlim(0, 53);


m = Model(solver=solver)


@variable(m,  VOLUME_MIN  <= x1[1:(TF+1)]  <= VOLUME_MAX)
@variable(m,  VOLUME_MIN  <= x2[1:(TF+1)]  <= VOLUME_MAX)
@variable(m,  CONTROL_MIN <= u1[1:TF]  <= CONTROL_MAX)
@variable(m,  CONTROL_MIN <= u2[1:TF]  <= CONTROL_MAX)
@variable(m, u3[1:TF] >= 0)
@variable(m, u4[1:TF] >= 0)

@objective(m, Min, sum(COST[i]*(u1[i] + u2[i]) for i = 1:TF))

for i in 1:TF
    @constraint(m, x1[i+1] - x1[i] + u1[i] + u3[i] - alea_year[i] == 0)
    @constraint(m, x2[i+1] - x2[i] + u2[i] + u4[i] - u1[i] - u3[i] == 0)
end

@constraint(m, x1[1] == x0[1])
@constraint(m, x2[1] == x0[2])

status = solve(m)
println(status)

LP_COST = getobjectivevalue(m)
println("LP value: ", LP_COST)

plot(getvalue(x1))
plot(getvalue(x2))
xlabel("Week")
ylabel("Stocks")
grid()
ylim(0, 100)
xlim(0, 53);

abs((LP_COST - SDDP_COST)/LP_COST)