# include <cmath>
# include <cstdlib>
# include <cstring>
# include <fstream>
# include <iomanip>
# include <iostream>

using namespace std;

# include "fsolve.hpp"
# include "midpoint_adaptive.hpp"

//****************************************************************************80

void mad_compute ( void f ( double t, double y[], double dydt[] ), double t0, 
  double tmax, int m, double y0[], double tau0, double reltol, double abstol, 
  int nmax, int &nstep, int &n_rejected, int &n_fsolve, double t[], double y[] )

//****************************************************************************80
//
//  Purpose:
//
//    mad_compute() estimates the solution of an ordinary differential equation.
//
//  Discussion:
//
//    mad_compute() uses an adaptive time stepping algorithm for the
//    implicit midpoint method, applied to a system of ordinary differential
//    equations.
//
//    The implicit midpoint method has local truncation error O(h^3).
//
//    The time step is adaptively controlled with respect to relative and
//    absolute tolerances applied to the estimated local truncation error (LTE).
//  
//    The adaptivity with respect to relative and absolute error tolerances
//    is described on page 168 in the Hairer reference.
//
//  Licensing:
//
//    This code is distributed under the MIT license.
//
//  Modified:
//
//    10 July 2024
//
//  Author:
//
//    John Burkardt, Catalin Trenchea
//
//  Reference:
//
//    Catalin Trenchea, John Burkardt,
//    Refactorization of the midpoint rule,
//    Applied Mathematics Letters,
//    Volume 107, September 2020.
//
//    Ernst Hairer, Syvert Norsett, Gerhard Wanner,
//    Solving ordinary differential equations, I. Nonstiff problems, 
//    Springer Series in Computational Mathematics, Number 8. 
//    Springer-Verlag, Berlin, 1987.
//
//  Input:
//
//    void f ( double t, double y[], double dydt[] ): evaluates the right 
//    hand side of the differential equation.
//
//    double t0, tmax: the interval of integration.
//
//    int m: the dimension of a single solution vector.
//
//    double y0[m]: the vector of initial conditions at the starting time. 
//
//    double tau0: the initial timestep.
//
//    double reltol, abstol: the tolerances for the local truncation error.
//
//    int nmax: the maximum number of time steps allowed.
//
//  Output:
//
//    int &nstep: the number of time steps taken.
//
//    int &n_rejected: the number of rejected time steps.
//
//    int &n_fsolve: the number of failed calls to fsolve().
//
//    double t[0:nmax], y[0:nmax,1:m]: contains in entries 0:nstep the 
//    the sequence of times and solution estimates actually computed.
//
//  Local:
//
//    int count: the number of steps, whether accepted or rejected.
//    It should be the case that n <= count.
//
//    double kappa: a factor for the stepsize update.
//    0 < kappa <= 1.  The value kappa = 0.85 is suggested.
//
//    double theta: the theta-method parameter.
//    theta = 0.0 for the backward Euler method.
//    theta = 0.5 for the midpoint method.
//    theta = 1.0 for the forward Euler method.
//
{
  int count;
  double factor;
  bool fsolve_fail;
  int i;
  double kappa;
  int n;
  double *tau;
  double theta;
  double tn;
  double tnmax;
  double tnmin;
  double *yn;

  tau = new double[nmax+1];
  yn = new double[m];

  kappa = 0.85;
  theta = 0.5;

  n_rejected = 0;
  n_fsolve = 0;
//
//  On starting this loop with a given value of n,
//  we have t(0:n-1), y(0:n-1,1:m) and tau(0:n-1).
//  We want to take step n, and compute t(n), y(n,1:m), tau(n).
//
  for ( n = 0; n <= nmax; n++ )
  {
//
//  n = 0: set n data
//
    if ( n == 0 )
    {
      t[n] = t0;
      for ( i = 0; i < m; i++ )
      {
        y[i+n*m] = y0[i];
      }
      tau[n] = tau0;
      count = 0;
    }
//
//  n = 1, 2: y1, y2
//  Steps 1 and 2 use fixed tau.
//
//  n = 3, ..., nmax: use data from step n-1 to compute data for step n.
//  Tau gets updated.
//  Trying a step to t(n) = t(n-1) + theta * tau(n-1).
//  If step fails, may need to reduce theta and try again.
//
    else
    {
      while ( true )
      {
        mad_step ( m, t[n-1], y+(n-1)*m, tau[n-1], f, theta, reltol, 
          tn, yn, fsolve_fail );

        if ( fsolve_fail )
        {
          n_fsolve = n_fsolve + 1;
          continue;
        }

        t[n] = tn;
        for ( i = 0; i < m; i++ )
        {
          y[i+n*m] = yn[i];
        }
        count = count + 1;

        if ( n == 1 || n == 2 )
        {
          tau[n] = tau0;
          break;
        }
//
//  Estimate local truncation error.
// 
        mad_lte ( m, y+(n-3)*m, tau+(n-3), reltol, abstol, 
          tnmin, tnmax );
//
//  If the LTE test was met, we accept T and Y, and set the next time step.
//
        if ( tnmin <= 1.0 )
        {
          factor = kappa * pow ( 1.0 / tnmax, 1.0/3.0 );
          factor = fmin ( factor, 1.5 );
          factor = fmax ( factor, 0.02 );
          tau[n] = factor * tau[n-1];
          break;
        }
        else
        {
          n_rejected = n_rejected + 1;
        }
      }
    }
//
//  Exit if we reached or surpassed final time b.
//
    if ( tmax <= t[n] )
    {
      break;
    }
  }

  nstep = n;
//
//  Free memory.
//
  delete [] tau;
  delete [] yn;

  return;
}
//****************************************************************************80

void mad_lte ( int m, double y[], double tau[], double reltol, double abstol, 
  double &tnmin, double &tnmax )

//****************************************************************************80
//
//  Purpose:
//
//    mad_lte() estimates local truncation error for midpoint adaptive method.
//
//  Discussion:
//
//    This function is called by mad_compute() to estimate the local truncation
//    error incurred during a single step of the implicit midpoint method.
//
//    We assume the step has been taken over an interval which we will 
//    represent as going from (t1,y1) to (t2,y2).  We may regard y2 as the
//    estimated solution of the local ODE:
//      dy/dt = f(t,y), y(t1) = y1.
//
//    We suppose that a function y*() is the exact solution of this system
//    and we define the local truncation error for this step as
//      lte = y*(t2) - y2
//    We wish to accept this step as long as lte is small, that is:
//      lte = y*(t2) - y2 < abstol + reltol * max ( y1, y2 )
//
//    Of course, to carry out this check, we need to estimate lte or y*(t2).
//    This is done using the Milne Device, described in the Milne reference.
//    Essentially, lte is estimated using a weighted difference of solution
//    estimates y2 and an estimate formed by an Adams-Bashforth AB2-like 
//    method.  See page 5 of the Burkardt/Trenchea reference for details.
//
//  Licensing:
//
//    This code is distributed under the MIT license.
//
//  Modified:
//
//    09 July 2024
//
//  Author:
//
//    John Burkardt, Catalin Trenchea
//
//  Reference:
//
//    John Burkardt, Catalin Trenchea,
//    Refactorization of the midpoint rule,
//    Applied Mathematics Letters,
//    Volume 107, 106438, 2020.
//
//    William Milne,
//    Numerical Integration of Ordinary Differential Equations,
//    American Mathematical Monthly,
//    Volume 33, number 9, pages 455–460, 1926.
//
//  Input:
//
//    int m: the spatial dimension of the solution.
//
//    double y[4*m]: the last four solution vectors.
//
//    double tau[3]: the last three stepsizes.
//
//    double reltol, abstol: local truncation error tolerances.
//
//  Output:
//
//    double tnmin, tnmax: the minimum and maximum entries of the
//    local truncation error estimator.
//
{
  double c1;
  double c2;
  double c3;
  int i;
  double r;
  double tn;
  double *uab2;
//
//  Evaluate the AB2-like solution.
//
  c1 = ( tau[2] + tau[1] ) * ( tau[2] + tau[1] + tau[0] ) 
    / ( tau[1] * ( tau[1] + tau[0] ) );

  c2 = - tau[2] * ( tau[2] + tau[1] + tau[0] ) / ( tau[1] * tau[0] );

  c3 = tau[2] * ( tau[2] + tau[1] ) / ( tau[0] * ( tau[1] + tau[0] ) );

  uab2 = new double[m];
  for ( i = 0; i < m; i++ )
  {
    uab2[i] = c1 * y[i+2*m] + c2 * y[i+1*m] + c3 * y[i+0*m];
  }
//
//  Evaluate R, the variable error coefficient in the LTE.
//
  r = 1.0 / 24.0 + 1.0 / 8.0 * ( 1.0 + tau[1] / tau[2] ) 
    * ( 1.0 + 2.0 * tau[1] / tau[2] + tau[0] / tau[2] );
//
//  Use AB2 estimator to approximate the LTE vector.
//
  tnmax = 0;
  tnmin = HUGE_VAL;

  for ( i = 0; i < m; i++ )
  {
    tn = ( y[i+3*m] - uab2[i] ) * 24.0 * r / ( 24.0 * r - 1.0 ) 
      / ( abstol + fabs ( y[i+3*m] ) * reltol );

    tn = fabs ( tn );
    if ( tnmax < tn )
    {
      tnmax = tn;
    }
    if ( tn < tnmin )
    {
      tnmin = tn;
    }
  }

  delete [] uab2;

  return;
}
//****************************************************************************80

void mad_phase_plot ( int n, int m, double t[], double y[], string label )

//****************************************************************************80
//
//  Purpose:
//
//    mad_phase_plot() writes data and command files for a phase plot.
//
//  Discussion:
//
//    We assume a simple phase plane plot of y(0) versus y(1).
//
//  Licensing:
//
//    This code is distributed under the MIT license.
//
//  Modified:
//
//    10 July 2024
//
//  Author:
//
//    John Burkardt
//
//  Input:
//
//    int n: the number of steps.
//
//    int m: the size of a solution vector.
//
//    double t[0:n]: the sequence of times.
//
//    double y[n*m]: N solutions, each of length M.
//
//    string label: a title.
//
{
  string command_filename;
  ofstream command;
  string data_filename;
  ofstream data;
  int i;
  int j;
  string png_filename;
//
//  Create the data file.
//
  data_filename = label + "_phase_data.txt";
  data.open ( data_filename.c_str ( ) );

  for ( j = 0; j < n; j++ )
  {
    data << "  " << t[j];
    for ( i = 0; i < m; i++ )
    {
      data << "  " << y[i+j*m];
    }
    data << "\n";
  }

  data.close ( );

  cout << "\n";
  cout << "  Graphics data stored in '" + data_filename + "'.\n";
//
//  Create the command file.
//
  command_filename = label + "_phase_commands.txt";
  command.open ( command_filename.c_str ( ) );

  png_filename = label + "_phase.png";

  command << "# " + command_filename + "\n";
  command << "#\n";
  command << "# Usage:\n";
  command << "#  gnuplot < " + command_filename + "\n";
  command << "#\n";
  command << "set term png\n";
  command << "set output '" + png_filename + "'\n";
  command << "set xlabel '<-- y[0][t] -->'\n";
  command << "set ylabel '<-- y[1][t] -->'\n";
  command << "set title 'MAD phase plane'\n";
  command << "set grid\n";
  command << "set style data lines\n";
  command << "plot '" + data_filename + "' using 2:3 with lines lw 3\n";
  command << "quit\n";

  command.close ( );

  cout << "  Plot commands stored in '" + command_filename + "'\n"; 

  return;
}
//****************************************************************************80

void mad_solution_plot ( int n, int m, double t[], double y[], string label )

//****************************************************************************80
//
//  Purpose:
//
//    mad_solution_plot() writes data and command files for a solution plot.
//
//  Licensing:
//
//    This code is distributed under the MIT license.
//
//  Modified:
//
//    10 July 2024
//
//  Author:
//
//    John Burkardt
//
//  Input:
//
//    int n: the number of steps.
//
//    int m: the size of a solution vector.
//
//    double t[0:n[: the sequence of time values.
//
//    double y[0:n * 1:m]: the sequence of solutions.
//
//    string label: a title.
//
{
  string command_filename;
  ofstream command;
  string data_filename;
  ofstream data;
  int i;
  int j;
  string png_filename;
//
//  Create the data file.
//
  data_filename = label + "_solution_data.txt";
  data.open ( data_filename.c_str ( ) );

  for ( j = 0; j < n; j++ )
  {
    data << "  " << t[j];
    for ( i = 0; i < m; i++ )
    {
      data << "  " << y[i+j*m];
    }
    data << "\n";
  }

  data.close ( );
  cout << "\n";
  cout << "  Solution data stored in '" + data_filename + "'.\n";
/*
  Create the command file.
*/
  command_filename = label + "_solution_commands.txt";
  png_filename = label + "_solution.png";

  command.open ( command_filename.c_str() );

  command << "# " + command_filename + "\n";
  command << "#\n";
  command << "# Usage:\n";
  command << "#  gnuplot < " + command_filename + "\n";
  command << "#\n";
  command << "set term png\n";
  command << "set output '" + png_filename + "'\n";
  command << "set xlabel '<-- t -->'\n";
  command << "set ylabel '<-- y[t] -->'\n";
  command << "set title 'MAD solution Y(T)'\n";
  command << "set grid\n";
  command << "set style data lines\n";
  command << "plot \\\n";
  for ( i = 2; i <= m; i++ )
  {
    command << "'" + data_filename + "' using 1:" << i << " with lines lw 3,\\\n";
  }
  command << "'" + data_filename + "' using 1:" << m + 1 << " with lines lw 3\n";
  command << "quit\n";

  command.close ( );

  cout << "  Solution plot commands stored in '" + command_filename + "'.\n";

  return;
}
//****************************************************************************80

void mad_step ( int m, double to, double yo[], double tauo, 
  void f ( double t, double y[], double dydt[] ), double theta, double tol, 
  double &tn, double yn[], bool &fsolve_fail )

//****************************************************************************80
//
//  Purpose:
//
//    mad_step() extends the (t,y) solution sequence by one more step.
//
//  Discussion:
//
//    This code is requested to add one more pair of (t,y) data to a string
//    of estimated solutions to the differential equation.
//
//    On input, an initial condition and n solutions have already been 
//    computed and stored in the arrays t and y.  
//
//    A stepsize tau(n) and factor theta are provided.  These define a solution
//    ym of the implicit backward Euler method at tm = t(n) + theta * tau(n). 
//    Then ym is used to construct the implicit midpoint solution y(n+1,1:n) 
//    at time t(n+1) = t(n) + tau(n).
//
//  Licensing:
//
//    This code is distributed under the MIT license.
//
//  Modified:
//
//    29 June 2024
//
//  Author:
//
//    John Burkardt, Catalin Trenchea
//
//  Reference:
//
//    Catalin Trenchea, John Burkardt,
//    Refactorization of the midpoint rule,
//    Applied Mathematics Letters,
//    Volume 107, September 2020.
//
//  Input:
//
//    int m: the size of a single solution vector y.
//
//    double to, yo[m]: the "old" time and solution.
//
//    double tauo: the time step.
//
//    void f ( double t, double y[], double dydt[] ): the function that 
//    evaluates the right hand side of the differential equation.
//
//    double theta: the theta-method parameter.  Common values include:
//    theta = 0.0 for the backward Euler method.
//    theta = 0.5 for the midpoint method.
//    theta = 1.0 for the forward Euler method.
//
//  Output:
//
//    double &tn, yn[m]: time and solution information has been 
//    added as index n+1.
//
//    bool fsolve_fail: true if the backward Euler computation failed.
//
{
  int i;
  int info;
  double *fm;
  double tm;
  double *ym;

  fsolve_fail = false;
//
//  Make a forward solution prediction at to + theta * tau
//
  fm = new double[m];
  ym = new double[m];

  tm = to + theta * tauo;
  f ( tm, yo, fm );
  for ( i = 0; i < m; i++ )
  {
    ym[i] = yo[i] + theta * tauo * fm[i];
  }
//
//  Starting from the forward prediction, 
//  solve the implicit backward Euler equation for (tm, ym).
//
  info = fsolve_be ( f, m, to, yo, tm, ym, fm, tol );

  if ( info != 1 )
  {
    cout << "\n";
    cout << "mad_step(): Fatal error!\n";
    cout << "  fsolve_be() returns info = " << info << "\n";
    fsolve_fail = true;
    return;
  }
//
//  Now predict the solution (tn, yn).
//
  tn = to + tauo;
  for ( i = 0; i < m; i++ )
  {
    yn[i] = (       1.0 / theta ) * ym[i]
          + ( 1.0 - 1.0 / theta ) * yo[i];
  }
//
//  Free memory.
//
  delete [] ym;
  delete [] fm;

  return;
}
//****************************************************************************80

void mad_stats ( int n, int n_rejected, int n_fsolve, double t[] )

//****************************************************************************80
//
//  Purpose:
//
//    mad_stats() reports time and timestep statistics for mad_compute().
//
//  Discussion:
//
//    This function prints some information about the time nodes and time steps.
//
//  Licensing:
//
//    This code is distributed under the MIT license.
//
//  Modified:
//
//    14 July 2024
//
//  Author:
//
//    John Burkardt, Catalin Trenchea
//
//  Reference:
//
//    Catalin Trenchea, John Burkardt,
//    Refactorization of the midpoint rule,
//    Applied Mathematics Letters,
//    Volume 107, September 2020.
//
//  Input:
//
//    int n: the number of time steps.
//
//    int n_rejected: the number of rejected time steps.
//
//    int n_fsolve: the number of failed calls to fsolve.
//
//    double t[n+1]: the sequence of time values.
//
{
  double dt;
  double dt_max;
  double dt_min;
  int i;

  dt_max = 0.0;
  dt_min = HUGE_VAL;

  for ( i = 0; i < n - 1; i++ )
  {
    dt = t[i+1] - t[i];
    if ( dt_max < dt )
    {
      dt_max = dt;
    }
    if ( dt < dt_min )
    {
      dt_min = dt;
    }
  }

  cout << "\n";
  cout << "mad_stats():\n";
  cout << "  Number of timesteps n =   " << n << "\n";
  cout << "  Number of rejected steps  " << n_rejected << "\n";
  cout << "  Number of fsolve failures " << n_fsolve << "\n";
  cout << "  First time interval: [" << t[0] << "," << t[1] << "]\n";
  cout << "  Final time interval: [" << t[n-2] << "," << t[n-1] << "]\n";
  cout << "  Minimum timestep is " << dt_min << "\n";
  cout << "  Maximum timestep is " << dt_max << "\n";

  return;
}
//****************************************************************************80

void mad_timestep_plot ( int n, double t[], string label )

//****************************************************************************80
//
//  Purpose:
//
//    mad_timestep_plot() writes data and command files for a timestep plot.
//
//  Licensing:
//
//    This code is distributed under the MIT license.
//
//  Modified:
//
//    10 July 2024
//
//  Author:
//
//    John Burkardt
//
//  Input:
//
//    int n: the number of steps.
//
//    double t[0:n]: the sequence of time values.
//
//    string label: a title.
//
{
  string command_filename;
  ofstream command;
  string data_filename;
  ofstream data;
  int j;
  string png_filename;
//
//  Create the data file.
//
  data_filename = label + "_timestamp_data.txt";

  data.open ( data_filename.c_str ( ) );

  for ( j = 1; j < n; j++ )
  {
    data << "  " << j - 1
         << "  " << t[j] - t[j-1] << "\n";
  }

  data.close ( );

  cout << "\n";
  cout << "  Time step data stored in '" + data_filename + "'.\n";
//
//  Create the command file.
//
  command_filename = label + "_timestep_commands.txt";
  png_filename = label + "_timestep.png";

  command.open ( command_filename.c_str ( ) );

  command << "# " + command_filename + "\n";
  command << "#\n";
  command << "# Usage:\n";
  command << "#  gnuplot < " + command_filename + "\n";
  command << "#\n";
  command << "set term png\n";
  command << "set output '" + png_filename + "'\n";
  command << "set xlabel '<-- Index -->'\n";
  command << "set ylabel '<-- DT -->'\n";
  command << "set title 'MAD Timesteps'\n";
  command << "set grid\n";
  command << "set style data lines\n";
  command << "plot '" + data_filename + "' using 1:2 with lines lw 3\n";
  command << "quit\n";

  command.close ( );

  cout << "  Timestep plot commands stored in '" + command_filename + "'.\n";

  return;
}