using System;
using HeuristicLab.Problems.DynamicalSystemsModelling;
using Microsoft.VisualStudio.TestTools.UnitTesting;

namespace AutoDiffForDynamicalModelsTest {
  [TestClass]
  public class TestCvodes {
    [TestMethod]
    public unsafe void TestCvodesASAMethod() {
      // test vectors
      var arr = new double[] { 3.14, 2.71 };
      var vec = CVODES.N_VMake_Serial(2, arr);
      CVODES.N_VDestroy_Serial(vec);

      vec = CVODES.N_VNew_Serial(10);

      unsafe {
        int* content = (int*)*(int*)vec.ToPointer();
        long length = *content;
        Console.WriteLine(*content);
        int own_data = *(content + 2);
        Console.WriteLine(own_data);
        double* data = (double*)*(content + 3);
        double data0 = *data;
      }
      CVODES.N_VConst_Serial(2.0, vec);
      CVODES.N_VPrint_Serial(vec);
      Assert.AreEqual(20, CVODES.N_VL1Norm_Serial(vec));
      CVODES.N_VDestroy_Serial(vec);


      // linear oscillator
      int numberOfEquations = 2;
      var y = CVODES.N_VNew_Serial(numberOfEquations);
      // y must be initialized before calling CVodeInit
      // CVODES.N_VConst_Serial(100.0, y);
      CVODES.NV_Set_Ith_S(y, 0, 0.5); // x
      CVODES.NV_Set_Ith_S(y, 1, 1);  // v

      var cvode_mem = CVODES.CVodeCreate(CVODES.MultistepMethod.CV_ADAMS, CVODES.NonlinearSolverIteration.CV_FUNCTIONAL);

      var flag = CVODES.CVodeInit(cvode_mem, F, 0.0, y);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      flag = CVODES.CVodeSStolerances(cvode_mem, 1E-4, 1.0);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      var A = CVODES.SUNDenseMatrix(numberOfEquations, numberOfEquations);
      Assert.AreNotSame(A, IntPtr.Zero);


      var linearSolver = CVODES.SUNDenseLinearSolver(y, A);
      Assert.AreNotSame(linearSolver, IntPtr.Zero);

      flag = CVODES.CVDlsSetLinearSolver(cvode_mem, linearSolver, A);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      // flag = CVODES.CVDlsSetJacFn(cvode_mem, JacF);
      // Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      // var ns = 1; // number of parameters
      // var p = new double[1]; // set as user-data
      // var yS0 = CVODES.N_VCloneVectorArray_Serial(ns, y); // clone the output vector for each parameter, TODO: free
      // for(int i=0;i<ns;i++) {
      // 
      // }


      var q = CVODES.N_VNew_Serial(1);
      CVODES.NV_Set_Ith_S(q, 0, 0.0);
      flag = CVODES.CVodeQuadInit(cvode_mem, FQ, q);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);


      int steps = 10;
      flag = CVODES.CVodeAdjInit(cvode_mem, steps, CVODES.CV_HERMITE);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      /* step by step forward integration 
      double t = 0.0;
      double tout = 0.1; // first output time
      int nout = 100; // number of output times
      for (int iout = 0; iout < nout; iout++) {
        flag = CVODES.CVode(cvode_mem, tout, y, ref t, CVODES.CV_NORMAL);
        Assert.AreEqual(CVODES.CV_SUCCESS, flag);
        Console.WriteLine("{0} {1} {2}", t, CVODES.NV_Get_Ith_S(y, 0), CVODES.NV_Get_Ith_S(y, 1));
        // CVODES.N_VPrint_Serial(y);
        tout += 0.1; 
      }
      */

      // complete forward integration
      double tout = 100.0; // last time
      double time = 0.0;
      int ncheck = 0; // number of checkpoints
      flag = CVODES.CVodeF(cvode_mem, tout, y, ref time, CVODES.CV_NORMAL, ref ncheck);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      long numSteps = 0;
      flag = CVODES.CVodeGetNumSteps(cvode_mem, ref numSteps);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      var yB = CVODES.N_VNew_Serial(numberOfEquations);
      CVODES.N_VConst_Serial(0.0, yB);

      int numberOfParameters = 2;
      var qB = CVODES.N_VNew_Serial(numberOfParameters);
      CVODES.N_VConst_Serial(0.0, qB);

      int indexB = 0;
      flag = CVODES.CVodeCreateB(cvode_mem, CVODES.MultistepMethod.CV_BDF, CVODES.NonlinearSolverIteration.CV_NEWTON, ref indexB);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      var TB1 = tout;
      flag = CVODES.CVodeInitB(cvode_mem, indexB, FB, TB1, yB);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      var relTolB = 1E-6;
      var absTolB = 1E-8;
      flag = CVODES.CVodeSStolerancesB(cvode_mem, indexB, relTolB, absTolB);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      var AB = CVODES.SUNDenseMatrix(numberOfEquations, numberOfEquations);
      Assert.AreNotSame(linearSolver, IntPtr.Zero);

      var lsB = CVODES.SUNDenseLinearSolver(yB, AB);
      Assert.AreNotSame(linearSolver, IntPtr.Zero);

      flag = CVODES.CVDlsSetLinearSolverB(cvode_mem, indexB, lsB, AB);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      flag = CVODES.CVDlsSetJacFnB(cvode_mem, indexB, JacFB);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      flag = CVODES.CVodeQuadInitB(cvode_mem, indexB, FQB, qB);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      /* First get results at t = TBout1 */

      /* Call CVodeB to integrate the backward ODE CVODES. */
      var tBackOut = 50.0;
      flag = CVODES.CVodeB(cvode_mem, tBackOut, CVODES.CV_NORMAL);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      /* Call CVodeGetB to get yB of the backward ODE CVODES. */
      flag = CVODES.CVodeGetB(cvode_mem, indexB, ref time, yB);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      /* Call CVodeGetAdjY to get the interpolated value of the forward solution
         y during a backward integration. */
      flag = CVODES.CVodeGetAdjY(cvode_mem, tBackOut, y);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      /* Then at t = T0 */

      double t0 = 0.0;
      flag = CVODES.CVodeB(cvode_mem, t0, CVODES.CV_NORMAL);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      long nstB = 0;
      CVODES.CVodeGetNumSteps(CVODES.CVodeGetAdjCVodeBmem(cvode_mem, indexB), ref nstB);

      flag = CVODES.CVodeGetB(cvode_mem, indexB, ref time, yB);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);


      flag = CVODES.CVodeGetQuadB(cvode_mem, indexB, ref time, qB);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      flag = CVODES.CVodeGetAdjY(cvode_mem, t0, y);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      CVODES.N_VDestroy_Serial(y);
      CVODES.CVodeFree(ref cvode_mem);
      CVODES.SUNLinSolFree(linearSolver);
      CVODES.SUNMatDestroy(A);
    }

    [TestMethod]
    public void TestCvodesFSAMethod() {

      // linear oscillator
      int numberOfEquations = 2;
      var y = CVODES.N_VNew_Serial(numberOfEquations);
      // y must be initialized before calling CVodeInit
      // CVODES.N_VConst_Serial(100.0, y);
      CVODES.NV_Set_Ith_S(y, 0, 0.5); // x
      CVODES.NV_Set_Ith_S(y, 1, 1);  // v

      var cvode_mem = CVODES.CVodeCreate(CVODES.MultistepMethod.CV_BDF, CVODES.NonlinearSolverIteration.CV_NEWTON);

      var flag = CVODES.CVodeInit(cvode_mem, F, 0.0, y);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      flag = CVODES.CVodeSStolerances(cvode_mem, 1E-4, 1.0);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      var A = CVODES.SUNDenseMatrix(numberOfEquations, numberOfEquations);
      Assert.AreNotSame(A, IntPtr.Zero);
      Console.WriteLine(CVODES.SUNDenseMatrix_Get(A, 0, 0));


      var linearSolver = CVODES.SUNDenseLinearSolver(y, A);
      Assert.AreNotSame(linearSolver, IntPtr.Zero);

      flag = CVODES.CVDlsSetLinearSolver(cvode_mem, linearSolver, A);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      flag = CVODES.CVDlsSetJacFn(cvode_mem, JacF);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      var ns = 1; // number of parameters
      var p = new double[1]; // set as user-data
      var yS0 = CVODES.N_VCloneVectorArray_Serial(ns, y); // clone the output vector for each parameter, TODO: free
      unsafe {
        for (int i = 0; i < ns; i++) {
          var yS0_i = *((IntPtr*)yS0.ToPointer() + i);
          CVODES.N_VConst_Serial(0.0, yS0_i);
        }
      }


      flag = CVODES.CVodeSensInit(cvode_mem, ns, CVODES.CV_SIMULTANEOUS, FS, yS0);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      // flag = CVODES.CVodeSensSStolerances(cvode_mem, reltolS, abstolS);
      // Assert.AreEqual(CVODES.CV_SUCCESS, flag);
      flag = CVODES.CVodeSensEEtolerances(cvode_mem);
      Assert.AreEqual(CVODES.CV_SUCCESS, flag);


      // var q = CVODES.N_VNew_Serial(1);
      // CVODES.NV_Set_Ith_S(q, 0, 0.0);
      // flag = CVODES.CVodeQuadInit(cvode_mem, FQ, q);
      // Assert.AreEqual(CVODES.CV_SUCCESS, flag);


      // int steps = 10;
      // flag = CVODES.CVodeAdjInit(cvode_mem, steps, CVODES.CV_HERMITE);
      // Assert.AreEqual(CVODES.CV_SUCCESS, flag);

      // step by step forward integration 
      double t = 0.0;
      double tout = 0.1; // first output time
      int nout = 1000; // number of output times
      for (int iout = 0; iout < nout; iout++) {
        flag = CVODES.CVode(cvode_mem, tout, y, ref t, CVODES.CV_NORMAL);
        Assert.AreEqual(CVODES.CV_SUCCESS, flag);

        // get sensitivities
        flag = CVODES.CVodeGetSens(cvode_mem, ref t, yS0);
        Assert.AreEqual(CVODES.CV_SUCCESS, flag);
        unsafe {
          var ySP0 = *(IntPtr*)yS0.ToPointer(); // parameter of sensitivities for first component
          Console.WriteLine("{0} {1} {2} {3} {4}", t, CVODES.NV_Get_Ith_S(y, 0), CVODES.NV_Get_Ith_S(y, 1), CVODES.NV_Get_Ith_S(ySP0, 0), CVODES.NV_Get_Ith_S(ySP0, 1));
        }

        tout += 0.1;
      }



      CVODES.N_VDestroy_Serial(y);
      CVODES.CVodeFree(ref cvode_mem);
      CVODES.SUNLinSolFree(linearSolver);
      CVODES.SUNMatDestroy(A);
    }



    public static int F(
      double t, // realtype
      IntPtr y, // N_Vector
      IntPtr ydot, // N_Vector
      IntPtr user_data) {

      // the system
      // ∂y1/∂t = y2
      // ∂y2/∂t = -0.3 y1
      CVODES.NV_Set_Ith_S(ydot, 0, CVODES.NV_Get_Ith_S(y, 1));
      CVODES.NV_Set_Ith_S(ydot, 1, -0.3 * CVODES.NV_Get_Ith_S(y, 0));
      return 0;
      ;
    }

    // computes ∂f/∂y or an approximation
    private int JacF(double t, IntPtr y, IntPtr fy, IntPtr Jac, IntPtr user_data, IntPtr tmp1, IntPtr tmp2, IntPtr tmp3) {
      // t is current time
      // y is current y
      // fy is current f(t,y)
      // Jac is output

      // Matrix:
      //  ∂f1/∂y1  ∂f1/∂y2
      //  ∂f2/∂y1  ∂f2/∂y2

      CVODES.SUNDenseMatrix_Set(Jac, 0, 0, 0.0);
      CVODES.SUNDenseMatrix_Set(Jac, 0, 1, 1.0);
      CVODES.SUNDenseMatrix_Set(Jac, 1, 0, -0.3);
      CVODES.SUNDenseMatrix_Set(Jac, 1, 1, 0.0);
      return 0;
    }

    private int JacFB(double t, IntPtr y, IntPtr yB, IntPtr fyB, IntPtr Jac, IntPtr user_data, IntPtr tmp1, IntPtr tmp2, IntPtr tmp3) {
      throw new NotImplementedException();
      return 0;
    }

    // sensitivity function must compute the vector (∂f/∂y)s_i(t) + ∂f/∂p_i and store in ySdot[i]
    public static int FS(
      int Ns,
      double t,
      IntPtr y, // N_Vector
      IntPtr ydot, // N_Vector
      IntPtr yS, // N_Vector* one vector for each parameter
      IntPtr ySdot, // N_Vector* one vector for each parameter
      IntPtr user_data,
      IntPtr tmp1, // N_Vector
      IntPtr tmp2 // N_Vector
      ) {

      // (∂f /∂y)s_i(t) + ∂f /∂p_i
      var y1 = CVODES.NV_Get_Ith_S(y, 0); var y2 = CVODES.NV_Get_Ith_S(y, 1);

      unsafe {
        var yS_p0 = *(IntPtr*)yS.ToPointer();
        var s1 = CVODES.NV_Get_Ith_S(yS_p0, 0); var s2 = CVODES.NV_Get_Ith_S(yS_p0, 1);


        // (∂f /∂y)s_i(t)  == Jac s_i(t)
        var sd1 = s2;
        var sd2 = -0.3 * s1;

        var s_p0 = *(IntPtr*)ySdot.ToPointer();
        CVODES.NV_Set_Ith_S(s_p0, 0, sd1 + 0.0);
        CVODES.NV_Set_Ith_S(s_p0, 1, sd2 - 0.3);
      }

      return 0;
    }

    public static int FB(
        double t, // realtype
        IntPtr y, // N_Vector
        IntPtr yB, // N_Vector
        IntPtr yBdot, // N_Vector
        IntPtr user_data) {

      // y' = f(y,p)

      // yB = λ
      // λ' = -(∂f / ∂y)^T λ - (∂g / ∂y)
      // the goal is to find dG/dp where G = integrate( g(y,t,p), t=0, t=T) )

      // for F above:
      // ∂f / ∂ = 
      //  0.0 1.0
      // -0.3 0.0

      // we have no g!?
      // therefore λ' = 
      // λ1' = 0.3 λ2 - 0.0
      // λ2' = 1.0 λ1 - 0.0

      CVODES.NV_Set_Ith_S(yBdot, 0, 0.3 * CVODES.NV_Get_Ith_S(yB, 1));
      CVODES.NV_Set_Ith_S(yBdot, 1, CVODES.NV_Get_Ith_S(yB, 0));

      return 0;
      ;
    }

    private int FQ(double t, IntPtr y, IntPtr yQdot, IntPtr user_data) {
      // TODO: squared error
      CVODES.NV_Set_Ith_S(yQdot, 0, CVODES.NV_Get_Ith_S(y, 2));
      return 0;
    }


    private int FQB(double t, IntPtr y, IntPtr yB, IntPtr qBdot, IntPtr user_data) {
      throw new NotImplementedException();
    }
  }
}