//////////////////////////////////////////////////////////////////////////
////////////////                mpec.cxx             /////////////////////
//////////////////////////////////////////////////////////////////////////
////////////////           PSOPT  Example            /////////////////////
//////////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////
//////// Title:         Dynamic MPEC problem              ////////////////
//////// Last modified: 27 May 2011                       ////////////////
//////// Reference:     Betts (2010)                      ////////////////
//////// (See PSOPT handbook for full reference)          ////////////////
//////////////////////////////////////////////////////////////////////////
////////     Copyright (c) Victor M. Becerra, 2009        ////////////////
//////////////////////////////////////////////////////////////////////////
//////// This is part of the PSOPT software library, which////////////////
//////// is distributed under the terms of the GNU Lesser ////////////////
//////// General Public License (LGPL)                    ////////////////
//////////////////////////////////////////////////////////////////////////

#include "psopt.h"

//////////////////////////////////////////////////////////////////////////
///////////////////  Define the end point (Mayer) cost function //////////
//////////////////////////////////////////////////////////////////////////

adouble endpoint_cost(adouble* initial_states, adouble* final_states,
                      adouble* parameters,adouble& t0, adouble& tf,
                      adouble* xad, int iphase, Workspace* workspace)
{
   adouble y_f = final_states[ 0 ];

   return pow( y_f - (5.0/3.0) , 2.0);
}

//////////////////////////////////////////////////////////////////////////
///////////////////  Define the integrand (Lagrange) cost function  //////
//////////////////////////////////////////////////////////////////////////

adouble integrand_cost(adouble* states, adouble* controls, adouble* parameters,
                     adouble& time, adouble* xad, int iphase, Workspace* workspace)
{
    adouble y = states[   0 ];
    adouble s = controls[ 0 ];
    adouble p = controls[ 1 ];
    adouble q = controls[ 2 ];


    adouble retval;

    double rho = 1.e3;

    retval = y*y + rho*( p*(s+1.0)+ q*(1.0-s));

    return  retval;

}


//////////////////////////////////////////////////////////////////////////
///////////////////  Define the DAE's ////////////////////////////////////
//////////////////////////////////////////////////////////////////////////

void dae(adouble* derivatives, adouble* path, adouble* states,
         adouble* controls, adouble* parameters, adouble& time,
         adouble* xad, int iphase, Workspace* workspace)
{


   adouble y = states[ 0 ];
   adouble ydot;


   adouble s = controls[ 0 ];
   adouble p = controls[ 1 ];
   adouble q = controls[ 2 ];

   ydot = 2.0 -s;

   derivatives[ 0 ] = ydot;


   path[ 0 ] = -y - p + q;

}

////////////////////////////////////////////////////////////////////////////
///////////////////  Define the events function ////////////////////////////
////////////////////////////////////////////////////////////////////////////

void events(adouble* e, adouble* initial_states, adouble* final_states,
            adouble* parameters,adouble& t0, adouble& tf, adouble* xad,
            int iphase, Workspace* workspace)

{
   adouble y0 = initial_states[ 0 ];


   e[ 0 ] = y0;

}



///////////////////////////////////////////////////////////////////////////
///////////////////  Define the phase linkages function ///////////////////
///////////////////////////////////////////////////////////////////////////

void linkages( adouble* linkages, adouble* xad, Workspace* workspace)
{
  // No linkages as this is a single phase problem
}



////////////////////////////////////////////////////////////////////////////
///////////////////  Define the main routine ///////////////////////////////
////////////////////////////////////////////////////////////////////////////

int main(void)
{

////////////////////////////////////////////////////////////////////////////
///////////////////  Declare key structures ////////////////////////////////
////////////////////////////////////////////////////////////////////////////

    Alg  algorithm;
    Sol  solution;
    Prob problem;

////////////////////////////////////////////////////////////////////////////
///////////////////  Register problem name  ////////////////////////////////
////////////////////////////////////////////////////////////////////////////

    problem.name        		= "Dynamic MPEC problem";
    problem.outfilename                 = "mpec.txt";

////////////////////////////////////////////////////////////////////////////
////////////  Define problem level constants & do level 1 setup ////////////
////////////////////////////////////////////////////////////////////////////

    problem.nphases   					= 1;
    problem.nlinkages               = 0;

    psopt_level1_setup(problem);

/////////////////////////////////////////////////////////////////////////////
/////////   Define phase related information & do level 2 setup  ////////////
/////////////////////////////////////////////////////////////////////////////

    problem.phases(1).nstates   		= 1;
    problem.phases(1).ncontrols 		= 3;
    problem.phases(1).nevents   		= 1;
    problem.phases(1).npath     		= 1;
    problem.phases(1).nodes         << 20;

    psopt_level2_setup(problem, algorithm);

////////////////////////////////////////////////////////////////////////////
///////////////////  Declare MatrixXd objects to store results //////////////
////////////////////////////////////////////////////////////////////////////

    MatrixXd y, controls, s, p, q, t;

////////////////////////////////////////////////////////////////////////////
///////////////////  Enter problem bounds information //////////////////////
////////////////////////////////////////////////////////////////////////////

    double y0 = -1.0;

    problem.phases(1).bounds.lower.states(0) = -10.0;



    problem.phases(1).bounds.upper.states(0) = 10.0;



    problem.phases(1).bounds.lower.controls(0) = -1.0;
    problem.phases(1).bounds.lower.controls(1) = 0.0;
    problem.phases(1).bounds.lower.controls(2) = 0.0;

    problem.phases(1).bounds.upper.controls(0) = 1.0;
    problem.phases(1).bounds.upper.controls(1) = 10.0;
    problem.phases(1).bounds.upper.controls(2) = 10.0;

    problem.phases(1).bounds.lower.events(0) = y0;
    problem.phases(1).bounds.upper.events(0) = y0;


    problem.phases(1).bounds.upper.path(0) = 0.0;
    problem.phases(1).bounds.lower.path(0) = 0.0;


    problem.phases(1).bounds.lower.StartTime    = 0.0;
    problem.phases(1).bounds.upper.StartTime    = 0.0;

    problem.phases(1).bounds.lower.EndTime      = 2.0;
    problem.phases(1).bounds.upper.EndTime      = 2.0;



////////////////////////////////////////////////////////////////////////////
///////////////////  Register problem functions  ///////////////////////////
////////////////////////////////////////////////////////////////////////////


    problem.integrand_cost 	= &integrand_cost;
    problem.endpoint_cost 		= &endpoint_cost;
    problem.dae             	= &dae;
    problem.events 				= &events;
    problem.linkages				= &linkages;

////////////////////////////////////////////////////////////////////////////
///////////////////  Define & register initial guess ///////////////////////
////////////////////////////////////////////////////////////////////////////

    int nnodes    			 = problem.phases(1).nodes(0);
    int ncontrols           = problem.phases(1).ncontrols;
    int nstates             = problem.phases(1).nstates;


    problem.phases(1).guess.controls       = zeros(ncontrols,nnodes);
    problem.phases(1).guess.states         = y0*ones(1,nnodes);
    problem.phases(1).guess.time           = linspace(0.0,2.0,nnodes);


////////////////////////////////////////////////////////////////////////////
///////////////////  Enter algorithm options  //////////////////////////////
////////////////////////////////////////////////////////////////////////////


    algorithm.nlp_iter_max                = 1000;
    algorithm.nlp_tolerance               = 1.e-4;
    algorithm.nlp_method                  = "IPOPT";
    algorithm.scaling                     = "automatic";
    algorithm.derivatives                 = "automatic";
    algorithm.mesh_refinement             = "automatic";
    algorithm.collocation_method 			= "trapezoidal";
    algorithm.ode_tolerance               = 1.e-6;



////////////////////////////////////////////////////////////////////////////
///////////////////  Now call PSOPT to solve the problem   /////////////////
////////////////////////////////////////////////////////////////////////////

    psopt(solution, problem, algorithm);

////////////////////////////////////////////////////////////////////////////
///////////  Extract relevant variables from solution structure   //////////
////////////////////////////////////////////////////////////////////////////


    y        = solution.get_states_in_phase(1);
    controls = solution.get_controls_in_phase(1);
    t        = solution.get_time_in_phase(1);
    s        = controls.row(0); 
    p        = controls.row(1); 
    q        = controls.row(2); 


////////////////////////////////////////////////////////////////////////////
///////////  Save solution data to files if desired ////////////////////////
////////////////////////////////////////////////////////////////////////////

    Save(y,"y.dat");
    Save(controls,"controls.dat");
    Save(t,"t.dat");

////////////////////////////////////////////////////////////////////////////
///////////  Plot some results if desired (requires gnuplot) ///////////////
////////////////////////////////////////////////////////////////////////////

    plot(t,y,problem.name+": state", "time (s)", "state","y");

    plot(t,s,problem.name+": algebraic variable s","time (s)", "s", "s");

    plot(t,p,problem.name+": algebraic variable p","time (s)", "p", "p");

    plot(t,q,problem.name+": algebraic variable q","time (s)", "q", "q");


    plot(t,y,problem.name+": state", "time (s)", "state","y",
         "pdf", "y.pdf");

    plot(t,s,problem.name+": algebraic variable s","time (s)", "s", "s",
         "pdf", "s.pdf");

    plot(t,p,problem.name+": algebraic variable p","time (s)", "p", "p",
         "pdf", "p.pdf");

    plot(t,q,problem.name+": algebraic variable q","time (s)", "q", "q",
         "pdf", "q.pdf");



}

////////////////////////////////////////////////////////////////////////////
///////////////////////      END OF FILE     ///////////////////////////////
////////////////////////////////////////////////////////////////////////////
