#include "sopnamsp.h"
#include "difeq.h"
#include "ctimer.h"

/*!
  \ingroup NTools
  \class SOPHYA::DiffEqSolver
  \brief Classe abstraite de résolveur d'équation différentielles.

  Beaucoup de méthodes renvoient l'objet afin de pouvoir
  utiliser une notation chaînée du type:

  s.Step(...).Start(...).Solve(...)

  \warning statut EXPERIMENTAL , NON TESTE

  \sa SOPHYA::DiffEqFunction
  \sa SOPHYA::RK4DiffEq  SOPHYA::RK4CDiffEq 
*/

//!	Constructeur vide. 
DiffEqSolver::DiffEqSolver()
: mFunc(NULL), mOwnFunc(false),
  mYStart(1), mXStart(0), mStep(0.01)
{}

//!  Constructeur général. L'équation est donnée sous forme de DiffEqFunction
DiffEqSolver::DiffEqSolver(DiffEqFunction* f)
: mFunc(f), mOwnFunc(false),
  mYStart(mFunc->NFuncReal()), mXStart(0), mStep(0.01)
{}

/*!
  Constructeur pour le cas particulier d'une équation du premier
  ordre. Voir DiffEqFcn1. La fonction f correspond à l'équation
  y' = f(y).
*/
DiffEqSolver::DiffEqSolver(DIFEQFCN1 f)
: mFunc(new DiffEqFcn1(f)), mOwnFunc(true),
  mYStart(mFunc->NFuncReal()), mXStart(0), mStep(0.01)
{}

DiffEqSolver::~DiffEqSolver()
{
  if (mOwnFunc) delete mFunc;
}


/*!
  Permet de spécifier l'équation différentielle, sous la forme
  d'une DiffEqFunction. Retourne l'objet DiffEqSolver : notation
  chaînée possible
*/
DiffEqSolver&
DiffEqSolver::Func(DiffEqFunction* f)
{
  if (mFunc && mOwnFunc) delete mFunc;
  mFunc    = f;
  mOwnFunc = false;
  return *this;
}

/*!
  Permet de spécifier l'équation différentielle, sous la forme
  d'une fonction f telle que y' = f(y). Retourne l'objet DiffEqSolver : 
  notation chaînée possible.
*/
DiffEqSolver&
DiffEqSolver::Func(DIFEQFCN1 f)
{
  if (mFunc && mOwnFunc) delete mFunc;
  mFunc    = new DiffEqFcn1(f);
  mOwnFunc = true;
  return *this;
}

/*!
  Spécifie le pas d'intégration de l'équation différentielle (pour
  les méthodes où ce pas a un sens). La valeur par défaut est 0.01.
  Retourne l'objet DiffEqSolver : notation chaînée possible.
*/
DiffEqSolver&
DiffEqSolver::Step(double step)
{
  mStep = step;
  return *this;
}


/*!
  Spécifie le point de départ de l'intégration de l'équadif.
  Le vecteur \b y contient les valeurs intiales. Voir la description
  de chaque fonction-objet-équation différentielle pour connaître
  le rôle de chaque élément du vecteur. \b t est la valeur initiale 
  du temps.
  La dimension de y doit correspondre au nombre apparent de fonctions.
  Retourne l'objet DiffEqSolver : notation chaînée possible.
*/
DiffEqSolver&
DiffEqSolver::StartV(Vector const& yi, double t)
{
  ASSERT(mFunc != NULL);
  ASSERT(yi.NElts() == mFunc->NFunc());   // Nombre apparent
  mYStart = yi;
  mXStart = t;
  mFunc->AdjustStart(mYStart,t);
  ASSERT(mYStart.NElts() == mFunc->NFuncReal());
  return *this;
}

/*!  
  Spécifie le point de départ de l'intégration de l'équadif,
  pour une équation y' = f(y). On donne le temps et la valeur de y.
  Retourne l'objet DiffEqSolver : notation chaînée possible.
*/
DiffEqSolver&
DiffEqSolver::Start1(double y, double t)
{
  ASSERT(mFunc != NULL);
  ASSERT(mFunc->NFunc() == 1);   // Nombre apparent
  mYStart.Realloc(1);
  mYStart(0) = y;
  mXStart = t;
  mFunc->AdjustStart(mYStart,t);
  ASSERT(mYStart.NElts() == mFunc->NFuncReal());
  return *this;
}


/*!
  Spécifie le point de départ de l'intégration de l'équadif.
  Le tableau \b yi contient les valeurs intiales. Voir la description
  de chaque fonction-objet-équation différentielle pour connaître
  le rôle de chaque élément du tableau. t est la valeur initiale 
  du temps.
  Le nombre d'éléments de yi doit correspondre au nombre apparent
  de fonctions.
  Retourne l'objet DiffEqSolver : notation chaînée possible.
*/
DiffEqSolver&
DiffEqSolver::Start(double const* yi, double t)
{
  mYStart.Realloc(mFunc->NFunc());
  for (int i=0; i<mFunc->NFunc(); i++)
    mYStart(i) = yi[i];
  mXStart = t;
  mFunc->AdjustStart(mYStart,t);
  ASSERT(mYStart.NElts() == mFunc->NFuncReal());
  return *this;
}

/*!
  \fn virtual void DiffEqSolver::SolveArr(Matrix&  y, double* t, double tf, int n)

  Lance la résolution de l'équadif, jusqu'à la valeur
  tf du temps. N valeurs intermédiaires sont retournées dans 
  les tableaux y et t :
  - t[i] est le temps au pas i, pour 0<=i<n
  - y[i] sont les valeurs des fonctions au pas i.

  Voir la description	
  de chaque fonction-objet-équation différentielle pour connaître
  le rôle de chaque élément de y[i].
  t[n-1] vaut tf.
  Le nombre d'éléments de y[i] doit correspondre au nombre apparent
  de fonctions.
  Cette fonction doit être redéfinie pour chaque méthode de résolution
  d'équations. Elle est appelée par toutes les autres fonctions de
  résolution (autres SolveArr, et Solve).
*/


/*!
  Lance la résolution de l'équadif, jusqu'à la valeur
  tf du temps. Les valeurs finales sont retournées dans yf.
  Voir la description
  de chaque fonction-objet-équation différentielle pour connaître
  le rôle de chaque élément de yf.
  Le nombre d'éléments de yf doit correspondre au nombre apparent
  de fonctions.
*/
void
DiffEqSolver::SolveV(Vector& yf, double tf)
{
  double t;
  Matrix m(1, mFunc->NFuncReal());
  SolveArr(m, &t, tf, 1);
  yf.Realloc(mFunc->NFunc());
  for (int i=0; i<mFunc->NFunc(); i++)
    yf(i) = m(0,i);
}


/*!
  Lance la résolution de l'équadif, jusqu'à la valeur
  tf du temps, pour une équation du premier ordre y' = f(y).
  La valeur finale de y est retournée dans yf.
*/
void
DiffEqSolver::Solve1(double& yf, double tf)
{
  ASSERT(mFunc->NFunc() == 1);
  double t;
  Matrix m(1,mFunc->NFuncReal());
  SolveArr(m, &t, tf, 1);
  yf = m(0,0);
}


/*!
  Lance la résolution de l'équadif, jusqu'à la valeur	
  tf du temps. Les valeurs finales sont retournées dans yf.
  Voir la description
  de chaque fonction-objet-équation différentielle pour connaître
  le rôle de chaque élément de yf.
  Le nombre d'éléments de yf doit correspondre au nombre apparent
  de fonctions.
*/
void
DiffEqSolver::Solve(double* yf, double tf)
{
  double t;
  Matrix m(1, mFunc->NFuncReal());
  SolveArr(m, &t, tf, 1);
  for (int i=0; i<mFunc->NFunc(); i++)
    yf[i] = m(0,i);
}

/*!
  Lance la résolution de l'équadif (du premier ordre), jusqu'à la valeur
  tf du temps. N valeurs intermediaires sont retournées dans 
  les tableaux y et t :
  t[i] est le temps au pas i, pour 0<=i<n
  y[i] est la valeur de la fonction au pas i.
  t[n-1] vaut tf.
*/
void
DiffEqSolver::SolveArr1(double*  y, double* t, double tf, int n)

{
  ASSERT(mFunc->NFunc() == 1);
  Matrix m(n, mFunc->NFuncReal());
  SolveArr(m, t, tf, n);
  for (int i=0; i<n; i++)
    y[i] = m(i,0);
}

/*!
  Lance la résolution de l'équadif, jusqu'à la valeur
  tf du temps. N valeurs intermédiaires sont retournées dans 
  les tableaux y et t :
  t[i] est le temps au pas i, pour 0<=i<n
  y[i] sont les valeurs des fonctions au pas i.
  Voir la description
  de chaque fonction-objet-équation différentielle pour connaître
  le rôle de chaque élément de y[i].
  t[n-1] vaut tf.
  Le nombre d'éléments de y[i] doit correspondre au nombre apparent
  de fonctions.
*/
void
DiffEqSolver::SolveArr2(double** y, double* t, double tf, int n)
{
   Matrix m(n, mFunc->NFuncReal());
   SolveArr(m, t, tf, n);
   for (int i=0; i<n; i++)
     for (int j=0; j<mFunc->NFunc(); j++)
       y[i][j] = m(i,j);
}


/*!
  \ingroup NTools 
  \class SOPHYA::DiffEqFunction
  
  \brief Classe de fonctions pour la résolution d'équations différentielles.
  
  On résoud de facon générale un système de n équations différentielles
  du premier ordre, donnant les dérivées fpi de n fonctions fi. 
  Cette méthode permet de résoudre toutes les sortes d'équations :
  pour une équation du second ordre on crée une fonction intermédiaire
  qui est la dérivée de la fonction cherchée.

  \warning statut EXPERIMENTAL , NON TESTE

  \sa SOPHYA::DiffEqSolver
*/



/*!
  \fn DiffEqFunction::DiffEqFunction(int n, int napp)
  Constructeur. N est le nombre réel de fonctions dans le
  système, et napp le nombre apparent (pour l'utilisateur).
  En effet, dans certains cas, on peut avoir besoin de fonctions
  supplémentaires masquées à l'utilisateur, par exemple la fonction
  constante valant 1, si l'équadif fait intervenir le temps (t).
*/

/*!
  \fn virtual void ComputeV(Vector& fpi, Vector const& fi)
  Calcule les valeurs des dérivées fpi à partir des valeurs 
  des fonctions fi. A redéfinir.
*/

/*!
  \fn virtual void Compute(double& fp, double f)
  Dans le cas où il y a une seule fonction, calcule la dérivée
  fp à partir de la valeur de la fonction f. A redéfinir.
*/


//++
// virtual void AdjustStart(Vector& start, double tstart)
//	Pour ajuster le vecteur de départ quand il y a des 
//	fonctions à usage interne...
//--



/*! 
  \ingroup NTools
  \class SOPHYA::DiffEqFcn1
  
  \brief Objet-fonction générique pour la résolution d'équations
  différentielles de la forme y' = f(y).
  
  On fournit une fonction de type <br>
  <tt> typedef double(*DIFEQFCN1)(double); </tt> <br>
  qui retourne y' en fonction de y.

  \warning statut EXPERIMENTAL , NON TESTE
*/


//!	Constructeur. On fournit la fonction f tq y'=f(y)
DiffEqFcn1::DiffEqFcn1(DIFEQFCN1 fcn)
: DiffEqFunction(1), mFcn(fcn)
{}


//! Implementation de Compute qui va utiliser la fonction fournie au constructeur.
void
DiffEqFcn1::Compute(double& fp, double f)
{
   fp = (*mFcn)(f);
}

/*!
  \ingroup NTools
  \class   SOPHYA::DiffEqFcnT1

  Objet-fonction générique pour la résolution d'équations
  différentielles de la forme y' = f(y,t).
  On fournit une fonction de type
  <tt>	typedef \tt double(*DIFEQFCNT1)(double, double); </tt>
  qui retourne y' en fonction de y et t.

  \verbatim
  Note : le système résolu est alors en fait
  y'[0] = fcn(y[0], y[1])
  y'[1] = 1
  \endverbatim

  \warning statut EXPERIMENTAL , NON TESTE
*/

//! Constructeur. On fournit la fonction f tq y' = f(y,t)
DiffEqFcnT1::DiffEqFcnT1(DIFEQFCNT1 fcn)
: DiffEqFunction(2, 1), mFcn(fcn)
{}

void
DiffEqFcnT1::ComputeV(Vector& fpi, Vector const& fi)
{
  fpi(0) = (*mFcn)(fi(0), fi(1));
  fpi(1) = 1;
}

void
DiffEqFcnT1::AdjustStart(Vector& start, double tstart)
{
  start.Realloc(2);
  start(1) = tstart;
}

/*!
  \ingroup NTools
  \class   SOPHYA::DiffEqFcn2
  
  Objet-fonction générique pour la résolution d'équations
  différentielles de la forme y'' = f(y',y).
  On fournit une fonction de type
  <tt>	typedef double(*DIFEQFCN2)(double, double); </tt>
  qui retourne y'' en fonction de y' et y.
  \verbatim
  Note : le système résolu est en fait
  y'[0] = y[1]
  y'[1] = f(y[1], y[0])
  \endverbatim

  \warning statut EXPERIMENTAL , NON TESTE
*/

//!   Constructeur. On fournit la fonction f tq y''=f(y',y)
DiffEqFcn2::DiffEqFcn2(DIFEQFCN2 fcn)
: DiffEqFunction(2), mFcn(fcn)
{}

void
DiffEqFcn2::ComputeV(Vector& fpi, Vector const& fi)
{
  fpi(0) = fi(1);
  fpi(1) = (*mFcn)(fi(1), fi(0));
}

/*!
  \ingroup NTools
  \class SOPHYA::DiffEqFcnT2
  
  Objet-fonction générique pour la résolution d'équations
  différentielles de la forme y'' = f(y',y,t).
  On fournit une fonction de type <br>
  <tt> typedef double(*DIFEQFCNT2)(double, double, double); </tt> <br>
  qui retourne y'' en fonction de y', y et t.
  \verbatim
  Note : le système résolu est alors en fait
  y'[0] = y[1]
  y'[1] = f(y[1], y[0], y[2])
  y'[2] = 1
  \endverbatim

  \warning statut EXPERIMENTAL , NON TESTE
*/

//!	Constructeur. On fournit la fonction f tq y'' = f(y',y,t)
DiffEqFcnT2::DiffEqFcnT2(DIFEQFCNT2 fcn)
: DiffEqFunction(3,2), mFcn(fcn)
{}

void
DiffEqFcnT2::ComputeV(Vector& fpi, Vector const& fi)
{
  fpi(0) = fi(1);
  fpi(1) = (*mFcn)(fi(1), fi(0), fi(2));
  fpi(2) = 1;
}

void
DiffEqFcnT2::AdjustStart(Vector& start, double tstart)
{
  start.Realloc(3);
  start(2) = tstart;
}


/*!
  \ingroup NTools
  \class  SOPHYA::DiffEqFcnV
  
  Objet-fonction générique pour la résolution d'équations
  différentielles 3D de la forme y'' = f(y',y), ou y est
  un vecteur de dimension 3.
  On fournit une fonction de type <br>
  <tt> typedef void(*DIFEQFCNV)(Vector& y2, Vector const& y1, <br>
  Vector const& y); </tt> <br>
  qui retourne y'' en fonction de y' et y.
  \verbatim
  Note : le système résolu est alors en fait
  v(0-2) = y, v(3-5) = y'

  v'[0] = v[3]
  v'[1] = v[4]
  v'[2] = v[5]
  v'[3-5] = f(v[3-5], v[0-2])
  \endverbatim

  \warning statut EXPERIMENTAL , NON TESTE
*/

DiffEqFcnV::DiffEqFcnV(DIFEQFCNV fcn)
: DiffEqFunction(6), mFcn(fcn),
  tmp1(3), tmp2(3), tmp3(3)
{}

void
DiffEqFcnV::ComputeV(Vector& fpi, Vector const& fi)
{
  fpi(0) = fi(3);
  fpi(1) = fi(4);
  fpi(2) = fi(5);
  
  tmp2(0) = fi(3); tmp2(1) = fi(4); tmp2(2) = fi(5);
  tmp3(0) = fi(0); tmp3(1) = fi(1); tmp3(2) = fi(2);
  
  (*mFcn)(tmp1, tmp2, tmp3);
  
  fpi(3) = tmp1(0); fpi(4) = tmp1(1); fpi(5) = tmp1(2);
}


/*!
  \ingroup NTools
  \class SOPHYA::RK4DiffEq
  
  Classe de résolution d'équadif par la méthode de
  Runge-Kutta d'ordre 4.
  Voir DiffEqSolver pour les méthodes.

  \warning statut EXPERIMENTAL , NON TESTE
*/

RK4DiffEq::RK4DiffEq()
: DiffEqSolver(), k1(10), k2(10), k3(10), k4(10)
{}

//!	Constructeur général
RK4DiffEq::RK4DiffEq(DiffEqFunction* f)
: DiffEqSolver(f), k1(f->NFuncReal()), 
  k2(f->NFuncReal()), k3(f->NFuncReal()), k4(f->NFuncReal())
{}


//!	Constructeur pour y' = f(y)
RK4DiffEq::RK4DiffEq(DIFEQFCN1 f)
: DiffEqSolver(f), k1(1), k2(1), k3(1), k4(1)
{}

void
RK4DiffEq::SolveArr(Matrix& y, double* t, double tf, int n)
{
  //TIMEF;
  // On calcule le nombre de sous-pas par pas
   
   int nStep = (mStep > 0) ? int((tf - mXStart)/n/mStep) : 1;
   if (nStep <= 1) nStep = 1;
   
   double dx = (tf - mXStart)/(n*nStep);
   
   Vector yt(mYStart,false);

   k1.Realloc(mFunc->NFuncReal());
   k2.Realloc(mFunc->NFuncReal());
   k3.Realloc(mFunc->NFuncReal());
   k4.Realloc(mFunc->NFuncReal());
   
      
   for (int i=0; i<n; i++) {
     for (int j=0; j<nStep; j++)
       RKStep(yt, yt, dx);
     for (int k=0; k<mFunc->NFunc(); k++)
       y(i,k) = yt(k);
     t[i] = (i+1)*dx*nStep + mXStart;  
   }
}

void
RK4DiffEq::RKStep(Vector& newY, Vector const& y0, double dt)
{
    mFunc->ComputeV(k1, y0);
    k1 *= dt;
    mFunc->ComputeV(k2, y0 + k1/2.);
    k2 *= dt;
    mFunc->ComputeV(k3, y0 + k2/2.);
    k3 *= dt;
    mFunc->ComputeV(k4, y0 + k3);
    k4 *= dt;
    
    newY = y0 + (k1 + k2*2. + k3*2. + k4)/6.;
}