#!/usr/bin/env python3
""" Attempt to compute a lower bound using fixed cost wbot for the bidding sequence
see Typst.app

This LP is unbounded. There is a serious problem with it.
"""

import sys, math, pulp, scipy 

if (len(sys.argv) != 2):
    prog = sys.argv[0]
    print(f"Usage: {prog} [R]" )
    print(f"or:    {prog} [low:high:step]")
    print("Will print a row per robustness guarantee R in range with R and lower bound on corresponding consistency")
    sys.exit(1)

n = 200
p = n-2  # prediction
bids = [math.exp(i / n) for i in range(n)]


def consistency(R, prefixSumTrick=True):
    pb = pulp.LpProblem("consistency", pulp.LpMaximize)

    wbot = float(-1 / scipy.special.lambertw(-1 / R, -1)) 
    
    y = []
    z = []
    for i in range(n):
        y.append(pulp.LpVariable(f"y{i}", lowBound = 0))
        z.append(pulp.LpVariable(f"z{i}", lowBound = 0))
    
    if prefixSumTrick:
        # improve LP through prefix sums
        yy = []     # prefix sum
        zz = []     # prefix sum 
        zb = []     # weighted prefix sum 
        for i in range(n):
            yy.append(pulp.LpVariable(f"yy{i}", lowBound = 0))
            zz.append(pulp.LpVariable(f"zz{i}", lowBound = 0))
            zb.append(pulp.LpVariable(f"zR{i}", lowBound = 0))
        
        # we use the fact that in Python yy[-1] is the last element 0 (same for zz, zb)
        yy.append(0)
        zz.append(0)
        zb.append(0)


        for i in range(n):
            pb += yy[i] == yy[i-1] + y[i] 
            pb += zz[i] == zz[i-1] + z[i]
            pb += zb[i] == zb[i-1] + z[i] * 1 / bids[i]

        pb += yy[n - 1] + wbot * zb[n-1] - R * zz[n-1] + z[p] * R 
        #    max  sum_j y_j + sum_j z_j (wbot / b_j - R) + z_p R  


        # "such that" forall i <= k :& 
        for k in range(n):
            for i in range(k + 1):
                pb += yy[k] - yy[i - 1] - (bids[k] * zb[n - 1] - zb[i]) <= 0
                # sum_(j:i <= j <= k) y_j 
                #    - sum_(j > i) z_j b_k / b_j <= 0   &[x_(i,k)]\
    else:
        # same LP without prefix sums


        pb += pulp.lpSum(y) \
            + pulp.lpSum(z[j] * (wbot / bj - R) for j, bj in enumerate(bids)) \
            + z[p] * R 

        for k in range(n):
            for i in range(k + 1):
                pb += pulp.lpSum(y[j] for j in range(i, k+1)) \
                    - pulp.lpSum(z[j] * bids[k] / bids[j] for j in range(i + 1, n)) <= 0


    pb += z[p] <= 1
    #    &z_p <= 1 &[C] \

    status = pb.solve(pulp.PULP_CBC_CMD(msg=0))
    # print(pulp. LpStatus[status], file=sys.stderr)
    assert (status == pulp.LpStatusOptimal)
    return pb.objective.value()

arg = sys.argv[1]
if ":" in arg:
    low, high, step = map(float, arg.split(":"))
    assert (low < high)
    assert (step > 0)
else:
    R = float(arg)
    low = R
    high = R 
    step = 1

R = low 
while (R <= high):
    print(R, consistency(R, False))
    R += step
