import lpLib as LP
import numpy as np
wheat = 0
corn = 1
sugar = 2
sugarLow = 3

#Parameter  = Value
s_1  = 170  # Sales price, wheat
s_2  = 150  # Sales price, corn
s_3  = 36   # Sales price, sugar beats favourable
s_4  = 10   # Sales price, sugar beats low
c_1  = 150  # Planting cost, wheat/acre
c_2  = 230  # Planting cost, corn/acre
c_3  = 260  # Planting cost, sugar beats/acre
b_1  = 238  # Cost of buying from the wholesaler, wheat
b_2  = 210  # Cost of buying from the wholesaler, corn
r_1  = 2.5  # Normal yield per acre, wheat
r_2  = 3    # Normal yield  per acre, corn
r_3  = 20   # Normal yield  per acre, sugar beets
d_1  = 200  # Demand/requirements for Tom's cattle, corn
d_2  = 240  # Demand/requirements for Tom's cattle, wheat
p_1  = 1/3  # Probability, good weather
p_2  = 1/3  # Probability, normal weather
p_3  = 1/3  # Probability, bad weather
s = [s_1,s_2,s_3,s_4] # Vector of sales prices
c = [c_1, c_2, c_3]   # Vector of planting costs
b = [b_1,b_2]         # Vector of bying from the wholesaler
r = [r_1,r_2,r_3]     # Vector of normal yields
d = [d_1,d_2]         # Vector of demands for the cattle 
p = [p_1, p_2, p_3]   # Probability vector
u = [1.2, 1, 0.8]     # Yield factor
n_x = 3 # x-variable = number of acres allocated for each type of seed
n_w = 4 # w-variable = amount sold
n_y = 2 # y-variable = amount buyed, if required
k = 3   # number of scenarios

varNames = ["x_Wheat","x_Corn","x_Sugar","w_Wheat","w_Corn","w_Sugar","w_SugarLow","y_Wheat","y_Corn"]
yields = ["High","Medium","Low"]
def spSolve(p,x_EV):
    '''
    the spSolve function allows to pass the following parameters:
    p = Probability vector
    x_EV = Expected value solution for x. Optional, if given, specify as constraint
    '''
    objective = []
    for j in range(n_x):
        objective.append(-c[j]) # Planting costs, x-variable
         
    for i in range(k):          # For each scenario:
        for j in range(n_w):
            objective.append(p[i]*s[j])  # Selling value
             
        for j in range(n_y):             
            objective.append(-p[i]*b[j]) # Cost of bying, if needed for the cattle   
            
    LP.coefficients(objective, False)  # Maximization problem
    nVar = len(objective)
     
    # Constraints, use of land:
    A_i = np.zeros([nVar])
    for j in range(n_x):
        A_i[j] = 1
    LP.upperBound(500,A_i)
    wOffset = n_x
    yOffset = wOffset + n_w

    for i in range(k):
        #Sugar beats, scenario $i$:
        #  w_3,i + w_4,i \le r_3 * u_i * x_3\\
        #  w_3,i \le 6000
        A_i = np.zeros([nVar]) # Create new empty row for A-matrix, then fill:
        A_i[wOffset+sugar+i*(n_w + n_y)] = 1
        A_i[wOffset+sugarLow+i*(n_w + n_y)] = 1
        A_i[sugar] = -u[i]*r[sugar]        
        LP.upperBound(0,A_i)
        A_i = np.zeros([nVar])
        A_i[wOffset+sugar+i*(n_w + n_y)] = 1
        LP.upperBound(6000,A_i)
        # Wheat:
        # x_1 r_1 u_i - w_{1,i} + y_{1,i} \ge d_1
        A_i = np.zeros([nVar])
        A_i[wheat] = r[wheat]*u[i]
        A_i[wOffset+wheat+i*(n_w + n_y)] = -1  
        A_i[yOffset+wheat+i*(n_w + n_y)] = 1
        LP.lowerBound(d[wheat],A_i)
        # Corn:
        # x_2 r_2 u_i - w_{2,i} + y_{2,i} \ge d_2
        A_i = np.zeros([nVar])
        A_i[corn] = r[corn]*u[i]
        A_i[wOffset+corn+i*(n_w + n_y)] = -1  
        A_i[yOffset+corn+i*(n_w + n_y)] = 1
        LP.lowerBound(d[corn],A_i)

        if len(x_EV) > 0: # The expected value solution is given as constraint
            for i in range(len(x_EV)):
                A_i = np.zeros([nVar])
                A_i[i] = 1
                LP.equalityConstraint(x_EV[i],A_i)
    ret = LP.lpSolve(True)
    l = 0
    
    for j in range(n_x):
        print(varNames[l],"=",ret.x[l])
        l += 1
    for i in range(k):          
        for j in range(n_w):
            print(varNames[l-i*(n_w+n_y)]+","+yields[i],"=",ret.x[l])
            l += 1      
        for j in range(n_y):
            print(varNames[l-i*(n_w+n_y)]+","+yields[i],"=",ret.x[l])
            l += 1            
              
    return ret

'''
a) and b)
'''
print("\nProblem a) and b):\n")
for i in range(k):
   print("When it is known that the yield will be", yields[i])
   pKnown = np.zeros([k])
   pKnown[i] = 1
   res = spSolve(pKnown,[])
   print("Z =", int(res.fun))
   print("Press any key to continue...")
   input()
  

'''
To find the expected value solution, we use the p-vector = [0,1,0], i.e.,
only the "medium" scenario for the yield is considered, which then
represent the expected value of U
'''
print("\nc) Expected value solution:")
res=spSolve([0,1,0],[])
x_EV=[]
for i in range(n_x):
    x_EV.append(res.x[i])
print("Press any key to continue...")
input()   
print("Result when yield is good, but it was planned for medium yield")
res=spSolve([1,0,0],x_EV)
print("Z =",int(res.fun))
print("Press any key to continue...")
input()   
print("Result when yield is bad, but it was planned for medium yield")
res=spSolve([0,0,1],x_EV)
print("Z =",int(res.fun))
print("Press any key to continue...")
input()   
'''
Now we calculate the stochastic solution. We specify the probabilty
vector p, and leave the constraint vector empty, i.e., no expected value constraint
'''
print("\nd) Stochastic solution")
res = spSolve(p,[])
Z12 = res.fun
print("Z12 =", int(Z12))
print("Press any key to continue...")
input()
 

res=spSolve(p,x_EV)
print("\ne) Z_EEV =",int(res.fun))
print("\nVSS =",int(Z12-res.fun))

print("Press any key to continue...")
input()    
print("\nf) Value of perfect information:")
Z = 0
for i in range(k):
    print("Perfect information, scenario", yields[i] )
    p_PFI = np.zeros([k])
    p_PFI[i] = 1
    res = spSolve(p_PFI,[])
    print("Z =",int(res.fun))
    Z += res.fun/k
    print("Press any key to continue...")
    input()
print("EVPI =",int(Z - Z12))
    
print("g) Too difficult, see course compendium for ideas...")