#Install package
import sys
!{sys.executable} -m pip install numpy
!{sys.executable} -m pip install matplotlib

# import packages

import numpy as np
import scipy
from scipy import optimize
from scipy.optimize import minimize

T= 24; # Time
x_max= 5; # battery maximum storage 
P_a= np.array([2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 2, 2, 2, 2, 2, 2, 2, 4, 4, 4, 4, 3, 3, 2]); # unitary buying price of energy at time s
P_v= np.array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 1, 1]); # unitary selling price of energy at time s

d =np.array([-1, -1, -1, -1, -1, -2, 2 ,6, 8, 4, 2,  0,  0, -1, -2, -1,  0,  1, 3, 7, 9, 3, 2,  0]); # net demand of energy (already known)
#d =np.array([-2, -2, -2, -2, -2, -3, 1 ,4 ,6 ,3 ,1, -1, -1, -2, -3, -2, -1,  0, 2 ,5 ,5, 2, 1, -1]);
#d =np.array([-3 ,-3 ,-3 ,-3 ,-3 ,-4 ,0 ,2 ,4, 2, 0 ,-2 ,-2 ,-3 ,-4 ,-3 ,-2 ,-1 ,1 ,3 ,5 ,1 ,0 ,-2]);

d=d.reshape((T,1))
P_a=P_a.reshape(T,1)
P_v=P_v.reshape(T,1)

# with T=2 (for testing purposes ...)
# T = 2;
#x_max= 5;

#P_a=np.array([-1, -1])
#P_v=P_a
#d =np.array([-1, -1])
#d=d.reshape(T,1)
#P_a=P_a.reshape(T,1)
#P_v=P_v.reshape(T,1)

def solve_deterministic_pb():
    global d, P_a,P_v
    """ print x,a,v and J(x,a,v) """ 
    
    """ 
    #######  Inequality constraints: #######
    
    xi<=xmax for i = 0.. T: concatenate A1=Identity matrix & A2 & A2
    ---------------
    -x<=0 (size T+1 ! )
    -a<=0 (size T)
    -v<=0 (size T) : identity matrix A3 (size 3T+1)
    ---------------
    
    #A X <= B 
    Indications: with A of shape((len(x)+1 + len(a)+1 + len(v)+1 , len(x) + len (a) + len(v))) #e.g. (10, 7) for T=2
    with B= [x_max*np.ones(T+1,1)],
             np.zeros((3*T+1,1))];
    """

    # xi<=xmax for i = 0.. T: concatenate A1=Identity matrix & A2 & A2
    A1=...
    A2=...
    A=np.concatenate((A1,A2,A2), axis=1)

    # -x, -a, -v <= 0,0,0
    A3=...
    A=...

    #print(np.shape(A))

    # xi<=xmax for i = 0.. T
    B1=...
    # -x, -a, -v <= 0,0,0
    B3=...
    
    B=np.concatenate((B1.transpose(),B3), axis=0);
    #print(B)

    """ 
    ####### Equality constraints: #######
    ---------------
    Aeq1:
    x(s+1)-x(s)-a(s)+v(s)=-d(s) for all s=0...T-1
    -1.  1. on the diagonal (-x(s)+x(s+1))
    -1 for a(s)
    1 for v(s) 
    
    ---------------
    Aeq2
     x(0)=0 
   
    ---------------
    
    # Aeq X = Beq 
    with Beq= [-d,0] # -d for Aeq1 and 0 for Aeq2
    
    """
    # Aeq1:
    M1=... #-x(s) and x(s+1) s=0, T-1
    M2=... #
    M=np.concatenate((-M1,M2), axis=1)+np.concatenate((M2,M1),axis=1)
    
    Aeq1=... # M & -a(s) & v(s) s=0, T-1
    Aeq2=... # #x(0)=0, (with concatenate, don't use axis=1 for only one line, use double parentheses)
    #print(Aeq2)
    Aeq=np.concatenate((Aeq1,Aeq2.transpose()),axis=0)
    #print(Aeq)
   
    Beq=np.concatenate((-d,np.zeros((1,1)))) 
    #print(Beq)
    """ 
    solving min c^T * X s.t. A X <= B and Aeq X= Beq
    no x in min: 0; 
    a P_a -v P_v
    """ 
    # min c^T [x,a,v]
    c=...#cost function: sum_s P_a a(s) - P_v(v)
    #print(c)  #columns of A must be equal to size of c
    X=optimize.linprog(...)
   
    print("x:", X.x[:T+1])
    print(np.shape(X.x[:T+1]))
    print("a:", X.x[T+1:2*T+1])
    print(np.shape(X.x[T+1:2*T+1]))
    print("v:", X.x[2*T+1:])
    print(np.shape(X.x[2*T+1:]))
    print("Optimal cost", X.fun)

solve_deterministic_pb()

from scipy.optimize import minimize

def interpolate_2(y,val):

    def objective(alpha):
        return ...

    # initial guess
    n = 3
    alpha0 = np.zeros(n)

    # show initial objective
    #print('Initial SSE Objective: ' + str(objective(alpha0)))

    # optimize
 
    solution = minimize(...) 
    #print('Output SSE Objective: ' + str((solution.fun)))
    #print('Optimal minimizers: ' + str((solution.x)))
    return solution.x

#% Exercise 2: with y=[0 1 2] and z=[1 3 7], the solution is alpha= [1 1 1]
y=np.array([0,1,2])
val=np.array([1,3,7])
interpolate_2(y,val)

def DP_solve(t,y,alpha):
    def cost(X):
        # X[0]=y X[1]=a, X[2]= v 
        # inf P_a*a -Pv*v + V(t+1,z) (V(t+1, z) en fonction de plusieurs valeurs de z
        return ...

    # x,a,v>=0, x<=xmax -> bounds
    # z -a +v = y -d[t] -> contraints
    
    #print(d[t])
   
    cons = ({'type': 'eq', 'fun': lambda x: ...})
    bnds = (...)
    res = minimize(...);
    X=res.x
    #print("objective solution: ",res.fun)
    #print("min solution: ",res.x)
    z= X[0];
    a= X[1];
    v= X[2];
    
    return z,a,v,res.fun
    

t=0, 
y=0, 
alpha=np.array([ 1, -2,  1])
DP_solve(t,y,alpha)
# the solution is z= 0.5, a= 0, v= 0.5, the optimal value is -0.25.

J=10
def DP_backward():
    
    global T;
    global J;
    global x_max;
    
    alpha= np.zeros((3,T+1));#V(T+1,z)=0
    y_grid= np.zeros(J)
    
    for j in range(J):
        y_grid[j]=... #yj={0,...,x_max}, j=0...J-1
        
    for t in reversed(range(T)): # Backward in time 
    
        Val_update=np.zeros(J) #V(t+1,z)
        for j in range(J): # for all value y_j, we evaluate V(t,y_j) (first with alpha=0 )
            _,_,_,val= DP_solve(...)
            
            Val_update[j]= val;

        #print(y_grid,"z",z)
        alpha[:,t]= interpolate_2(...); # (then with alpha = interpolation of V(t,y_j) and so on...)

    return alpha

DP_backward()

alpha= DP_backward(); # for any z, interpolation of V(t+1,z)
def DP_forward():
    
    global T;
    global P_a;
    global P_v;

    x= np.zeros(T+1); #initialization
    a= np.zeros(T);
    v= np.zeros(T);
    for t in range(0,T):
        z,a0,v0,_= ... # solve
        x[t+1]= ...;# update
        a[t]= ... ;
        v[t]= ... ;

    val= ... #cost
    return x, a, v, val

x,a,v,val=DP_forward()
print("X",x,"a",a,"v",v)
print("val",val)