import sys
import pulp
import time
from scipy.special import lambertw
from math import exp

if (len(sys.argv) != 2):
    print("Usage: lb.py [R]")
    sys.exit(1)

R = float(sys.argv[1])
wbot = -1. / lambertw( -1. / R, -1).real
print(f"value of wbot : {wbot}")
start_build = time.time()
n = 300
def b(x):
    return exp(8*x/n)
p =  n-2

pb = pulp.LpProblem("robustness", pulp.LpMaximize)

gamma = pulp.LpVariable("gamma", lowBound = 0)
Y = [pulp.LpVariable(f"M{i}", lowBound=0) for i in range(n+1)]
Z = [pulp.LpVariable(f"P{i}", lowBound=0) for i in range(n+1)]
ZB = [pulp.LpVariable(f"ZB{i}", lowBound=0) for i in range(n+1)]
y = [pulp.LpVariable(f"pi{i}", lowBound=0) for i in range(n+1)]
z = [pulp.LpVariable(f"mu{i}", lowBound=0) for i in range(n+1)]
zp = pulp.LpVariable(f"zp", lowBound=0)

pb += Y[0] == y[0]
for i in range(1, n+1):
    pb += Y[i] == Y[i-1] + y[i]

pb += ZB[0] == z[0] * 1 / b(0)
for i in range(1, n+1):
    pb += ZB[i] == ZB[i-1] + z[i] * 1./ b(i)

pb += Z[0] == z[0]
for i in range(1, n+1):
    pb += Z[i] == Z[i-1] + z[i]

pb += z[0] == 0
pb += y[0] == 0

pb += Y[n] + wbot * ZB[n] - R * Z[n] + zp * wbot / b(p)
pb += zp <= 1 

for i in range(1, n + 1):
    for k in range(i, n + 1):
        sum_y = Y[k] - Y[i-1]
        sum_z = b(k) * (ZB[n] - ZB[i])
        if (i < p) :
            pb += sum_y - sum_z - zp * b(k)/b(p) <= 0
        else :
            pb += sum_y - sum_z <= 0

build_time = time.time() - start_build

# --- Solve ---
start_solve = time.time()

pb.writeLP("tmp.lp")
status = pb.solve()
solve_time = time.time() - start_solve

print("Status:", pulp.LpStatus[status])
print("Objective:", pulp.value(pb.objective))
print("Build time: %.3fs" % build_time)
print("Solve time: %.3fs" % solve_time)