content: # Linear and nonlinear control of startup of a CSTR.
  import mpctools as mpc
  import numpy as np
  import matplotlib.pyplot as plt
  from matplotlib import gridspec
  import time
  
  # Define some parameters and then the CSTR model.
  Nx = 3
  Nu = 2
  Nd = 1
  # Ny = Nx
  Delta = .25
  # eps = 1e-6 # Use this as a small number.
  
  T0 = 350
  c0 = 1
  r = .219
  k0 = 7.2e10
  E = 8750
  U = 54.94
  rho = 1000
  Cp = .239
  dH = -5e4
  
  def ode(x,u,d):
      # Grab the states, controls, and disturbance. We would like to write
      #    
      # [c, T, h] = x[0:Nx]
      # [Tc, F] = u[0:Nu]
      # [F0] = d[0:Nd]
      #    
      # but this doesn't work in Casadi 3.0. So, we're stuck with the following:
      c = x[0]
      T = x[1]
      h = x[2]
      Tc = u[0]
      F = u[1]
      F0 = d[0]
  
      # Now create the ODE.
      rate = k0*c*np.exp(-E/T)
          
      dxdt = np.array([
          F0*(c0 - c)/(np.pi*r**2*h) - rate,
          F0*(T0 - T)/(np.pi*r**2*h)
              - dH/(rho*Cp)*rate
              + 2*U/(r*rho*Cp)*(Tc - T),    
          (F0 - F)/(np.pi*r**2)
      ])
      return dxdt
  
  # Turn into casadi function and simulator.
  ode_casadi = mpc.getCasadiFunc(ode,[Nx,Nu,Nd],["x","u","d"],funcname="ode")
  ode_rk4_casadi = mpc.getCasadiFunc(ode,[Nx,Nu,Nd],["x","u","d"],
                                     funcname="ode_rk4",rk4=False,Delta=Delta)
  cstr = mpc.DiscreteSimulator(ode, Delta, [Nx,Nu,Nd], ["x","u","d"])
  
  # Steady-state values.
  cs = .878
  Ts = 324.5
  hs = .659
  Fs = .1
  Tcs = 300
  F0s = .1
  
  # Update the steady-state values a few times to make sure they don't move.
  for i in range(10):
      [cs,Ts,hs] = cstr.sim([cs,Ts,hs],[Tcs,Fs],[F0s]).tolist()
  xs = np.array([cs,Ts,hs])
  us = np.array([Tcs,Fs])
  ds = np.array([F0s])
  
  # Now get a linearization at this steady state.
  #ss = mpc.util.getLinearizedModel(ode_casadi, [xs,us,ds], ["A","B","Bp"], Delta)
  #A = ss["A"]
  #B = ss["B"]
  #Bp = ss["Bp"]
  #C = np.eye(Nx)
  
  # Weighting matrices for controller.
  Q = .5*np.diag(xs**-2)
  R = 2*np.diag(us**-2)
  
  # model_casadi = mpc.getCasadiFunc(ode,[Nx,Nu,Nd],["x","u","d"],funcname="cstr")
  
  #[K, Pi] = mpc.util.dlqr(A,B,Q,R)
  
  # Define casadi functions.
  Fnonlinear = ode_rk4_casadi
  
  # def measurement(x,d):
  #    return x
  # h = mpc.getCasadiFunc(measurement,[Nx,Nd],["x","d"],funcname="h")
  
  #def linmodel(x,u,d):
  #    Ax = mpc.mtimes(A,x-xs) + xs
  #    Bu = mpc.mtimes(B,u-us)
  #    Bpd = mpc.mtimes(Bp,d-ds)
  #    return Ax + Bu + Bpd
  #Flinear = mpc.getCasadiFunc(linmodel,[Nx,Nu,Nd],["x","u","d"],funcname="F")
  
  def stagecost(x,u,xsp,usp,Q,R):
      # Return deviation variables.
      dx = x - xsp
      du = u - usp
      # Calculate stage cost.
      return mpc.mtimes(dx.T,Q,dx) + mpc.mtimes(du.T,R,du)
  largs = ["x","u","x_sp","u_sp","Q","R"]
  l = mpc.getCasadiFunc(stagecost,[Nx,Nu,Nx,Nu,(Nx,Nx),(Nu,Nu)],largs,
                        funcname="l")
  
  def costtogo(x,xsp):
      # Deviation variables.
      dx = x - xsp
      
      # Calculate cost to go.
      return mpc.mtimes(dx.T,10*Q,dx)
  Pf = mpc.getCasadiFunc(costtogo,[Nx,Nx],["x","s_xp"],funcname="Pf")
  
  # First see what happens if we try to start up the reactor under no control.
  Nsim = 100
  x0 = np.array([.05*cs,.75*Ts,.5*hs])
  xcl = {}
  ucl = {}
  xcl["uncont"] = np.zeros((Nsim+1,Nx))
  xcl["uncont"][0,:] = x0
  ucl["uncont"] = np.tile(us,(Nsim,1))
  for t in range(Nsim):
      xcl["uncont"][t+1,:] = cstr.sim(xcl["uncont"][t,:],ucl["uncont"][t,:],ds)
  
  # Build a solver for the linear and nonlinear models.
  Nt = 15
  sp = {"x" : np.tile(xs, (Nt+1,1)), "u" : np.tile(us, (Nt,1))}
  #xguesslin = np.zeros((Nt+1,Nx))
  #xguesslin[0,:] = x0
  #for t in range(Nt):
  #    xguesslin[t+1,:] = A.dot(xguesslin[t,:] - xs) + xs
  #guesslin = {"x" : xguesslin, "u" : np.tile(us,(Nt,1))}
  guessnonlin = sp.copy()
  
  # Control bounds.
  umax = np.array([.05*Tcs,.15*Fs])
  dumax = .2*umax # Maximum for rate-of-change.
  bounds = dict(uub=[us + umax],ulb=[us - umax])
  ub = {"u" : np.tile(us + umax, (Nt,1)), "Du" : np.tile(dumax, (Nt,1))}
  lb = {"u" : np.tile(us - umax, (Nt,1)), "Du" : np.tile(-dumax, (Nt,1))}
  
  N = {"x":Nx, "u":Nu, "p":Nd, "t":Nt, "c":3}
  p = np.tile(ds, (Nt,1)) # Parameters for system.
  nmpc_commonargs = {
      "N" : N,
      "Delta": Delta,
      "x0" : x0,
      "lb" : lb,
      "ub" : ub,
      "p" : p,
      "verbosity" : 0,
      "Pf" : Pf,
      "l" : l,
      "sp" : sp,
      "uprev" : us,
      "funcargs" : {"l" : largs},
      "extrapar" : {"Q" : Q, "R" : R}, # In case we want to tune online.
  }
  solvers = {}
  # solvers["lmpc"] = mpc.nmpc(f=Flinear,guess=guesslin,**nmpc_commonargs)
  solvers["nmpc"] = mpc.nmpc(f=Fnonlinear,guess=guessnonlin,**nmpc_commonargs)
  
  # Also build steady-state target finders.
  contVars = [0,2]
  #sstarg_commonargs = {
  #    "N" : N,
  #    "lb" : {"u" : np.tile(us - umax, (1,1))},
  #    "ub" : {"u" : np.tile(us + umax, (1,1))},
  #    "verbosity" : 0,
  ##     "h" : h,
  #    "p" : np.array([ds]),
  #}
  #sstargs = {}
  # sstargs["lmpc"] = mpc.sstarg(f=Flinear,**sstarg_commonargs)
  # sstargs["nmpc"] = mpc.sstarg(f=Fnonlinear,**sstarg_commonargs)
  
  # Now simulate the process under control.
  tcl = {}
  
  for method in solvers.keys():
      xcl[method] = np.zeros((Nsim+1,Nx))
      xcl[method][0,:] = x0
      tcl[method] = np.zeros((Nsim+1,1))
      thisx = x0    
      ucl[method] = np.zeros((Nsim,Nu))
      # ysp = np.tile(xs,(Nsim+1,1))
      
      xsp = np.zeros((Nsim+1,Nx))
      usp = np.zeros((Nsim,Nu))    
      
      # ysp[int(Nsim/3):int(2*Nsim/3),:] = xs*np.array([.85,.75,1.15])    
      for t in range(Nsim):
          # Figure out setpoints.
  #        if t == 0 or not np.all(ysp[t,:] == ysp[t-1,:]):        
  #            thisysp = ysp[t,:]            
  #            sstargs[method].fixvar("y",0,thisysp[contVars],contVars)
  #            sstargs[method].guess["u",0] = us
  #            sstargs[method].guess["x",0] = thisysp
  #            sstargs[method].guess["y",0] = thisysp
  #            sstargs[method].solve()            
  #            
  #            print "%10s %3d: %s" % ("sstarg",t,sstargs[method].stats["status"])
  #            if sstargs[method].stats["status"] != "Solve_Succeeded":
  #                print "***Target finder failed!"
  #                break
  #            
  #            xsp[t,:] = np.squeeze(sstargs[method].var["x",0])
  #            usp[t,:] = np.squeeze(sstargs[method].var["u",0])
  #            
  #            solvers[method].par["x_sp"] = [xsp[t,:]]*(Nt + 1)
  #            solvers[method].par["u_sp"] = [usp[t,:]]*Nt
          
          # Fix initial condition and solve.
          t0 = time.time()
          solvers[method].fixvar("x",0,thisx)
          solvers[method].solve()
          print "%10s %3d: %s" % (method,t,solvers[method].stats["status"])
          if solvers[method].stats["status"] != "Solve_Succeeded":
              print "***Solver failed!"         
              break
          else:
              solvers[method].saveguess()
              
          thisu = np.squeeze(solvers[method].var["u"][0])
          ucl[method][t,:] = thisu
          t1 = time.time()
          tcl[method][t] = t1-t0
          thisx = cstr.sim(thisx,thisu,ds)
          xcl[method][t+1,:] = thisx
          
          # Update previous u.
          solvers[method].par["u_prev",0] = ucl[method][t,:]
  
  # Define plotting function.
  def cstrplot(x,u,xsp=None,contVars=[],title=None,colors={},labels={},
               markers={},keys=None,bounds=None,ilegend=0):
      if keys is None:
          keys = x.keys()
      for k in keys:    
          u[k] = np.concatenate((u[k],u[k][-1:,:]))
      ylabelsx = ["$c$ (mol/L)", "$T$ (K)", "$h$ (m)"]
      ylabelsu = ["$T_c$ (K)", "$F$ (kL/min)"]
      
      gs = gridspec.GridSpec(Nx*Nu,2)    
      fig = plt.figure(figsize=(10,6),facecolor="none")
      leglines = []
      leglabels = []
      for i in range(Nx):
          ax = fig.add_subplot(gs[i*Nu:(i+1)*Nu,0])
          for k in keys:
              t = np.arange(0,x[k].shape[0])*Delta
              args = {"color":colors.get(k,"black"), "label":labels.get(k,k),
                      "marker":markers.get(k,"")}
              [line] = ax.plot(t,x[k][:,i],markeredgecolor="none",**args)
              if i == ilegend:
                  leglines.append(line)
                  leglabels.append(args["label"])
          if i in contVars and xsp is not None:
              ax.step(t,xsp[:,i],linestyle="--",color="black",where="post")
          ax.set_ylabel(ylabelsx[i])
          mpc.plots.zoomaxis(ax,yscale=1.1)
          mpc.plots.prettyaxesbox(ax)
          mpc.plots.prettyaxesbox(ax,
              facecolor="white",front=False)
      ax.set_xlabel("Time (min)")
      for i in range(Nu):
          ax = fig.add_subplot(gs[i*Nx:(i+1)*Nx,1])
          for k in keys:
              t = np.arange(0,u[k].shape[0])*Delta
              args = {"color":colors.get(k,"black"), "label":labels.get(k,k)}
              ax.step(t,u[k][:,i],where="post",**args)
          if bounds is not None:
              for b in set(["uub", "ulb"]).intersection(bounds.keys()):
                  ax.plot(np.array([t[0],t[-1]]),np.ones((2,))*bounds[b][i],
                          '--k')
          ax.set_ylabel(ylabelsu[i])
          mpc.plots.zoomaxis(ax,yscale=1.25)
          mpc.plots.prettyaxesbox(ax)
          mpc.plots.prettyaxesbox(ax,
              facecolor="white",front=False)
      ax.set_xlabel("Time (min)")
      fig.legend(leglines,leglabels,loc="lower center",ncol=len(keys))
      fig.tight_layout(pad=.5,rect=(0,.075,1,1))
      if title is not None:
          fig.canvas.set_window_title(title)
      return fig
  
  x = xcl['nmpc']
  u = ucl['nmpc']
  ptimes = tcl['nmpc']
  # Make plots.
  keys = ["uncont", "nmpc"]
  colors = {"lmpc":"blue", "nmpc":"green", "uncont":"red"}
  labels = {"lmpc":"LMPC", "nmpc":"NMPC", "uncont":"Uncontrolled"}
  markers = {"lmpc":"s", "nmpc":"o", "uncont":"^"}
  plotbounds = dict([(k,bounds[k][0]) for k in ["ulb","uub"]])
  fig = cstrplot(xcl, ucl, colors=colors, contVars=contVars, labels=labels,
                 keys=keys, markers={}, bounds=plotbounds, ilegend=2)
  fig.show()
  # mpc.plots.showandsave(fig,"cstr_startup.pdf",facecolor="none")
  