include("colMatrix.jl")


# Initialize the NMPC framework
function nMPC0(u0List::Vector{Float64},x0List::Vector{Float64},z0List::Vector{Float64},pList::Vector{Float64},NFE) #This only builds the nlp as it lacks the sufficient info to solve
    col = [0, 0.155051,0.644949,1.0] #colpoint 1 is the zeroth colpoint in this formulation

    colMat,colCont,colObj = colMatrix.collocationMatrix(col) #colMat is square but first row will be ignored   

    # setting Timescale and number of steps

    dt = 0.05  #! [hours]
    Nx = 5 # States
    Nu = 2 # Inputs
    Nz = 2 # Algebraic equations
    Np = 2 # Parameters
    NCP = 3 # Collocation points

    # Define initial guesses
    uG = u0List
    xG = x0List
    zG = z0List

    optimizer = optimizer_with_attributes(Ipopt.Optimizer, 
                "linear_solver" => "mumps",
                "print_level" => 5,
                "max_iter" => 1000,
                "tol" => 1e-8,
                "acceptable_tol" => 1e-8, 
                "mu_init" => 10^(-1))
    #optimizer = with_optimizer(Ipopt.Optimizer, linear_solver = "mumps")
    nlp = Model(optimizer)


    #-------Declare variables----------
    ubl = [0.0, -9000.0] # Lower bounds on u. Vindot and Qdot respectively
    ubu = [32.4, 0.0] # Upper bounds on u. Vindot and Qdot respectively


    #Input variables
    @variable(nlp, ubl[nu] <= u[nu=1:Nu,-1:NFE-1] <= ubu[nu]) # Declare input and corresponding bounds

    xbl = [0.0,0.0,0.0,273.0,273.0]
    xbu = [8.0,5.0,5.0,350.0,350.0]

    # State variables
    @variable(nlp, xbl[nx] <= x[nx=1:Nx,1:NFE,0:NCP] <= xbu[nx]) # Declare states and corresponding bounds
    @variable(nlp, xend[nx=1:Nx]) #Declare xend for the orthogonal collocation

    # Algebraic variables
    @variable(nlp, 0.0 <= z[1:Nz,1:NFE,1:NCP]) # Declare algebraic states

    # Slack variables
    @variable(nlp, e1[1,1:NFE,0:NCP]) #! Added more points
    @variable(nlp, e2[1,1:NFE,0:NCP]) #! Added more points

    # Setting initial values
    for nx in 1:Nx
        set_start_value.(x[nx,:,:],xG[nx])
    end

    for nu in 1:Nu
        set_start_value.(u[nu,:],uG[nu])
    end

    for nz in 1:Nz
        set_start_value.(z[nz,:,:],zG[nz])
    end

    # Parameter variables and initial state
    @NLparameter(nlp, par[np=1:Np,1:NFE] == pList[np]) # Declare parameters
    @NLparameter(nlp, x0[nx=1:Nx] == x0List[nx]) #Define initial states

    ## Model equations (Dynamic Model)

    # Collocation points are defined
    @NLconstraints(nlp, begin
        [nfe=1:NFE,ncp=1:NCP], sum(colMat[ncp+1,h+1]*x[1,nfe,h] for h in 0:NCP) - dt*((u[1,nfe-1])) == 0 # Vr
        [nfe=1:NFE,ncp=1:NCP], sum(colMat[ncp+1,h+1]*x[2,nfe,h] for h in 0:NCP) - dt*(-((u[1,nfe-1]*x[2,nfe,ncp])/x[1,nfe,ncp]) - (par[1,nfe]*x[2,nfe,ncp]*x[3,nfe,ncp])) == 0 # Ca
        [nfe=1:NFE,ncp=1:NCP], sum(colMat[ncp+1,h+1]*x[3,nfe,h] for h in 0:NCP) - dt*(((u[1,nfe-1]/(x[1,nfe,ncp]))*(Cbin-x[3,nfe,ncp])) - (par[1,nfe]*x[2,nfe,ncp]*x[3,nfe,ncp])) == 0 # Cb
        [nfe=1:NFE,ncp=1:NCP], sum(colMat[ncp+1,h+1]*x[4,nfe,h] for h in 0:NCP) - dt*(((u[1,nfe-1]/x[1,nfe,ncp])*(Tin-x[4,nfe,ncp])) - ((α*z[2,nfe,ncp]*(x[4,nfe,ncp]-x[5,nfe,ncp]))/(ρ*x[1,nfe,ncp]*Cp)) - ((par[1,nfe]*x[2,nfe,ncp]*x[3,nfe,ncp]*par[2,nfe])/(ρ*Cp))) == 0 # Tr
        [nfe=1:NFE,ncp=1:NCP], sum(colMat[ncp+1,h+1]*x[5,nfe,h] for h in 0:NCP) - dt*(((u[2,nfe-1]+(α*z[2,nfe,ncp]*(x[4,nfe,ncp]-x[5,nfe,ncp])))/(ρ*Vj*Cp))) == 0 # Tj
    end)

    # Model equations (Algebraic Model)
    # The algebraic variables are dependent of the states, so collocation points has to be used 
    @NLconstraints(nlp,begin
        [nfe=1:NFE,ncp=1:NCP], z[1,nfe,ncp] - ((Ca0*Vr0 + Cc0*Vr0 - x[2,nfe,ncp]*x[1,nfe,ncp])/x[1,nfe,ncp]) == 0 
        [nfe=1:NFE,ncp=1:NCP], z[2,nfe,ncp] - (π*(r^2) + ((0.002*x[1,nfe,ncp])/r)) == 0
    end)

    # Define constraints form Equation 25b, 25c, 25d, 25e
    @NLconstraint(nlp, [nfe=1:NFE,ncp=0:NCP], 322.0 <= x[4,nfe,ncp] + e1[1,nfe,ncp] <= 326.0) # 25c
    @NLconstraint(nlp, [nfe=1:NFE,ncp=0:NCP], x[1,nfe,ncp] + e2[1,nfe,ncp] <= 7.0) # 25d
    @NLconstraint(nlp, [nfe=1:NFE,ncp=0:NCP], -1.0 <= e1[1,nfe,ncp] <= 1.0) # 25e
    @NLconstraint(nlp, [nfe=1:NFE,ncp=0:NCP], -0.01 <= e2[1,nfe,ncp] <= 0.01)  # 25f


    # Collocation constraints
    @NLconstraint(nlp, [nx=1:Nx,nfe=2:NFE,ncp=0], x[nx,nfe,ncp] - sum(colCont[h+1]*x[nx,nfe-1,h] for h in 0:NCP) == 0)
    @NLconstraint(nlp, [nx=1:Nx,nfe=1,ncp=0], x[nx,nfe,ncp] - x0List[nx] == 0)
    @NLconstraint(nlp, [nx=1:Nx], xend[nx] - sum(colCont[h+1]*x[nx,NFE,h] for h in 0:NCP) == 0)

    # Defining u0 constraint
    @NLconstraint(nlp, [nu=1:Nu], u[nu,-1] - u0List[nu] == 0)

    # Defining the objective function that is to be minimized
    @NLobjective(nlp, Min, -(NFE/1)*(z[1,NFE,NCP]*x[1,NFE,NCP])
                        + sum(0.0154*(u[1,nfe]-u[1,nfe-1])^2 for nfe in 0:NFE-1)
                        + sum(5.5e-5*(u[2,nfe]-u[2,nfe-1])^2 for nfe in 0:NFE-1)
                        + sum(sum(1e6*e1[1,nfe,ncp]^2 for ncp in 0:NCP) for nfe in 1:NFE)
                        + sum(sum(1e10*e2[1,nfe,ncp]^2 for ncp in 0:NCP) for nfe in 1:NFE))
    return nlp
end

function nMPC(u0List::Vector{Float64},x0List::Array{Float64,1},z0List::Vector{Float64},pList::Array{Float64,1},NFE)
    Nx = 5
    Nu = 2
    Nz = 2
    Np = 2 
    NCP = 3
    nlp = nMPC0(u0List,x0List,z0List,pList,NFE)

    # Initial condition required for integration
    nlp.nlp_data.nlparamvalues[end-Nx+1:end] = x0

    ####Solving part:
    optimize!(nlp)

    optFlag = termination_status(nlp) == JuMP.MOI.TerminationStatusCode(4)


    # Return first optimal u and flag that can be used to indicate if optimal solution was found
    return value.(nlp[:u][:,0]).data,optFlag
end