# --------------------------------------------------------------------------
# Source file provided under Apache License, Version 2.0, January 2004,
# http://www.apache.org/licenses/
# (c) Copyright IBM Corp. 2015, 2018
# --------------------------------------------------------------------------

"""The model aims at minimizing the production cost for a number of products
while satisfying customer demand. Each product can be produced either inside
the company or outside, at a higher cost.
The inside production is constrained by the company's resources, while outside
production is considered unlimited.
The model first declares the products and the resources.
The data consists of the description of the products (the demand, the inside
and outside costs, and the resource consumption) and the capacity of the
various resources.
The variables for this problem are the inside and outside production for each
product.
"""

from docplex.mp.model import Model
from docplex.util.environment import get_environment
import numpy as np
import Simulation as sim
import torch
from scipy.stats import norm


# ----------------------------------------------------------------------------
# Initialize the problem data
# ----------------------------------------------------------------------------



# ----------------------------------------------------------------------------
# Build the model
# ----------------------------------------------------------------------------

def propogateStates(A,B,t,numProp=2):
    prev_uk_coeff = []
    #prev_uk_coeff.append((np.zeros((11, 4))))
    uk_coeff = []
    prev_Xk_coeff = np.identity(11)


    for i in range(numProp):
        xk_coeff = np.dot(A, prev_Xk_coeff) * t + prev_Xk_coeff
        prev_Xk_coeff = xk_coeff
        for j in range(len(prev_uk_coeff)):
            uk_coeff[j]=np.dot(A, prev_uk_coeff[j]) * t + prev_uk_coeff[j]

        uk_coeff.append(B*t)
        prev_uk_coeff = uk_coeff
    return prev_Xk_coeff,prev_uk_coeff

def solve_br(AAverage,BAverage,SigmaXk,SigmaU,Z=np.zeros((50)),xk=np.array([4.84,0,0,0,0,0,0,0,0,100,1001]),u0=np.array([0, 7269, 0, 0])):
    """ Takes as input:
        - a list of product tuples (name, demand, inside, outside)
        - a list of resource tuples (name, capacity)
        - a list of consumption tuples (product_name, resource_named, consumed)
    """

    simulator=sim.Simulation()
    t=simulator.t
    r=np.array([10,10,100,50,50,50,50,10,10,300]).reshape(10,-1)*.25
    xref=np.array([5,0,0,0,0,0,0,0,0,1000]).reshape(10,-1)
    xk=xk.reshape(11,-1)
    C=np.identity(11)
    C=np.delete(C,9,0)

    Probs = np.array([0.253, 0.5244, 0.8416, 1.644, 2.63, 5.6,20])
    #Probs = np.array([1.644, 2.63, 5.6,7,10,15,20])

    stateSize=10
    actionSize=4
    numProbs=Probs.shape[0]

    NUMSTEPS=3

    set_U_allStep = range(0, actionSize * NUMSTEPS)
    set_U_twoStep=range(0,actionSize*2)
    set_State= range(0, stateSize)
    set_P = range(0, numProbs)

    MVecX=np.array([1000,1000,1000,1000,1000,1000,1000,1000,1000,10000])
    MVecU = np.array([1000, 1000000, 1000, 1000])

    uUpper = np.array([45, 10000, 45, 45])
    uUpper=np.tile(uUpper,NUMSTEPS)
    uLower = np.array([-45, 0, -45,-45])
    uLower=np.tile(uLower,NUMSTEPS)



    # --- decision variables ---
    mdl = Model(name='BoundedRat')

    u_vars ={(i): mdl.continuous_var(lb=uLower[i], ub=uUpper[i],name="u_{0}".format(i))
     for i in set_U_allStep}

    p2_vars = {(i): mdl.binary_var(name="p2_{0}".format(i))
              for i in set_P}
    p3_vars = {(i): mdl.binary_var(name="p3_{0}".format(i))
              for i in set_P}
    """p4_vars = {(i): mdl.binary_var(name="p4_{0}".format(i))
              for i in set_P}
              """

    allP=[p2_vars,p3_vars]

    """uP1_vars = {(i): mdl.continuous_var(lb=0,name="uP1_{0}".format(i))
              for i in set_U_allStep}
    uP2_vars = {(i): mdl.continuous_var(lb=0, name="uP2_{0}".format(i))
              for i in set_U_allStep}
              """

    z1_vars = {(i): mdl.continuous_var(lb=0, name="z1_{0}".format(i))
               for i in set_State}

    z2_vars = {(i): mdl.continuous_var(lb=0, name="z2_{0}".format(i))
               for i in set_State}


    for N in range (2,NUMSTEPS+1):
        #r=r1*N
        p_vars=allP[N-2]
        SigmaU_step=np.tile(SigmaU,N)
        #newSigmaU=prevSigmaU+SigmaU
        #prevSigmaU=newSigmaU
        #SigmaU_step=np.hstack((SigmaU_step,newSigmaU))
        xkC_Step, ukC_Step = propogateStates(AAverage, BAverage, t, N)

        #UCoefficients
        coef1_Step=np.hstack((np.dot(C, ukC_Step[i])for i in range(0,N)))

        #XK Coefficients
        b1_Step=(r+xref-np.dot(C,np.dot(xkC_Step,xk))).reshape(-1)
        b2_Step=(r-xref+np.dot(C,np.dot(xkC_Step,xk))).reshape(-1)


        set_U_Step = range(0, actionSize * N)

        # --- constraints ---
        # demand satisfaction

        #Positive absolute value of probability, positive absolute value of U


        constraintsCC_PP_UP = {(s,p):
        mdl.add_constraint(
                ct=-p_vars[p]*MVecX[s]+mdl.sum((coef1_Step[s,u] + Probs[p] * SigmaU_step[s,u])*u_vars[u]  for u in set_U_Step)<= (b1_Step[s]-Probs[p]*np.dot(SigmaXk,abs(xk))[s]).item() ,
                ctname="chanceConstraintCC_PP_UP_{0}_{1}".format(s,p))
            for s in set_State for p in set_P}
        # Negative absolute value of probability, positive absolute value of U
        constraintsCC_PN_UP = {(s,p):
        mdl.add_constraint(
                ct=-p_vars[p]*MVecX[s]+mdl.sum((-1*coef1_Step[s,u] - Probs[p] * SigmaU_step[s,u])*u_vars[u]  for u in set_U_Step) <= (b2_Step[s]+Probs[p]*np.dot(SigmaXk,abs(xk))[s]).item() ,
                ctname="chanceConstraintPlusCC_PN_UP_{0}_{1}".format(s,p))
            for s in set_State for p in set_P}

        #Positive absolute value of probability, negative absolute value of U
        constraintsCC_PP_UN = {(s,p):
        mdl.add_constraint(
                ct=-p_vars[p]*MVecX[s]+mdl.sum((coef1_Step[s,u] - Probs[p] * SigmaU_step[s,u])*u_vars[u]  for u in set_U_Step)<= (b1_Step[s]-Probs[p]*np.dot(SigmaXk,abs(xk))[s]).item() ,
                ctname="chanceConstraintPlusCC_PP_UN_{0}_{1}".format(s,p))
            for s in set_State for p in set_P}
        # Negative absolute value of probability, negative absolute value of U
        constraintsCC_PN_UN= {(s,p):
        mdl.add_constraint(
                ct=-p_vars[p]*MVecX[s]+ mdl.sum((-1*coef1_Step[s,u] + Probs[p] * SigmaU_step[s,u])*u_vars[u]  for u in set_U_Step) <= (b2_Step[s]+Probs[p]*np.dot(SigmaXk,abs(xk))[s]).item() ,
                ctname="chanceConstraintPlusCC_PN_UN_{0}_{1}".format(s,p))
            for s in set_State for p in set_P}

        if N == 3:
            Ax = np.dot(C, np.dot(xkC_Step, xk))
            constraintsEQ_P = {
                mdl.add_constraint(
                    ct=mdl.sum((coef1_Step[s, u]) * u_vars[u] for u in set_U_Step) + Ax[s].item() - xref[s].item() -
                       z1_vars[s] + z2_vars[s] == 0,
                    ctname="ZDecisionVariables_{0}") for s in set_State}


    constraintsEQ_P2 = {
    mdl.add_constraint(
        ct=mdl.sum(p2_vars[p] for p in set_P) == numProbs-1,
        ctname="ProbDecisionVariables2_{0}_{1}")}

    constraintsEQ_P3 = {
        mdl.add_constraint(
            ct=mdl.sum(p3_vars[p] for p in set_P) == numProbs - 1,
            ctname="ProbDecisionVariables3_{0}_{1}")}

    """constraintsEQ_P4 = {
        mdl.add_constraint(
            ct=mdl.sum(p4_vars[p] for p in set_P) == numProbs - 1,
            ctname="ProbDecisionVariables4_{0}_{1}")}
            """



    constraintsEQ_P = {
        mdl.add_constraint(
            ct=(u_vars[u+m*4]-u_vars[u+m*4+4])<=5,
            ctname="ChangeU_{0}_{1}") for u in set([0,2,3]) for m in range(0,NUMSTEPS-1)}

    constraintsEQ_P = {
        mdl.add_constraint(
            ct=(-u_vars[u+m*4]+u_vars[u+m*4+4])<=5,
            ctname="CHANGEU_{0}_{1}")for u in set([0,2,3]) for m in range(0,NUMSTEPS-1)}

    constraintsEQ_P = {
        mdl.add_constraint(
            ct=(u_vars[u]-u0[u])<=5,
            ctname="ChangeU1_{0}")for u in set([0,2,3])}

    constraintsEQ_P = {
        mdl.add_constraint(
            ct=(-u_vars[u]+u0[u])<=5,
            ctname="CHANGEU1_{0}_{1}")for u in set([0,2,3])}


    constraintsEQ_P = {
        mdl.add_constraint(
            ct=(u_vars[1+m*4]-u_vars[1+m*4+4])<=100,
            ctname="ChangeUThrottle_{0}_{1}") for m in range(0,NUMSTEPS-1)}

    constraintsEQ_P = {
        mdl.add_constraint(
            ct=(-u_vars[1+m*4]+u_vars[1+m*4+4])<=100,
            ctname="CHANGEUThrottle_{0}_{1}") for m in range(0,NUMSTEPS-1)}

    constraintsEQ_P = {
        mdl.add_constraint(
            ct=(u_vars[1]-u0[1])<=100,
            ctname="ChangeU1Throttle_{0}")}

    constraintsEQ_P = {
        mdl.add_constraint(
            ct=(-u_vars[1]+u0[1])<=100,
            ctname="CHANGEU1Throttle_{0}_{1}")}


    """constraintsEQ_P = {
            mdl.add_constraint(
                ct=u_vars[u]-u0[u]-uP1_vars[u] +uP2_vars[u]==0,
                ctname="ZDecisionVariables_{0}") for u in set([0,1,2,3])}

    constraintsEQ_P = {
        mdl.add_constraint(
            ct=u_vars[u+m*4] - u_vars[u+m*4+4] - uP1_vars[u+m*4+4] + uP2_vars[u+m*4+4] == 0,
            ctname="ZDecisionVariables_{0}") for u in set([0,1,2,3]) for m in range(0,NUMSTEPS-1)}
            """



    ###################RELU

################DIVERSITY
    d1_vars = {(i): mdl.continuous_var(lb=-100,ub=100,name="d1_{0}".format(i))
               for i in set([0,1,2,3])}

    d2_vars = {(i): mdl.continuous_var(lb=-100, ub=100, name="d2_{0}".format(i))
               for i in set([0,1,2,3])}

    k1 = {(i): mdl.binary_var(name="k1_{0}".format(i))
          for i in set([0,1,2,3])}

    k2 = {(i): mdl.binary_var(name="k2_{0}".format(i))
          for i in set([0, 1, 2, 3])}

    M2 = 1000000
    constraintsEQ_P = {
        mdl.add_indicator(
            k1[u],d1_vars[u] <= u_vars[u] - u0[u],
            name="ZDecisionVariables_{0}") for u in set([0, 1, 2, 3])}
    constraintsEQ_P = {
        mdl.add_indicator(
            k1[u],d1_vars[u] <= u0[u]-u_vars[u],active_value=0,
            name="ZDecisionVariables_{0}") for u in set([0, 1, 2, 3])}

    constraintsEQ_P = {
        mdl.add_indicator(
            k2[u],d2_vars[u] <= u_vars[u] - u_vars[u+4],
            name="ZDecisionVariables_{0}") for u in set([0, 1, 2, 3])}
    constraintsEQ_P = {
        mdl.add_indicator(
            k2[u],d2_vars[u] <= u_vars[u+4]-u_vars[u],active_value=0,
            name="ZDecisionVariables_{0}") for u in set([0, 1, 2, 3])}

    prime_weight=np.array([.0001,.00001,.0001,.00001])
    prime_weight=np.tile(prime_weight,NUMSTEPS)
    # --- objective ---
    #
    divWeight=np.array([.2,.01,.2,.2])
    divWeight=np.tile(divWeight,NUMSTEPS)
    maxProb_Objective=mdl.sum((1-p2_vars[p2])*Probs[p2] for p2 in set_P)+mdl.sum((1-p3_vars[p3])*Probs[p3] for p3 in set_P)
    #minChangeInput_Objective=-mdl.sum((uP1_vars[z]+uP2_vars[z])*prime_weight[z] for z in set_U_allStep)
    minDistance = -mdl.sum((z1_vars[z] + z2_vars[z]) for z in set_State)
    maxDiversityObj = mdl.sum((d1_vars[z])*divWeight[z] for z in set([0, 1, 2, 3]))+mdl.sum((d2_vars[z])*divWeight[z] for z in set([0, 1, 2, 3]))
    #-.0000001*maxInfoGain_Objective
    objective =100* maxProb_Objective+2*maxDiversityObj+minDistance
    mdl.maximize(objective)
    sol=mdl.solve()
    #mdl.print_solution()
    #print(mdl.get_solve_status())

    solution=[sol[u_vars[u]] for u in set_U_allStep]
    allU=[sol[u_vars[u]] for u in set_U_allStep]

    prob1 = np.sum(np.array([(1 - sol[p2_vars[p2]]) * Probs[p2] for p2 in set_P]))
    prob1 = norm.cdf(prob1)
    prob2 = np.sum(np.array([(1 - sol[p3_vars[p3]]) * Probs[p3] for p3 in set_P]))
    prob2 = norm.cdf(prob2)

    prob = np.mean(np.array([prob1, prob2]))

    #xkC_Step, ukC_Step = propogateStates(AAverage, BAverage, t, 3)
    #coef1_Step = np.hstack((np.dot(C, ukC_Step[i]) for i in range(0, 3)))
    #print('here',np.dot(coef1_Step , np.asarray(allU)).reshape(10,1)+np.dot(C,np.dot(xkC_Step,xk)).reshape(10,1))

    #print("STEP")
    #print(u0)
    #print("u",solution)
    #print('infogain',[.01*fC3[i].item() * sol[o2[i]] for i in set_Neurons ])
    #print('minchange',[(sol[uP1_vars[z]]+sol[uP2_vars[z]])*prime_weight[z] for z in set_U_allStep])
    #print("prob",[(1-sol[p2_vars[p2]])*Probs[p2] for p2 in set_P]+[(1-sol[p3_vars[p3]])*Probs[p3] for p3 in set_P])
    #print("Z",[sol[z1_vars[z]] + sol[z2_vars[z]] for z in set_State])
    #print("zVec",zVec)
    #print('Must be greater than',[-(1-sol[e_vars[u]])*M+eps1 for u in set_Neurons])
    #print('Must be less than', [( sol[e_vars[u]]) * M + eps1 for u in set_Neurons])
    #print('e', [sol[e_vars[u]] for u in set_Neurons])
    #print('R', [sol[R_vars[u]] for u in set_Neurons])
    #print('o1', [sol[o1[u]] for u in set_Neurons])
    #print('o2', [sol[o2[u]] for u in set_Neurons])


    return solution


def print_production_solution(mdl, products):
    obj = mdl.objective_value
    print("* Production model solved with objective: {:g}".format(obj))
    print("* Total inside cost=%g" % mdl.total_inside_cost.solution_value)
    for p in products:
        print("Inside production of {product}: {ins_var}".format
              (product=p[0], ins_var=mdl.inside_vars[p].solution_value))
    print("* Total outside cost=%g" % mdl.total_outside_cost.solution_value)
    for p in products:
        print("Outside production of {product}: {out_var}".format
              (product=p[0], out_var=mdl.outside_vars[p].solution_value))


# ----------------------------------------------------------------------------
# Solve the model and display the result
# ----------------------------------------------------------------------------


if __name__ == '__main__':
    # Build the model
    model = solve_br()
    model.print_information()
