#include <iostream>
#include <fstream>
#include <string>

#include <petsc.h>
#include "simulation.h"
#include "vars.h"

int iteration=0;

// Collect all the values from the parallel Ybus into a sequential vector ybus[ngen*ngen] on each processor
static PetscErrorCode scatterYbus(Mat fy, PetscScalar *ybus, PetscInt ngen)
{
  PetscErrorCode ierr;
  PetscInt i, j;
  PetscInt rA, cA;
  ierr = MatGetSize(fy, &rA, &cA); CHKERRQ(ierr);
  Vec ACol[cA];
  PetscScalar *AColVal[cA];

  for (i = 0; i < cA; i++) {
    ierr = VecCreateMPI(PETSC_COMM_WORLD, PETSC_DECIDE, rA, &ACol[i]); CHKERRQ(ierr);
    ierr = MatGetColumnVector(fy, ACol[i], i); CHKERRQ(ierr);
    VecScatter ctx;
    Vec localACol;
    ierr = VecScatterCreateToAll(ACol[i], &ctx, &localACol); CHKERRQ(ierr);
    ierr = VecScatterBegin(ctx, ACol[i], localACol, INSERT_VALUES, SCATTER_FORWARD); CHKERRQ(ierr);
    ierr = VecScatterEnd(ctx, ACol[i], localACol, INSERT_VALUES, SCATTER_FORWARD); CHKERRQ(ierr);
    ierr = VecGetArray(localACol, &AColVal[i]); CHKERRQ(ierr);
    ierr = VecScatterDestroy(&ctx); CHKERRQ(ierr);
    ierr = VecRestoreArray(localACol, NULL); CHKERRQ(ierr);
  }
  for (i = 0; i < rA; i++) {
    for (j = 0; j < cA; j++) {
      ybus[i*ngen+j] = AColVal[i][j];
      //printf("%lf+%lfi\n", PetscRealPart( AColVal[i][j]), PetscImaginaryPart( AColVal[i][j]));
    }
  }
  //printf("\n");

  PetscFunctionReturn(0);
}

PetscErrorCode Simulation::IFunction(TS ts,PetscReal t,Vec X,Vec Xdot,Vec F,Userctx *ctx)
{
  PetscLogStagePush(ctx->fstage);
  PetscErrorCode ierr;
  PetscScalar *x, *xdot, *f;
  PetscInt i, j, k;
  PetscInt ngen = ctx->ngen;
  PetscScalar *ybus = ctx->ybus;
  PetscScalar *eqprime= ctx->eqprime;
  PetscScalar *pmech = ctx->pmech;
  PetscScalar *gen_d0 = ctx->gen_d0;
  PetscScalar *gen_h = ctx->gen_h;
  VecScatter vs = ctx->vs;

  //for (i = 0; i < ngen*ngen; i++)
    //printf("%lf+%lfi\n", PetscRealPart(ybus[i]), PetscImaginaryPart(ybus[i]));

  DM da;
  PetscInt xstart, xlen;

  int me;

  double t0, t1;

  PetscFunctionBeginUser;

  MPI_Comm_rank(PETSC_COMM_WORLD, &me);

  //printf("IFunction, me = %d:\n", me);

  ierr = TSGetDM(ts, &da); CHKERRQ(ierr);

  // Get pointers to vector data
  //ierr = DMDAVecGetArray(da, X, &x); CHKERRQ(ierr);
  //ierr = DMDAVecGetArray(da, Xdot, &xdot); CHKERRQ(ierr);
  t0 = MPI_Wtime();

  ierr = VecScatterBegin(ctx->vs, X, ctx->x, INSERT_VALUES, SCATTER_FORWARD); CHKERRQ(ierr);
  ierr = VecScatterEnd(ctx->vs, X, ctx->x, INSERT_VALUES, SCATTER_FORWARD); CHKERRQ(ierr);
  ierr = VecGetArray(ctx->x, &x); CHKERRQ(ierr);

  ierr = VecScatterBegin(ctx->vs, Xdot, ctx->xdot, INSERT_VALUES, SCATTER_FORWARD); CHKERRQ(ierr);
  ierr = VecScatterEnd(ctx->vs, Xdot, ctx->xdot, INSERT_VALUES, SCATTER_FORWARD); CHKERRQ(ierr);
  ierr = VecGetArray(ctx->xdot, &xdot); CHKERRQ(ierr);
  //t1 = MPI_Wtime();
  //if (me == 0) std::cout << "Scattering IFunction time: " << t1 - t0 << std::endl;

  ierr = DMDAVecGetArray(da, F, &f); CHKERRQ(ierr);
  //VecView(X, PETSC_VIEWER_STDOUT_WORLD);
  //VecView(Xdot, PETSC_VIEWER_STDOUT_WORLD);

  // Get local grid boundaries
  ierr = DMDAGetCorners(da, &xstart, NULL, NULL, &xlen, NULL, NULL); CHKERRQ(ierr);
  //printf("Function Called   ");
  //printf("IFunction:me, xstart, xstart+xlen=%d: [%d, %d]\n", me, xstart, xstart+xlen);

  //t0 = MPI_Wtime();

  // Set up DAE equations
  // Compute function over the locally owned part of the grid
  for (i = 0; i < ngen; i++) {
    if (2*i >= xstart && 2*i < xstart+xlen) {
      f[2*i] = xdot[2*i] - basrad * x[2*i+1];
      //std::cout << "IFunction:me,f[2*i]:" << me << ":[" << xstart << "," << xstart+xlen << "] " <<  "f[" << 2*i << "]=" << f[2*i] << std::endl;
    }
    if (2*i+1 >= xstart && 2*i+1 < xstart+xlen) {
      f[2*i+1] = xdot[2*i+1] - 0.5/gen_h[i]*(
                 pmech[i]-eqprime[i]*x[2*ngen+2*i+1]-gen_d0[i]*x[2*i+1]);
      //std::cout << "IFunction:me,f[2*i+1]:" << me << ":[" << xstart << "," << xstart+xlen << "] " <<  "f[" << 2*i+1 << "]=" << f[2*i+1] << std::endl;
      //std::cout <<"IFunction:x[2*ngen+2*i+1]:"<<x[2*ngen+2*i+1] << std::endl;
    }
  }
  for (i = 0; i < ngen; i++) {
    if (2*ngen+2*i >= xstart && 2*ngen+2*i < xstart+xlen) {
      f[2*ngen+2*i] = sin(x[2*i])*x[2*ngen+2*i] + cos(x[2*i])*x[2*ngen+2*i+1];
      for (k = 0; k < ngen; k++) {
        f[2*ngen+2*i] += -((cos(x[2*k])*eqprime[k])*PetscRealPart(ybus[k*ngen+i])
                         -(sin(x[2*k])*eqprime[k])*PetscImaginaryPart(ybus[k*ngen+i]));
      }
    }
    if (2*ngen+2*i+1 >= xstart && 2*ngen+2*i+1 < xstart+xlen) {
      f[2*ngen+2*i+1] = -cos(x[2*i])*x[2*ngen+2*i] + sin(x[2*i])*x[2*ngen+2*i+1];
      for (k = 0; k < ngen; k++) {
        f[2*ngen+2*i+1] += -((cos(x[2*k])*eqprime[k])*PetscImaginaryPart(ybus[k*ngen+i])
                           +(sin(x[2*k])*eqprime[k])*PetscRealPart(ybus[k*ngen+i]));
      }
    }
  }
  //VecView(F, PETSC_VIEWER_STDOUT_WORLD);

  //t1 = MPI_Wtime();
  //if (me == 0) std::cout << "Scattering IFunction time: " << t1 - t0 << std::endl;

  //for (i=0;i<4*ngen;i++) printf("me, xstart, xstart+xlen, PetscRealPart(f[i]), PetscImaginaryPart(f[i]):%d:[%d,%d] %lf+%lfi\n", me, xstart, xstart+xlen, PetscRealPart(f[i]), PetscImaginaryPart(f[i]));
  //for (i=0;i<4*ngen;i++) printf("me, xstart, xstart+xlen, PetscRealPart(pmech[i]), PetscImaginaryPart(pmech[i]):%d:[%d,%d] %lf+%lfi\n", me, xstart, xstart+xlen, PetscRealPart(pmech[i]), PetscImaginaryPart(pmech[i]));
  //printf("\n");

  ierr = VecRestoreArray(ctx->x, &x); CHKERRQ(ierr);
  ierr = VecRestoreArray(ctx->xdot, &xdot); CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, F, &f); CHKERRQ(ierr);

  PetscLogStagePop();
  PetscFunctionReturn(0);
}

//PetscErrorCode Simulation::IJacobian(TS ts, PetscReal t, Vec X, Vec Xdot, PetscReal a, Mat *A, Mat *B, MatStructure *flag, Userctx *ctx)
PetscErrorCode Simulation::IJacobian(TS ts, PetscReal t, Vec X, Vec Xdot, PetscReal a, Mat A, Mat B, Userctx *ctx) // version 3.5 change: all "*A" change to "A"
{
  PetscLogStage stage = ctx->jstage1;
  PetscLogStagePush(stage);

  VecScatter vs = ctx->vs;

  PetscErrorCode ierr;
  PetscInt ngen = ctx->ngen;
  PetscScalar *x;

  PetscInt i, j, icol;
  PetscScalar *ybus = ctx->ybus;
  PetscScalar *eqprime= ctx->eqprime;
  PetscScalar *pmech = ctx->pmech;
  PetscScalar *gen_d0 = ctx->gen_d0;
  PetscScalar *gen_h = ctx->gen_h;

  //for (i = 0; i < ngen*ngen; i++)
    //printf("%lf+%lfi\n", PetscRealPart(ybus[i]), PetscImaginaryPart(ybus[i]));

  // DM da;
  PetscInt xstart, xend , m, n;

  int me;
  double t0, t1, t2, t3, t4, t5, t6, t7;

  PetscFunctionBeginUser;

  //printf("   IJacobian:\n");

  // ierr = TSGetDM(ts, &da); CHKERRQ(ierr);

  // Get pointers to vector data
  //ierr = DMDAVecGetArray(da, X, &x); CHKERRQ(ierr);
  MPI_Comm_rank(PETSC_COMM_WORLD, &me);
  //t0 = MPI_Wtime();

  ierr = VecScatterBegin(ctx->vs, X, ctx->x, INSERT_VALUES, SCATTER_FORWARD); CHKERRQ(ierr);
  ierr = VecScatterEnd(ctx->vs, X, ctx->x, INSERT_VALUES, SCATTER_FORWARD); CHKERRQ(ierr);
  ierr = VecGetArray(ctx->x, &x); CHKERRQ(ierr);
  //t1 = MPI_Wtime();
  //if (me == 0) std::cout << "Scattering IJacobian time: " << t1 - t0 << std::endl;

  // Get local grid boundaries
  // ierr = DMDAGetCorners(da, &xstart, NULL, NULL, &xlen, NULL, NULL); CHKERRQ(ierr);


  ierr = MatGetOwnershipRange(A, &xstart, &xend); CHKERRQ(ierr);

  ////////////////////////////////////////////////////////////////////////////////////////
  if (iteration > 0) {
    ierr = MatRetrieveValues(A); CHKERRQ(ierr);
  }

  PetscScalar xVal, ybusVal;
  PetscScalar val[1];
  PetscInt ii;
  //printf("Jacobian Called  \n");
  t0 = MPI_Wtime();
  // Compute function over the locally owned part of the grid
  for (i = xstart; i < xend; i++) {
    //printf("IJacobian:i:%d\n", i);
    //if (xstart+xlen <= 2*ngen) { // the first two 2*ngen rows of Jacobain matrix
    if (i < 2*ngen) { // the first two 2*ngen rows of Jacobain matrix
    // only calculate the frist two 2*ngen in the first IJacobian iteration;
    // otherwise, just reuse the stored data calculated from the first iteration
      if (iteration == 0) {
        if (i % 2 == 0) { // row 0,2,4,...,2*ngen-2
          icol = i;
          val[0] = a;
          ierr = MatSetValues(A, 1, &i, 1, &icol, val, INSERT_VALUES); CHKERRQ(ierr);
          icol = i+1;
          val[0] = -basrad;
          ierr = MatSetValues(A, 1, &i, 1, &icol, val, INSERT_VALUES); CHKERRQ(ierr);
        } else { // row 1,3,5,...,2*ngen-1
          ii = (i-1)/2;
          icol = i;
          val[0] = a + gen_d0[ii]/(2.0*gen_h[ii]);
          ierr = MatSetValues(A, 1, &i, 1, &icol, val, INSERT_VALUES); CHKERRQ(ierr);
          icol = 2*ngen + i;
          val[0] = eqprime[ii]/(2.0*gen_h[ii]);
          ierr = MatSetValues(A, 1, &i, 1, &icol, val, INSERT_VALUES); CHKERRQ(ierr);
        }
      }
      t2 =  MPI_Wtime();
    } else { // the last two 2*ngen rows of Jacobian matrix
      if (i % 2 == 0) { // row 2*ngen,2*ngen+2,...,4*ngen-2
        ii = (i-2*ngen)/2; // ii=0,1,2...
        xVal = x[2*ii];
        for (j = 0; j < ngen; j++) {
          icol = j*2;
          ybusVal = ybus[j*ngen+ii];
          if (i == j*2+ngen*2) { // diaganol Jgx(ii)
            val[0] = x[i] * cos(xVal)
                   - x[i+1] * sin(xVal)
                   + PetscRealPart(ybusVal)*eqprime[ii]*sin(xVal)
                   + PetscImaginaryPart(ybusVal)*eqprime[ii]*cos(xVal);
          } else { // off-diaganol Jgx(ik)
            val[0] = PetscRealPart(ybusVal)*eqprime[j]*sin(x[2*j])
                   + PetscImaginaryPart(ybusVal)*eqprime[j]*cos(x[2*j]);
          }
          ierr = MatSetValues(A, 1, &i, 1, &icol, val, INSERT_VALUES); CHKERRQ(ierr);
        }
        icol = i;
        val[0] = sin(xVal);
        ierr = MatSetValues(A, 1, &i, 1, &icol, val, INSERT_VALUES); CHKERRQ(ierr);
        icol = i+1;
        val[0] = cos(xVal);
        ierr = MatSetValues(A, 1, &i, 1, &icol, val, INSERT_VALUES); CHKERRQ(ierr);

      } else { // row 2*ngen+1,2*ngen+3,...,4*ngen-1
        ii = (i-2*ngen-1)/2; // ii=0,1,2...
        xVal = x[2*ii];
        for (j = 0; j < ngen; j++) {
          icol = j*2;
          ybusVal = ybus[j*ngen+ii];
          if ((i-1) == j*2+ngen*2) { // diaganol Jgx(ii)
            val[0] = x[i-1] * sin(xVal)
                   + x[i] * cos(xVal)
                   + PetscImaginaryPart(ybusVal)*eqprime[ii]*sin(xVal)
                  - PetscRealPart(ybusVal)*eqprime[ii]*cos(xVal);
          } else {// off-diaganol Jgx(ik)
            val[0] = PetscImaginaryPart(ybusVal)*eqprime[j]*sin(x[2*j])
                   - PetscRealPart(ybusVal)*eqprime[j]*cos(x[2*j]);
          }
          ierr = MatSetValues(A, 1, &i, 1, &icol, val, INSERT_VALUES); CHKERRQ(ierr);
        }
        icol = i-1;
        val[0] = -cos(xVal);
        ierr = MatSetValues(A, 1, &i, 1, &icol, val, INSERT_VALUES); CHKERRQ(ierr);
        icol = i;
        val[0] = sin(xVal);
        ierr = MatSetValues(A, 1, &i, 1, &icol, val, INSERT_VALUES); CHKERRQ(ierr);
      }
    }
    t3 = MPI_Wtime();
  }

  t1 = MPI_Wtime();
 // if (me == 0) std::cout << "Scattering IJacobian time: " << t1 - t0 <<"," << t2 - t0<<"," << t3 - t2<<std::endl;


  /*ierr = MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);
  ierr = MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);
  MatView(A, PETSC_VIEWER_STDOUT_WORLD);
  exit(1);*/

  ierr = VecRestoreArray(ctx->x, &x); CHKERRQ(ierr);

  PetscLogStagePush(ctx->jstage2);

  ierr = MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);
  ierr = MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);
  //MatView(A, PETSC_VIEWER_STDOUT_WORLD);
  // *flag = SAME_NONZERO_PATTERN; // version 3.5 change

  PetscLogStagePop();

  if (iteration == 0) {
    ierr = MatSetOption(A,MAT_NEW_NONZERO_LOCATIONS,PETSC_FALSE); CHKERRQ(ierr);
    ierr = MatStoreValues(A); CHKERRQ(ierr);
  }

  iteration ++;
  //t1 = MPI_Wtime();
  //if (me == 0) std::cout << "Scattering IJacobian time: " << t1 - t0 << std::endl;

  PetscLogStagePop();
  PetscFunctionReturn(0);
}

PetscErrorCode Simulation::formInitialSolution(TS ts, Vec U, Userctx *ctx, PetscInt *t_step, PetscReal *t_width)
{
  PetscInt i;
  PetscScalar *psw1, *psw7;
  PetscScalar *pbus_a, *pbus_v, *pbus_pg, *pbus_qg;
  PetscScalar *pgen_mva, *pgen_d0, *pgen_h, *pgen_dtr;
  PetscScalar *theta = new PetscScalar[nbus];
  PetscErrorCode ierr;
  PetscInt flagF;
  PetscReal H_sum;
  PetscInt tst1;
  PetscScalar curr, phi, v, eprime;

  PetscScalar eterm, pelect, qelect, curq;
  PetscScalar mac_ang, mac_spd, dmac_ang, dmac_spd;
  PetscScalar cur_re, cur_im, psi_re, psi_im;
  PetscInt simu_k;

  DM da;
  PetscInt xstart, xlen, Mx;
  PetscScalar *u;
  PetscReal hx, x, r;

  PetscScalar *eqprime= ctx->eqprime;
  PetscScalar *pmech = ctx->pmech;

  PetscFunctionBeginUser;

  ierr = TSGetDM(ts, &da); CHKERRQ(ierr);

  // Get pointers to vector data
  ierr = DMDAVecGetArray(da, U, &u); CHKERRQ(ierr);

  // Get local grid boundaries
  ierr = DMDAGetCorners(da, &xstart, NULL, NULL, &xlen, NULL, NULL); CHKERRQ(ierr);

  simu_k = 0;

  ierr = VecGetArray(sw1, &psw1); CHKERRQ(ierr);
  ierr = VecGetArray(sw7, &psw7); CHKERRQ(ierr);

  ierr = VecGetArray(bus_a, &pbus_a); CHKERRQ(ierr);
  ierr = VecGetArray(bus_v, &pbus_v); CHKERRQ(ierr);
  ierr = VecGetArray(bus_pg, &pbus_pg); CHKERRQ(ierr);
  ierr = VecGetArray(bus_qg, &pbus_qg); CHKERRQ(ierr);

  ierr = VecGetArray(gen_mva, &pgen_mva); CHKERRQ(ierr);
  ierr = VecGetArray(gen_d0, &pgen_d0); CHKERRQ(ierr);
  ierr = VecGetArray(gen_h, &pgen_h); CHKERRQ(ierr);
  ierr = VecGetArray(gen_dtr, &pgen_dtr); CHKERRQ(ierr);

  for (i = 0; i < nswtch-1; i++) {
    t_step[i] = (PetscInt) PetscRealPart((psw1[i+1] - psw1[i]) / psw7[i]);
    t_width[i] = PetscRealPart((psw1[i+1] - psw1[i]) / (PetscScalar) t_step[i]);
    simu_k += t_step[i];
    //std::cout << i+1 << " " << t_step[i] << " " << t_width[i] << " " << simu_k << std::endl;
  }
  ierr = VecRestoreArray(sw1, &psw1); CHKERRQ(ierr);
  ierr = VecRestoreArray(sw1, &psw7); CHKERRQ(ierr);

  simu_k ++;
  //allocSimuData(); // Allocate memories for simulation related arrays

  for (i = 0; i < nbus; i++) {
    theta[i] = pbus_a[i] * M_PI / 180.0;
    //printf("theta[%d]: %f+%fi\n", i, PetscRealPart(theta[i]), PetscImaginaryPart(theta[i]));
  }

  // Calculate parameters
  flagF = 0;
  H_sum = 0.0;
  for (i = 0; i < ngen; i++) {
    pgen_mva[i] = basmva / pgen_mva[i];
    ctx->gen_d0[i] = pgen_d0[i] / pgen_mva[i];
    ctx->gen_h[i] = pgen_h[i] / pgen_mva[i];
    //std::cout << ctx->gen_d0[i] << ", " <<  ctx->gen_h[i] << std::endl;
    tst1 = gen_bus[i];
    eterm = pbus_v[tst1];
    pelect = pbus_pg[tst1];
    qelect = pbus_qg[tst1];
    curr = sqrt(pelect * pelect + qelect * qelect) / eterm * pgen_mva[i];
    phi = atan2(PetscRealPart(qelect), PetscRealPart(pelect));
    v = eterm * exp(PETSC_i * theta[tst1]);
    curr = curr * exp(PETSC_i * (theta[tst1] - phi));
    eprime = v + PETSC_i * pgen_dtr[i] * curr;
    mac_ang = atan2(PetscImaginaryPart(eprime), PetscRealPart(eprime));
    mac_spd = 0;
    eqprime[i] = abs(eprime);
    pmech[i] = pelect;
    //std::cout << std::endl <<  "eterm: " << eterm  << std::endl;
    //std::cout << "curr: " << curr  << std::endl;
    //std::cout << "phi: " << phi  << std::endl;
    //std::cout << "v: " << v  << std::endl;
    //std::cout << "eprime: " << eprime  << std::endl;
    //std::cout << "mac_ang: " << mac_ang  << std::endl;
    //std::cout << "mac_spd: " << mac_spd  << std::endl;
    if (i*2 >= xstart && i*2 < xstart+xlen)
      u[i*2] = mac_ang;
    //std::cout <<"formInitialSolution: i*2,u[i*2]:" << i*2 << ": " << u[i*2] << std::endl;
    if (i*2+1 >= xstart && i*2+1 < xstart+xlen)
      u[i*2+1] = mac_spd;
    //std::cout << "formInitialSolution: i*2+1,u[i*2+1]:" <<i*2+1 << ": " << u[i*2+1] << std::endl;
    if (2*ngen+i*2 >= xstart && 2*ngen+i*2 < xstart+xlen)
      u[2*ngen+i*2] = sin(mac_ang)*PetscRealPart(curr)
                      - cos(mac_ang)*PetscImaginaryPart(curr);
    //std::cout << "formInitialSolution: 2*ngen+i*2,u[2*ngen+i*2]:" << 2*ngen+i*2 << ": " << u[2*ngen+i*2] << std::endl;
    if (2*ngen+i*2+1 >= xstart && 2*ngen+i*2+1 < xstart+xlen)
      u[2*ngen+i*2+1] = cos(mac_ang)*PetscRealPart(curr)
                        + sin(mac_ang)*PetscImaginaryPart(curr);
    //std::cout << "formInitialSolution: 2*ngen+i*2+1,u[2*ngen+i*2+1]:"<< 2*ngen+i*2+1 << ": " << u[2*ngen+i*2+1] << std::endl;
  }

  //for (i = 0; i < 4*ngen; i++)
  //  for (i = xstart; i < xstart+xlen; i++)
      //std::cout << "formInitialSolution: i, u[i]:"<< i << ": " << u[i] << std::endl;

  delete [] theta;

  ierr = VecRestoreArray(bus_a, &pbus_a); CHKERRQ(ierr);
  ierr = VecRestoreArray(bus_v, &pbus_v); CHKERRQ(ierr);
  ierr = VecRestoreArray(bus_pg, &pbus_pg); CHKERRQ(ierr);
  ierr = VecRestoreArray(bus_qg, &pbus_qg); CHKERRQ(ierr);

  ierr = VecRestoreArray(gen_mva, &pgen_mva); CHKERRQ(ierr);
  ierr = VecRestoreArray(gen_d0, &pgen_d0); CHKERRQ(ierr);
  ierr = VecRestoreArray(gen_h, &pgen_h); CHKERRQ(ierr);
  ierr = VecRestoreArray(gen_dtr, &pgen_dtr); CHKERRQ(ierr);

  // Set Userctx values: eqprime, pmech, gen_d0, gen_h

  // Restore vectors
  ierr = DMDAVecRestoreArray(da, U, &u); CHKERRQ(ierr);
  //VecView(U, PETSC_VIEWER_STDOUT_WORLD);

  PetscFunctionReturn(0);
}

PetscErrorCode Simulation::simu(Mat prefy11,Mat fy11,Mat posfy11)
{
  PetscErrorCode ierr;

  // Simulation related variables
  PetscInt i;
  PetscInt t_step[20];
  PetscReal t_width[20];

  // Petsc DAE related variables
  TS ts; // Nonlinear solver
  Vec x; // Solution, residual vectors
  Mat J; // Jacobian matrix
  PetscInt steps;
  PetscReal ftime;
  Userctx user;

  DM da;

  double t0, t1, tt0, tt1, tt2, tt3, tt4;

  PetscFunctionBegin;

  tt0 = MPI_Wtime(); //

  //---------------------------------------------------------------------
  // view prefy11, fy11, posfy11 values
  //---------------------------------------------------------------------
  //MatView(prefy11, PETSC_VIEWER_STDOUT_WORLD);
  //MatView(fy11, PETSC_VIEWER_STDOUT_WORLD);
  //MatView(posfy11, PETSC_VIEWER_STDOUT_WORLD);

  //---------------------------------------------------------------------
  // Start of Petsc DAE solver main driver
  //---------------------------------------------------------------------

  //---------------------------------------------------------------------
  // Create distributed array (DMDA) to manage parallel grid and vectors
  //---------------------------------------------------------------------
  ierr = DMDACreate1d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, 4*ngen, 1, 1, NULL, &da);
  CHKERRQ(ierr);

  //---------------------------------------------------------------------
  // Extract global vectors from DMDA;
  //---------------------------------------------------------------------
  ierr = DMCreateGlobalVector(da, &x); CHKERRQ(ierr);
  //VecView(x, PETSC_VIEWER_STDOUT_SELF);

  //---------------------------------------------------------------------
  // Initialize user application context
  //---------------------------------------------------------------------
  user.ngen = ngen;
  user.ybus = new PetscScalar[ngen*ngen];
  user.eqprime = new PetscScalar[ngen];
  user.pmech = new PetscScalar[ngen];
  user.gen_d0 = new PetscScalar[ngen];
  user.gen_h = new PetscScalar[ngen];

  PetscLogStageRegister("IFunction",&user.fstage);
  PetscLogStageRegister("IJacobian",&user.jstage1);
  PetscLogStageRegister("IJacobian assemble",&user.jstage2);

  ////////////////////////////////////////////////////////////////////////
  // 1. prefy11 DAE:
  ////////////////////////////////////////////////////////////////////////

  // Set Userctx values: corresponding Ybus (first with prefy11)
  // Collect all the values from the parallel Ybus into a sequential vector ybus[ngen*ngen] on each processor
  t0 = MPI_Wtime();
  scatterYbus(prefy11, user.ybus, ngen);
  t1 = MPI_Wtime();
  if (me == 0) std::cout << "Scattering prefy11 time: " << t1 - t0 << std::endl;

  //---------------------------------------------------------------------
  // Create timestepping solver context
  //---------------------------------------------------------------------
  //TSCreate(PETSC_COMM_WORLD,&ts);
  //TSSetProblemType(ts,TS_NONLINEAR);
  //TSSetType(ts,TSROSW);
  //TSSetType(ts,TSARKIMEX);
  t0 = MPI_Wtime();
  ierr = TSCreate(PETSC_COMM_WORLD, &ts); CHKERRQ(ierr);
  //ierr = TSSetType(ts, TSBEULER); CHKERRQ(ierr);
  ierr = TSSetType(ts, TSTHETA); CHKERRQ(ierr);
  //ierr = TSThetaSetTheta(ts, 1.0); CHKERRQ(ierr);
  ierr = TSThetaSetTheta(ts, 0.5); CHKERRQ(ierr);
  //ierr = TSSetType(ts, TSEULER); CHKERRQ(ierr);
  ierr = TSSetIFunction(ts, NULL, (TSIFunction) IFunction, &user); CHKERRQ(ierr);
  /*ierr = MatCreate(PETSC_COMM_WORLD, &J); CHKERRQ(ierr); // J: Jacobian matrix
  ierr = MatSetSizes(J, PETSC_DECIDE, PETSC_DECIDE, 4*ngen, 4*ngen); CHKERRQ(ierr);
  ierr = MatSetFromOptions(J); CHKERRQ(ierr);
  ierr = MatSetUp(J); CHKERRQ(ierr);*/

  ierr = MatCreate(PETSC_COMM_WORLD, &J); CHKERRQ(ierr); // J: Jacobian matrix
  ierr = MatSetType(J, MATMPIAIJ); CHKERRQ(ierr);
  ierr = MatSetSizes(J, PETSC_DECIDE, PETSC_DECIDE, 4*ngen, 4*ngen); CHKERRQ(ierr);
  ierr = MatSetFromOptions(J); CHKERRQ(ierr);


  /* PetscInt nnz[4*ngen];
  for (i = 0; i < 4*ngen; i++) {
     if (i < 2*ngen) {
        nnz[i] = 2;
     }else{
        nnz[i] = ngen+2;
     }
  }
  ierr = MatSeqAIJSetPreallocation(J,ngen+2,nnz);  CHKERRQ(ierr); */

  ierr = MatMPIAIJSetPreallocation(J,ngen+2,NULL,ngen+2,NULL);  CHKERRQ(ierr);

  //ierr = MatSetUp(J); CHKERRQ(ierr);

  //ierr = DMCreateMatrix(da, MATAIJ, &J); CHKERRQ(ierr);
  ierr = TSSetIJacobian(ts, J, J, (TSIJacobian) IJacobian, &user); CHKERRQ(ierr);

  // Use TSGetDM() to access. Setting here allows easy use of geometric multigrid.
  ierr = TSSetDM(ts, da); CHKERRQ(ierr);
  t1 = MPI_Wtime();
  if (me == 0) std::cout << "timestepping solver context time: " << t1 - t0 << std::endl;

  //---------------------------------------------------------------------
  // Set initial conditions
  //---------------------------------------------------------------------
  t0 = MPI_Wtime();
  ierr = formInitialSolution(ts, x, &user, t_step, t_width); CHKERRQ(ierr);
  ftime = t_step[0] * t_width[0];
  //printf("ftime = %lf\n", ftime);
  ierr = TSSetDuration(ts, PETSC_DEFAULT, ftime); CHKERRQ(ierr);
  ierr = TSSetSolution(ts, x); CHKERRQ(ierr);
  ierr = TSSetInitialTimeStep(ts, 0.0, t_width[0]); CHKERRQ(ierr);

  //---------------------------------------------------------------------
  // Use initial solution vector to set up vector scatter
  //---------------------------------------------------------------------
  ierr = VecScatterCreateToAll(x, &(user.vs), &(user.x)); CHKERRQ(ierr);
  ierr = VecDuplicate(user.x, &(user.xdot)); CHKERRQ(ierr);

  //---------------------------------------------------------------------
  // Set runtime options
  //---------------------------------------------------------------------
  ierr = TSSetFromOptions(ts);CHKERRQ(ierr);
  t1 = MPI_Wtime();
  if (me == 0) std::cout << "Set initial conditions time: " << t1 - t0 << std::endl;

  //---------------------------------------------------------------------
  // Solve nonlinear system
  //---------------------------------------------------------------------
  //ierr = VecView(x,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
  t0 = MPI_Wtime();
  ierr = TSSolve(ts,x);CHKERRQ(ierr);
  //  ierr = TSView(ts,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
  t1 = MPI_Wtime();
  if (me == 0) std::cout << "TSSolve time: " << t1 - t0 << std::endl;
  ierr = TSGetSolveTime(ts,&ftime);CHKERRQ(ierr);
  ierr = TSGetTimeStepNumber(ts,&steps);CHKERRQ(ierr);
  ierr = PetscPrintf(PETSC_COMM_WORLD,"%D steps, ftime %f\n",steps,ftime);CHKERRQ(ierr);
  // version 3.5 change: %G format is no longer supported, use %g and caste the argument to double,all change to %f
  //ierr = VecView(x,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
  t1 = MPI_Wtime();
  if (me == 0) std::cout << "Solve nonlinear system time: " << t1 - t0 << std::endl;


  tt2 = MPI_Wtime(); //
  if (me == 0) std::cout << "Total prefy11 time: " << tt2 - tt0 << std::endl;


  ////////////////////////////////////////////////////////////////////////
  // 2. fy11 DAE:
  ////////////////////////////////////////////////////////////////////////
  // Set Userctx values: corresponding Ybus (second with fy11)
  t0 = MPI_Wtime();
  scatterYbus(fy11, user.ybus, ngen);
  t1 = MPI_Wtime();
  if (me == 0) std::cout << "Scattering fy11 time: " << t1 - t0 << std::endl;
  ierr = TSCreate(PETSC_COMM_WORLD, &ts); CHKERRQ(ierr);
  //ierr = TSSetType(ts, TSBEULER); CHKERRQ(ierr);
  ierr = TSSetType(ts, TSTHETA); CHKERRQ(ierr);
  //ierr = TSThetaSetTheta(ts, 1.0); CHKERRQ(ierr);
  ierr = TSThetaSetTheta(ts, 0.5); CHKERRQ(ierr);
  //ierr = TSSetType(ts, TSEULER); CHKERRQ(ierr);
  ierr = TSSetIFunction(ts, NULL, (TSIFunction) IFunction, &user); CHKERRQ(ierr);
  ierr = TSSetIJacobian(ts, J, J, (TSIJacobian) IJacobian, &user); CHKERRQ(ierr);
  ierr = TSSetDM(ts, da); CHKERRQ(ierr);
  //ierr = formInitialSolution(ts, x, &user, t_step, t_width); CHKERRQ(ierr);
  ftime = t_step[1] * t_width[1];
  ierr = TSSetDuration(ts, PETSC_DEFAULT, ftime); CHKERRQ(ierr);
  ierr = TSSetInitialTimeStep(ts, 0.0, t_width[1]); CHKERRQ(ierr);
  ierr = TSSetFromOptions(ts);CHKERRQ(ierr);
  ierr = TSSolve(ts,x);CHKERRQ(ierr);
  //ierr = TSView(ts,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
  ierr = TSGetSolveTime(ts,&ftime);CHKERRQ(ierr);
  ierr = TSGetTimeStepNumber(ts,&steps);CHKERRQ(ierr);
  ierr = PetscPrintf(PETSC_COMM_WORLD,"%D steps, ftime %f\n",steps,ftime);CHKERRQ(ierr);
  //ierr = VecView(x,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);

  tt3 = MPI_Wtime(); //
  if (me == 0) std::cout << "Total fy11 time: " << tt3 - tt2 << std::endl;



  ////////////////////////////////////////////////////////////////////////
  // 2. posfy11 DAE:
  ////////////////////////////////////////////////////////////////////////
  // Set Userctx values: corresponding Ybus (third with posfy11)
  t0 = MPI_Wtime();
  scatterYbus(posfy11, user.ybus, ngen);
  t1 = MPI_Wtime();
  if (me == 0) std::cout << "Scattering posfy11 time: " << t1 - t0 << std::endl;
  ierr = TSCreate(PETSC_COMM_WORLD, &ts); CHKERRQ(ierr);
  //ierr = TSSetType(ts, TSBEULER); CHKERRQ(ierr);
  ierr = TSSetType(ts, TSTHETA); CHKERRQ(ierr);
  //ierr = TSThetaSetTheta(ts, 1.0); CHKERRQ(ierr);
  ierr = TSThetaSetTheta(ts, 0.5); CHKERRQ(ierr);
  //ierr = TSSetType(ts, TSEULER); CHKERRQ(ierr);
  ierr = TSSetIFunction(ts, NULL, (TSIFunction) IFunction, &user); CHKERRQ(ierr);
  ierr = TSSetIJacobian(ts, J, J, (TSIJacobian) IJacobian, &user); CHKERRQ(ierr);
  ierr = TSSetDM(ts, da); CHKERRQ(ierr);
  //ierr = formInitialSolution(ts, x, &user, t_step, t_width); CHKERRQ(ierr);
  ftime = t_step[2] * t_width[2];
  ierr = TSSetDuration(ts, PETSC_DEFAULT, ftime); CHKERRQ(ierr);
  ierr = TSSetInitialTimeStep(ts, 0.0, t_width[2]); CHKERRQ(ierr);
  ierr = TSSetFromOptions(ts);CHKERRQ(ierr);
  t0 = MPI_Wtime();
  ierr = TSSolve(ts,x);CHKERRQ(ierr);
  //  ierr = TSView(ts,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
  t1 = MPI_Wtime();
  if (me == 0) std::cout << "TSSolve posfy11 time: " << t1 - t0 << std::endl;
  ierr = TSGetSolveTime(ts,&ftime);CHKERRQ(ierr);
  ierr = TSGetTimeStepNumber(ts,&steps);CHKERRQ(ierr);
  ierr = PetscPrintf(PETSC_COMM_WORLD,"%D steps, ftime %f\n",steps,ftime);CHKERRQ(ierr);
  //ierr = VecView(x,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);

  tt4 = MPI_Wtime(); //
  if (me == 0) std::cout << "Total posfy11 time: " << tt4 - tt3 << std::endl;

  PetscViewer viewer;
  PetscViewerASCIIOpen(PETSC_COMM_WORLD, "u.dat", &viewer);
  PetscViewerSetFormat(viewer,PETSC_VIEWER_DEFAULT);
  VecView(x,viewer);
  PetscViewerDestroy(&viewer);

  //---------------------------------------------------------------------
  // Free work space.  All PETSc objects should be destroyed when they
  // are no longer needed.
  //---------------------------------------------------------------------
  ierr = VecScatterDestroy(&(user.vs));
  ierr = VecDestroy(&(user.x));
  ierr = VecDestroy(&(user.xdot));
  ierr = MatDestroy(&J); CHKERRQ(ierr);
  ierr = VecDestroy(&x); CHKERRQ(ierr);
  ierr = TSDestroy(&ts); CHKERRQ(ierr);
  ierr = DMDestroy(&da); CHKERRQ(ierr);

  delete [] user.ybus;
  delete [] user.eqprime;
  delete [] user.pmech;
  delete [] user.gen_d0;
  delete [] user.gen_h;

  tt1 = MPI_Wtime(); //
  if (me == 0) std::cout << "Total simu time: " << tt1 - tt0 << std::endl;

  PetscFunctionReturn(0);
}