package com.dwave.math;

 /**
  * Runge-Kutta methods propagate a solution over an interval by combining the
  * information from several Euler-sytle steps(each involving one evaluation of
  * the right-hand <i>f</i>'s), and then using the information obtained to match
  * a Taylor series expansion up to some higher order.
  * <p>
  * Translated to Java directly from Numerical Recipes in C.
  */

public final class RungeKutta {
    private final static float
           a2  = 0.2f,
           a3  = 0.3f,
           a4  = 0.6f,
           a5  = 1.0f,
           a6  = 0.875f,
           b21 = 0.2f,
           b31 = 0.075f,
           b32 = 0.225f,
           b41 = 0.3f,
           b42 = -0.9f,
           b43 = 1.2f,
           b51 = -11.0f / 54.0f,
           b52 = 2.5f,
           b53 = -70.0f / 27.0f,
           b54 = 35.0f / 27.0f,
           b61 = 1631.0f / 55296.0f,
           b62 = 175.0f / 512.0f,
           b63 = 575.0f / 13824.0f,
           b64 = 44275.0f / 110592.0f,
           b65 = 253.0f / 4096.0f,
           c1  = 37.0f / 378.0f,
           c3  = 250.0f / 621.0f,
           c4  = 125.0f / 594.0f,
           c6  = 512.0f / 1771.0f,
           dc5 = -277.0f / 14336.0f;

    static long totalTime = 0L;

  /**
   *  Private constructor to guard against instantiation
   */
  private RungeKutta() {}


               //////////////////////////////
              ///// ALGORITHM ROUTINES /////
             //////////////////////////////

 /**
  * Given values for the variables y[1...n] and their derivatives dydk[1..n] known at x, use
  * the fourth-order Runge-Kutta method to advance the solution over an interval h and
  * return the incremented variables as yout[1..n], which need not be a distinct array from y.
  * The user supplies the routine dervis(x,y,dydx), which returns derivatives dydx at x.
  *
  * @param y      psi
  * @param dydx   derivative first order terms
  * @param x      time
  * @param h      stepsize
  * @param yout   output
  * @param derivs method for calculating the derivative
  */
	public static boolean rk4(float y[],float dydx[],float x,
                              float h,float yout[],ODEInterface derivs) {

    if(y.length != dydx.length || y.length != yout.length) {
			System.out.println("RK.rk4: error");
			System.out.println("RK.rk4: y.length " + y.length);
			System.out.println("RK.rk4: dydx.length " + dydx.length);
			System.out.println("RK.rk4: yout.length " + yout.length);
			return false;
		}

		int i;
		float xh, hh, h6, dym[], dyt[], yt[];

		dym = new float[y.length];
		dyt = new float[y.length];
		yt = new float[y.length];
		hh = 0.5f * h;
		h6 = h / 6.0f;
		xh = x + hh;

		for (i=0;i<y.length;i++)
      yt[i] = y[i] + hh * dydx[i];
		derivs.deriv(xh,yt,dyt);
		for (i=0;i<y.length;i++)
      yt[i] = y[i] + hh * dyt[i];
		derivs.deriv(xh,yt,dym);
		for (i=0;i<y.length;i++) {
			yt[i] = y[i] + h * dym[i];
			dym[i] += dyt[i];
		}
		derivs.deriv(x + h,yt,dyt);
		for (i=0;i<y.length;i++)
      yout[i] = y[i] + h6 * (dydx[i] + dyt[i] + 2.0f * dym[i]);
		return true;
	}  

  /**
   *  Given values for y and their derivitives dydx known at x, use the fifth-order
   *  Cash-Karp Runge-Kutta method to advance the solution over an interval h and return
   *  the incremented variables as yout.  Also return an estimate of the local
   *  truncation error in yout using the embedded fourth order method.  The user
   *  supplies the routine derivs wshich returns derivatives dydx at x.
   *
   *  @param  y       psi
   *  @param  dydx    derivative first order terms
   *  @param  x       time
   *  @param  h       time step
   *  @param  yout    results
   *  @param  yerr    estimated error
   *  @param  derivs  method for calculating the derivative
   */
  public static void rkck(float[] y, float[] dydx, float x, float h, float[] yout,
                            float[] yerr, ODEInterface derivs) {
    float  dc1 = c1 - 2825.0f / 27648.0f,
            dc3 = c3 - 18575.0f / 48384.0f,
            dc4 = c4 - 13525.0f / 55296.0f,
            dc6 = c6 - 0.25f;
    float[] ak2, ak3, ak4, ak5, ak6, ytemp;

    ak2   = new float[y.length];
    ak3   = new float[y.length];
    ak4   = new float[y.length];
    ak5   = new float[y.length];
    ak6   = new float[y.length];
    ytemp = new float[y.length];

    for(int i = 0; i < y.length; i++)
      ytemp[i] = y[i] + b21 * h * dydx[i];
    derivs.deriv(x+a2 * h, ytemp, ak2);

    for(int i = 0; i < y.length; i++)
      ytemp[i] = y[i] + h * (b31 * dydx[i] + b32 * ak2[i]);
    derivs.deriv(x+a3*h, ytemp, ak3);

    for(int i = 0; i < y.length; i++)
      ytemp[i] = y[i] + h * (b41 * dydx[i] + b42 * ak2[i] + b43 * ak3[i]);
    derivs.deriv(x+a4*h, ytemp, ak4);

    for(int i = 0; i < y.length; i++)
      ytemp[i] = y[i] + h * (b51 * dydx[i] + b52 * ak2[i] + b53 * ak3[i] + b54 * ak4[i]);
    derivs.deriv(x+a5*h, ytemp, ak5);

    for(int i = 0; i < y.length; i++)
      ytemp[i] = y[i] + h * (b61 * dydx[i] + b62 * ak2[i] + b63 * ak3[i] + b64 * ak4[i] + b65 * ak5[i]);
    derivs.deriv(x+a6*h, ytemp, ak6);

    for(int i = 0; i < y.length; i++)
      yout[i] = y[i] + h * (c1 * dydx[i] + c3 * ak3[i] + c4 * ak4[i] + c6 * ak6[i]);
    for(int i = 0; i < y.length; i++)
      yerr[i] = h * (dc1 * dydx[i] + dc3 * ak3[i] + dc4 * ak4[i] + dc5 * ak5[i] + dc6* ak6[i] );
  }


               ////////////////////
              ///// STEPPERS /////
             ////////////////////

 /**
  * Fifth-order Runge-Kutta step with monitoring of local truncation error to ensure
  * accuracy and adjust stepsize. Input are the dependent variable float[] y[1..n]
  * and it's derivative dydx[1..n] at the starting value of the independent variable x.
  * Also input are the stepsize to be attempted htry, the required accuracy eps, and the
  * float[] yscal[1..n] against which the error is scaled. On output, y and x are
  * replaced by their new values, hdid is the stepsize that was actually accomplished, and
  * hnext is the estimated next stepsize. derivs is the user-supplied routine that computes
  * the right-hand side derivatives.
  *
  * @param y        psi
  * @param dydx     derivative first order terms
  * @param x        time
  * @param htry     stepsize to be attempted
  * @param eps      acccuracy
  * @param yscal    error scale
  * @param hdid     stepsize accomplished
  * @param hnext    next estimated stepsize
  * @param derivs   method for calculating the derivative
  */
  public static boolean rkqs(float[] y, float[] dydx, float[] x, float htry, float eps,
                              float[] yscal, float[] hdid, float[] hnext, ODEInterface derivs) {
		float PGROW = -0.20f;
		float PSHRINK = -0.25f;
		float SAFETY = 0.9f;
		float ERRCON = 1.89e-4f;

    float errmax, h, htemp, xnew;
    float[] yerr, ytemp;

    yerr = new float[y.length];
    ytemp = new float[y.length];
    h = htry;

    for(;;) {
      rkck(y, dydx, x[0], h, ytemp, yerr, derivs);
      errmax = 0.0f;
      for(int i = 0; i < y.length; i++)
        errmax = Math.max(errmax, Math.abs(yerr[i]/yscal[i]));
      errmax /= eps;
//      System.out.println("RKQS pass number "+num++);
//      System.out.println("RKQS time = "+x[0]);
      if(errmax <= 1.0f)
        break;
      htemp = SAFETY * h * (float)Math.pow(errmax, PSHRINK);
      h = (h >= 0.0f ? (float)Math.max(htemp, 0.1 * h) : (float)Math.min(htemp, 0.1 * h));
      xnew = x[0] + h;

      if (xnew == x[0]) {
        System.out.println("Stepsize underflow at rkqs");
        return false;
      }
    }
    if(errmax > ERRCON)
      hnext[0] = SAFETY * h *(float) Math.pow(errmax, PGROW);
    else
      hnext[0] = 5 * h;
    x[0] += (hdid[0] = h);
    for(int i = 0; i < y.length; i++)
      y[i] = ytemp[i];
    return true;
  }      

 /**
  * Older fifth-order Runge-Kutta step with monitoring of local truncation error to
  * ensure accuracy and adjust stepsize. Input are the dependent variable float[] y[1..n]
  * and it's derivative dydx[1..n] at the starting value of the independent variable x.
  * Also input are the stepsize to be attempted htry, the required accuracy eps, and the
  * float[] yscal[1..n] against which the error is scaled. On output, y and x are
  * replaced by their new values, hdid is the stepsize that was actually accomplished, and
  * hnext is the estimated next stepsize. derivs is the user-supplied routine that computes
  * the right-hand side derivatives.
  *
  * @param y        psi
  * @param dydx     derivative first order terms
  * @param x        time
  * @param htry     stepsize to be attempted
  * @param eps      acccuracy
  * @param yscal    error scale
  * @param hdid     stepsize accomplished
  * @param hnext    next estimated stepsize
  * @param derivs   method for calculating the derivative
  */
	public static boolean rkqc(float y[],float dydx[],float x[],float htry,float eps,
		                          float yscal[],float hdid[],float hnext[],ODEInterface derivs) {

		int i;
		float xsav, hh, h, temp, errmax;
		float[] dysav, ysav, ytemp;
		float PGROW = -0.20f;
		float PSHRNK = -0.25f;
		float FCOR = 0.06666666f;
		float SAFETY = 0.9f;
		float ERRCON = 6.0e-4f;

		dysav = new float[y.length];
		ysav = new float[y.length];
		ytemp = new float[y.length];
		xsav = x[0];
		for (i=0;i<y.length;i++) {
			ysav[i] = y[i];
			dysav[i] = dydx[i];
		}
		h = htry;
		while (true) {
			hh = 0.5f * h;
			rk4(ysav,dysav,xsav,hh,ytemp,derivs);
			x[0] = xsav + hh;
			derivs.deriv(x[0],ytemp,dydx);
			rk4(ytemp,dydx,x[0],hh,ytemp,derivs);
			x[0] = xsav + h;
			if (x[0] == xsav) {
				System.out.println("h too small in RK.rkqc()");
				return false;
			}
			rk4(ysav,dysav,xsav,h,ytemp,derivs);
			errmax = 0.0f;

			for (i=0;i<y.length;i++) {
				ytemp[i] = y[i] - ytemp[i];
				temp = Math.abs(ytemp[i] / yscal[i]);

				if (errmax < temp)
          errmax = temp;
			}
			errmax /= eps;

			if (errmax <= 1.0f) {
				hdid[0] = h;
				hnext[0] = (errmax > ERRCON ? SAFETY * h * (float)Math.exp(PGROW * (float)Math.log(errmax)) : 4.0f * h);
				break;
			}
			h = SAFETY * h *(float) Math.exp(PSHRNK * (float)Math.log(errmax));
		}
		for (i=0;i<y.length;i++)
      y[i] += ytemp[i] * FCOR;
		return true;
	}

               ///////////////////
              ///// DRIVERS /////
             ///////////////////

 /**
  * Starting from intial values ystart[1..nvar] known at x1 use fourth-order Runge-Kutta
  * to advance nstep equal increments to x2.  The user-supplied routine derivs(x, y, dydx)
  * evaluates derivatives.  Results are stored in the global variables y[1..nvar][1..nstep+1]
  * and xx[1..nstep+1].
  *
  * @param ystart   initial values
  * @param x1       start time
  * @param x2       end time
  * @param nstep    number of steps to take
  * @param derivs   method for calculating the derivative
  */
  public static boolean rkdumb(float[] ystart, float x1, float x2,
                                int nstep, ODEInterface derivs) {
    int nvar;
    float x, h;
    float[] ytemp, yout, dy, xx;
    float[][] y;

    nvar = ystart.length;
    ytemp = new float[nvar];
    yout = new float[nvar];
    dy = new float[nvar];
    y = new float[nvar][nstep + 1];
    xx = new float[nstep + 1];

    for(int i = 0; i < nvar; i++) {
      ytemp[i] = ystart[i];
      y[i][1] = ytemp[i];
    }
    xx[0] = x1;
    x = x1;
    h = (x2 - x1) / nstep;

    for(int i = 0; i < nstep; i++) {
      derivs.deriv(x, ytemp, dy);
      rk4(ytemp, dy, x, h, yout, derivs);
      if((float) (x+h) == x) {
        System.out.println("Step size too small in routine rkdumb");
        return false;
      }
      x += h;
      xx[i+1] = x;
      for(int j = 0; j < nvar; j++) {
        ytemp[j] = yout[j];
        y[j][i+1] = ytemp[j];
      }
    }
    return true;
  }

 /**
  * Runge-Kutta driver with adaptive stepsize control. Integrate starting values ystart[1..nvar] from
  * x1 to x2 with accuracy eps, storing intermediate results in global varialbles. h1 should be set
  * as a guessed first stepsize, hmin as the minimum allowed stepsize(can be zero).  On output, nok
  * and nbad are the number of good and bad(but retried and fixed) steps taken, and ystart is replaced
  * by values at the end of the integration interval.  derivs is the user-supplied routine for
  * calculating the right-hand side derivative, while rkqc is the name of the stepper  routine
  * to be used.<br><br>
  * [NOTE] User storage for intermediate results: preset kmax (can be 0) and dxsav in the calling
  * program. if kmax != 0 results are stored at approximate intervals dxsav in the array xp[1..kount],
  * yp[1..ystart.length][1..kount], where kount is output by odeint.  Defining declarations for
  * these variables should be in the calling program.
  *
  * @param ystart     psi
  * @param x1         initial time
  * @param x2         finishing time
  * @param eps        accuracy
  * @param h1         guessed first step size
  * @param hmin       minimum stepsize allowed
  * @param nok        number of good steps
  * @param nbad       number of bad steps
  * @param kmax       number of step to be stored
  * @param dxsav      approximate storage interval
  * @param derivs     mehtod for calculating the derivative
  */
	public static boolean odeint(float[] ystart,float x1,float x2,float eps,
		                             float h1,float hmin,int[] nok,int[] nbad,
                                   int kmax, float dxsav, ODEInterface derivs) {
    float xsav, h;
    float[] yscal, y, dydx, xp, x, hnext, hdid;

    int MAXSTEP = 10000,
        KMAXX = 200,
        kount = 0;
    float TINY = 1.0e-30f;

    if(kmax == 0)
      kmax = KMAXX;

    x = new float[1];
    hnext = new float[1];
    hdid = new float[1];
    yscal = new float[ystart.length];
    y = new float[ystart.length];
    dydx = new float[ystart.length];
    xp = new float[kmax];

    x[0] = x1;
		h = (x2 > x1) ? Math.abs(h1) : -Math.abs(h1);
    nok[0] = nbad[0] = 0;

    for(int i = 0; i < ystart.length; i++)
      y[i] = ystart[i];

//    if(kmax > 0)
      xsav = x[0] - dxsav * 2.0f;
//    else
//      xsav = 0;

    for(int nstep = 0; nstep < MAXSTEP; nstep++) {
      derivs.deriv(x[0], y, dydx);
      for(int i = 0; i < ystart.length; i++)
        yscal[i] = Math.abs(y[i]) + Math.abs(dydx[i] * h) + TINY;
      if(kmax > 0 && kount < kmax - 1 && Math.abs(x[0] - xsav) > Math.abs(dxsav)) {
        xp[kount++] = x[0];
        xsav = x[0];
      }
      if((x[0] + h - x2) * (x[0] + h - x1) > 0.0f)
        h = x2 - x[0];
      rkqs(y, dydx, x, h, eps, yscal, hdid, hnext, derivs);

      if(hdid[0] == h)
        nok[0]++;
      else
        nbad[0]++;
      if((x[0] - x2) * (x2 - x1) >= 0.0f) {
        for(int i = 0; i < ystart.length; i++)
          ystart[i] = y[i];
        if(kmax > 0) {
          xp[kount++] = x[0];
        }
        return true;
      }
      if(Math.abs(hnext[0]) <= hmin) {
        System.out.println("Step size too small in odeint");
        return false;
      }
      h = hnext[0];
    }
    System.out.println("Too many steps in routine odeint");
    return false;
	}
}