#include "domain.hpp"

using namespace std;

Domain::Domain(const char *file)
: _file(file),_delim(" = "),_inputDir("unspecified_"),_outputDir(" "),
  _bulkDeformationType("linearElastic"),
  _momentumBalanceType("quasidynamic"),
  _sbpType("mfc_coordTrans"),_operatorType("matrix-based"),
  _sbpCompatibilityType("fullyCompatible"),
  _gridSpacingType("variableGridSpacing"),
  _isMMS(0),_loadICs(0),
  _order(4),_Ny(-1),_Nz(-1),_Ly(-1),_Lz(-1),
  _vL(1e-9),
  _q(NULL),_r(NULL),_y(NULL),_z(NULL),_y0(NULL),_z0(NULL),_dq(-1),_dr(-1),
  _bCoordTrans(-1)
{
  // load data from file
  loadData(_file);

  // check domain size and set grid spacing in y direction
  if (_Ny > 1) {
    _dq = 1.0 / (_Ny - 1.0);
  }
  else {
    _dq = 1;
  }

  // set grid spacing in z-direction
  if (_Nz > 1) {
    _dr = 1.0 / (_Nz - 1.0);
  }
  else {
    _dr = 1;
  }

  checkInput(); // perform some basic value checking to prevent NaNs
  setFields();
  setScatters();
}

// destructor
Domain::~Domain()
{
  // free memory
  VecDestroy(&_q);
  VecDestroy(&_r);
  VecDestroy(&_y);
  VecDestroy(&_z);
  VecDestroy(&_y0);
  VecDestroy(&_z0);

  // set map iterator, free memory from VecScatter
  map<string,VecScatter>::iterator it;
  for (it = _scatters.begin(); it != _scatters.end(); it++) {
    VecScatterDestroy(&(it->second));
  }
}

// construct coordinate transform, setting vectors q, r, y, z
PetscErrorCode Domain::setFields()
{
  PetscErrorCode ierr = 0;
  #if VERBOSE > 1
    string funcName = "Domain::setFields";
    ierr = PetscPrintf(PETSC_COMM_WORLD,"Starting %s in %s.\n",funcName.c_str(),FILENAME);
    CHKERRQ(ierr);
  #endif

  // generate vector _y with size _Ny*_Nz
  ierr = VecCreate(PETSC_COMM_WORLD,&_y); CHKERRQ(ierr);
  ierr = VecSetSizes(_y,PETSC_DECIDE,_Ny*_Nz); CHKERRQ(ierr);
  ierr = VecSetFromOptions(_y); CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject) _y, "y"); CHKERRQ(ierr);

  // duplicate _y into _z, _q, _r
  ierr = VecDuplicate(_y,&_z); CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject) _z, "z"); CHKERRQ(ierr);
  ierr = VecDuplicate(_y,&_q); CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject) _q, "q"); CHKERRQ(ierr);
  ierr = VecDuplicate(_y,&_r); CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject) _r, "r"); CHKERRQ(ierr);

  // construct coordinate transform
  PetscInt Ii,Istart,Iend,Jj = 0;
  PetscScalar *y,*z,*q,*r;
  ierr = VecGetOwnershipRange(_q,&Istart,&Iend);CHKERRQ(ierr);

  // return pointers to local data arrays (the processor's portion of vector data)
  ierr = VecGetArray(_y,&y); CHKERRQ(ierr);
  ierr = VecGetArray(_z,&z); CHKERRQ(ierr);
  ierr = VecGetArray(_q,&q); CHKERRQ(ierr);
  ierr = VecGetArray(_r,&r); CHKERRQ(ierr);

  // set vector entries for q, r (coordinate transform) and y, z (no transform)
  for (Ii=Istart; Ii<Iend; Ii++) {
    q[Jj] = _dq*(Ii/_Nz);
    r[Jj] = _dr*(Ii-_Nz*(Ii/_Nz));

    // matrix-based, fully compatible, allows curvilinear coordinate transformation
    if (_sbpType.compare("mfc_coordTrans") ) {
      y[Jj] = (_dq*_Ly)*(Ii/_Nz);
      z[Jj] = (_dr*_Lz)*(Ii-_Nz*(Ii/_Nz));
    }
    else {
      // hardcoded transformation (not available for z)
      if (_bCoordTrans > 0) {
	y[Jj] = _Ly * sinh(_bCoordTrans * q[Jj]) / sinh(_bCoordTrans);
      }
      // no transformation
      y[Jj] = q[Jj]*_Ly;
      z[Jj] = r[Jj]*_Lz;
    }
    Jj++;
  }

  // restore arrays
  ierr = VecRestoreArray(_y,&y); CHKERRQ(ierr);
  ierr = VecRestoreArray(_z,&z); CHKERRQ(ierr);
  ierr = VecRestoreArray(_q,&q); CHKERRQ(ierr);
  ierr = VecRestoreArray(_r,&r); CHKERRQ(ierr);

  return ierr;
}


// scatters values from one vector to another
// used to get slip on the fault from the displacement vector, i.e., slip = u(1:Nz); shear stress on the fault from the stress vector sxy; surface displacement; surface heat flux
PetscErrorCode Domain::setScatters()
{
  PetscErrorCode ierr = 0;

  // set _y0 to be zero vector with length _Nz
  ierr = VecCreate(PETSC_COMM_WORLD,&_y0); CHKERRQ(ierr);
  ierr = VecSetSizes(_y0,PETSC_DECIDE,_Nz); CHKERRQ(ierr);
  ierr = VecSetFromOptions(_y0); CHKERRQ(ierr);
  ierr = VecSet(_y0,0.0); CHKERRQ(ierr);

  // set _z0 to be zero vector with length _Ny
  ierr = VecCreate(PETSC_COMM_WORLD,&_z0); CHKERRQ(ierr);
  ierr = VecSetSizes(_z0,PETSC_DECIDE,_Ny); CHKERRQ(ierr);
  ierr = VecSetFromOptions(_z0); CHKERRQ(ierr);
  ierr = VecSet(_z0,0.0); CHKERRQ(ierr);

  // set up scatter context to take values for y = 0 from body field and put them on a Vec of size Nz
  PetscInt *indices;
  IS is;  // index set
  ierr = PetscMalloc1(_Nz,&indices); CHKERRQ(ierr);

  // we want to scatter from index 0 to _Nz - 1, i.e. take the first _Nz components of the vector to scatter from
  for (PetscInt Ii = 0; Ii<_Nz; Ii++) {
    indices[Ii] = Ii;
  }

  // creates data structure for an index set containing a list of integers
  ierr = ISCreateGeneral(PETSC_COMM_WORLD, _Nz, indices, PETSC_COPY_VALUES, &is); CHKERRQ(ierr);

  // creates vector scatter context, scatters values from _y (at indices is) to _y0 (at indices is)
  ierr = VecScatterCreate(_y, is, _y0, is, &_scatters["body2L"]); CHKERRQ(ierr);

  // free memory
  ierr = PetscFree(indices); CHKERRQ(ierr);
  ierr = ISDestroy(&is); CHKERRQ(ierr);

  //===============================================================================
  // set up scatter context to take values for y = Ly from body field and put them on a Vec of size Nz
  // indices to scatter from
  PetscInt *fi;
  IS isf;
  ierr = PetscMalloc1(_Nz,&fi); CHKERRQ(ierr);

  // we want to scatter from index _Ny*_Nz - _Nz to _Ny*_Nz - 1, i.e. the last _Nz entries of the vector to scatter from
  for (PetscInt Ii = 0; Ii<_Nz; Ii++) {
    fi[Ii] = Ii + (_Ny*_Nz-_Nz);
  }
  ierr = ISCreateGeneral(PETSC_COMM_WORLD, _Nz, fi, PETSC_COPY_VALUES, &isf); CHKERRQ(ierr);

  // indices to scatter to
  PetscInt *ti;
  IS ist;
  ierr = PetscMalloc1(_Nz,&ti); CHKERRQ(ierr);
  for (PetscInt Ii = 0; Ii<_Nz; Ii++) {
    ti[Ii] = Ii;
  }
  ierr = ISCreateGeneral(PETSC_COMM_WORLD, _Nz, ti, PETSC_COPY_VALUES, &ist); CHKERRQ(ierr);
  ierr = VecScatterCreate(_y, isf, _y0, ist, &_scatters["body2R"]); CHKERRQ(ierr);

  // free memory
  ierr = PetscFree(fi); CHKERRQ(ierr);
  ierr = PetscFree(ti); CHKERRQ(ierr);
  ierr = ISDestroy(&isf); CHKERRQ(ierr);
  ierr = ISDestroy(&ist); CHKERRQ(ierr);

  
  //============================================================================== 
  // set up scatter context to take values for z = 0 from body field and put them on a Vec of size Ny
  // indices to scatter from
  IS isf2;
  /* creates a data structure for an index set with a list of evenly spaced integers
   * locally owned portion of index set has length _Ny
   * first element of locally owned index set is 0
   * change to the next index is _Nz (the stride)
   * takes indices [0, _Nz, 2*_Nz, ..., (_Ny-1)*_Nz]
   */
  ierr = ISCreateStride(PETSC_COMM_WORLD, _Ny, 0, _Nz, &isf2); CHKERRQ(ierr);

  // indices to scatter to
  PetscInt *ti2;
  IS ist2;
  ierr = PetscMalloc1(_Ny,&ti2); CHKERRQ(ierr);

  // length _Ny
  for (PetscInt Ii=0; Ii<_Ny; Ii++) {
    ti2[Ii] = Ii;
  }
  ierr = ISCreateGeneral(PETSC_COMM_WORLD, _Ny, ti2, PETSC_COPY_VALUES, &ist2); CHKERRQ(ierr);
  ierr = VecScatterCreate(_y, isf2, _z0, ist2, &_scatters["body2T"]); CHKERRQ(ierr);

  // free memory
  ierr = PetscFree(ti2); CHKERRQ(ierr);
  ierr = ISDestroy(&isf2); CHKERRQ(ierr);
  ierr = ISDestroy(&ist2); CHKERRQ(ierr);


  //==============================================================================
  // set up scatter context to take values for z = Lz from body field and put them on a Vec of size Ny
  // indices to scatter from
  IS isf3;
  // takes indices [_Nz - 1, 2*_Nz - 1, ..., _Ny*_Nz - 1]
  ierr = ISCreateStride(PETSC_COMM_WORLD, _Ny, _Nz - 1, _Nz, &isf3); CHKERRQ(ierr);

  // indices to scatter to
  PetscInt *ti3;
  IS ist3;
  ierr = PetscMalloc1(_Ny,&ti3); CHKERRQ(ierr);
  for (PetscInt Ii = 0; Ii<_Ny; Ii++) {
    ti3[Ii] = Ii;
  }
  ierr = ISCreateGeneral(PETSC_COMM_WORLD, _Ny, ti3, PETSC_COPY_VALUES, &ist3); CHKERRQ(ierr);
  ierr = VecScatterCreate(_y, isf3, _z0, ist3, &_scatters["body2B"]); CHKERRQ(ierr);

  // free memory
  ierr = PetscFree(ti3); CHKERRQ(ierr);
  ierr = ISDestroy(&isf3); CHKERRQ(ierr);
  ierr = ISDestroy(&ist3); CHKERRQ(ierr);

  return ierr;
}