// -*- mode: c++ -*-

#include "mesh.hpp"


// ------------------------------------------------------------------------
// -- PUBLIC METHODS
// ------------------------------------------------------------------------

void Mesh::build_surface_mesh(H5::H5File fp, H5::H5File ep, std::string datasetname)
{

  get_mesh_metadata(fp, SURFACE_NCOL);

  handle_surface_node_data(fp);

  build_surface_cells(ep, datasetname);
  
}

void Mesh::build_subsurface_mesh(H5::H5File fp)
{

  get_mesh_metadata(fp, SUBSURFACE_NCOL);

  handle_subsurface_node_data(fp);

  build_subsurface_cells();
  
}

void Mesh::build_environment(H5::H5File ep,
			     std::string datasetname,
			     std::string datasetname_next,
			     int flag)
{

  std::string groupname;
  
  // Read in velocities in x direction
    
  if (flag == SUBSURFACE_FLAG) {
    groupname = "darcy_velocity.cell.0";
  } else if (flag == SURFACE_FLAG) {
    groupname = "surface-velocity.cell.0";
  }

  read_velocity(ep, groupname, datasetname, datasetname_next, flag, 0);

  // Read in velocities in y direction

  if (flag == SUBSURFACE_FLAG) {
    groupname = "darcy_velocity.cell.1";
  } else if (flag == SURFACE_FLAG) {
    groupname = "surface-velocity.cell.1";
  }

  read_velocity(ep, groupname, datasetname, datasetname_next, flag, 1);

  // Read velocities in z direction
  
  if (flag == SUBSURFACE_FLAG) {
    groupname = "darcy_velocity.cell.2";
  } else if (flag == SURFACE_FLAG) {
    groupname = "surface-surface_subsurface_flux.cell.0";
  }

  read_velocity(ep, groupname, datasetname, datasetname_next, flag, 2);

  
  // Surface environment needs adjustments to the vertical velocity
  // and needs to read in water depth
  
  if (flag == SURFACE_FLAG) {
    
    // Rescale the exchange flux from mole to m/s
    for (int i = 0; i < n_cells_surface; i ++) {
      surface_cells[i].set_velocity(2, surface_cells[i].get_velocity(2) * 1.8e-5 * AREADIV);
      surface_cells[i].set_velocity_next(2, surface_cells[i].get_velocity_next(2) * 1.8e-5 * AREADIV);
    }

    // Read in water depth
    read_water_depth(ep, datasetname, datasetname_next);
    
  }
  
}

// ------------------------------------------------------------------------
// GETTERS AND SETTERS
// ------------------------------------------------------------------------

Mesh::Mesh()
{
  n_cells_surface = 0;
  n_cells_subsurface = 0;
  high_coordinates.resize(DIM);
  low_coordinates.resize(DIM);
}

int Mesh::get_n_cells(int flag)
{
  int n_cells;
  
  if (flag == SUBSURFACE_FLAG) {
    n_cells = n_cells_subsurface;
  } else {
    n_cells = n_cells_surface;
  }
  
  return n_cells;
}

Cell* Mesh::get_cell(int index, int flag)
{

  if (flag == SUBSURFACE_FLAG) {
    return &(subsurface_cells[index]);
  } else {
    return &(surface_cells[index]);
  } 
}

void Mesh::set_high_coordinate(int index, double value)
{
  high_coordinates[index] = value;
}

double Mesh::get_high_coordinate(int index)
{
  return high_coordinates[index];
}

void Mesh::set_low_coordinate(int index, double value)
{
  low_coordinates[index] = value;
}

double Mesh::get_low_coordinate(int index)
{
  return low_coordinates[index];
}

void Mesh::set_fuzz(double value)
{
  fuzz = value;
}

// ------------------------------------------------------------------------
// -- PRIVATE METHODS
// ------------------------------------------------------------------------

void Mesh::handle_surface_node_data(H5::H5File fp)
{

  H5::DataSet dataset_n = fp.openDataSet(MESH_NODES);
  H5::DataType datatype_n = dataset_n.getDataType();
  H5::DataSpace dataspace_n = dataset_n.getSpace();

  int rank = dataspace_n.getSimpleExtentNdims();

  hsize_t *dims_n = new hsize_t[rank];
  dataspace_n.getSimpleExtentDims(dims_n);

#if VERBOSE
  for (int i = 0; i < rank; i ++)
    std::cout << "[~~~>] Dimension in " << i << " of nodes: " << dims_n[i] << std::endl;
#endif
  
  double *data_n = new double[dims_n[0]*dims_n[1]];
  dataset_n.read(data_n, H5::PredType::IEEE_F64LE);
  
  nodes.resize(dims_n[0]);

  for (size_t i = 0; i < dims_n[0]; i ++) {
    nodes[i].resize(dims_n[1]);
    for (size_t j = 0; j < dims_n[1]; j ++)
      nodes[i][j] = data_n[i * dims_n[1] + j];
  }
  
}

void Mesh::build_surface_cells(H5::H5File ep, std::string datasetname)
{

  std::vector<double> coords1(DIM-1);
  std::vector<double> coords2(DIM-1);
  std::vector<double> coords3(DIM-1);
  std::vector<double> coords_avg(DIM-1);

  read_elevation(ep, datasetname);

  for (int i = 0; i < n_cells_surface; i ++) {

    int n1 = element_node_indices[i][1];
    int n2 = element_node_indices[i][2];
    int n3 = element_node_indices[i][3];

    std::vector<double> edge_length(DIM);

    for (size_t k = 0; k < DIM; k ++)
      edge_length[k] = DBL_MAX;


    double smallest_distance;
    double largest_distance;
    
    for (size_t j = 0; j < DIM-1; j ++) {

      coords1[j] = nodes[n1][j];
      coords2[j] = nodes[n2][j];
      coords3[j] = nodes[n3][j];

      coords_avg[j] = (coords1[j] + coords2[j] + coords3[j]) / 3.0;

      edge_length[0] = fabs(coords1[j] - coords_avg[j]);
      edge_length[1] = fabs(coords2[j] - coords_avg[j]);
      edge_length[2] = fabs(coords3[j] - coords_avg[j]);

      smallest_distance = edge_length[0];
      largest_distance  = edge_length[0];

      for (size_t k = 1; k < DIM; k ++) {
	smallest_distance = std::min(smallest_distance, edge_length[k]);
	largest_distance  = std::max(largest_distance, edge_length[k]);
      }
      
    }

    surface_cells[i].set_characteristic_length(smallest_distance);
    surface_cells[i].set_maximum_length(largest_distance);

    for (size_t k = 0; k < DIM-1; k ++)
      surface_cells[i].set_coordinate(k, coords_avg[k]);
    
    surface_cells[i].set_cell_id((int) i);

    // find the coupled cell in the subsurface
    std::vector<double> distance(DIM);
    int coupled_neighbor = -1; // store current coupled neighbor here
    double dz = DBL_MAX;       

    for (int k = 0; k < n_cells_subsurface; k ++) {

      for (size_t j = 0; j < DIM; j ++)
	distance[j] = subsurface_cells[k].get_coordinate(j) - surface_cells[i].get_coordinate(j);

      // Because the mesh is a stacked mesh, the only candidates for the subsurface
      // neighbor are the cells located in the same column.  These cells will have
      // the same centroid in the xy-plane as the surface cell.  The problem reduces
      // to finding the nearest cell in the vertical direction in the column.
      
      double xy_distance = sqrt(distance[0]*distance[0] + distance[1]*distance[1]);

      if (xy_distance <= fuzz) {
	dz = std::min(dz, fabs(distance[2]));
	coupled_neighbor = k;
	if (dz <= 2.0 * subsurface_cells[k].get_characteristic_length())
	  break;
      }
      
    }

    surface_cells[i].set_coupled_neighbor(coupled_neighbor);
    subsurface_cells[coupled_neighbor].set_coupled_neighbor(i);
    
  }

}

void Mesh::handle_subsurface_node_data(H5::H5File fp)
{
  
  H5::DataSet dataset_n = fp.openDataSet(MESH_NODES);
  H5::DataType datatype_n = dataset_n.getDataType();
  H5::DataSpace dataspace_n = dataset_n.getSpace();

  int rank = dataspace_n.getSimpleExtentNdims();

  hsize_t *dims_n = new hsize_t[rank];
  dataspace_n.getSimpleExtentDims(dims_n);

#if VERBOSE
  for (int i = 0; i < rank; i ++)
    std::cout << "[~~~>] Dimension in " << i << " of nodes: " << dims_n[i] << std::endl;
#endif
  
  double *data_n = new double[dims_n[0]*dims_n[1]];
  dataset_n.read(data_n, H5::PredType::IEEE_F64LE);

  nodes.resize(dims_n[0]);

  for (size_t i = 0; i < dims_n[0]; i ++) {
    nodes[i].resize(dims_n[1]);
    for (size_t j = 0; j < dims_n[1]; j ++)
      nodes[i][j] = data_n[i * dims_n[1] + j];
  }
  
}

void Mesh::build_subsurface_cells()
{

  std::vector<double> coords1(DIM);
  std::vector<double> coords2(DIM);
  std::vector<double> coords3(DIM);
  std::vector<double> coords4(DIM);
  std::vector<double> coords5(DIM);
  std::vector<double> coords6(DIM);

  std::vector<double> coords_avg(DIM);

  for (size_t i = 0; i < DIM; i ++) {
    set_high_coordinate(i, DBL_MIN);
    set_low_coordinate(i, DBL_MAX);
  }

  double smallest_distance;
  double largest_distance;
    
  for (int i = 0; i < n_cells_subsurface; i ++) {

    int n1 = element_node_indices[i][1];
    int n2 = element_node_indices[i][2];
    int n3 = element_node_indices[i][3];
    int n4 = element_node_indices[i][4];
    int n5 = element_node_indices[i][5];
    int n6 = element_node_indices[i][6];

    std::vector<double> edge_length(N_VERTICES);

    for (size_t k = 0; k < N_VERTICES; k ++)
      edge_length[k] = DBL_MAX;
    
    for (size_t j = 0; j < DIM; j++) {
      
      coords1[j] = nodes[n1][j];
      coords2[j] = nodes[n2][j];
      coords3[j] = nodes[n3][j];
      coords4[j] = nodes[n4][j];
      coords5[j] = nodes[n5][j];
      coords6[j] = nodes[n6][j];

      coords_avg[j] = (coords1[j] + coords2[j] + coords3[j] + coords4[j] + coords5[j] + coords6[j]) / 6.0;

      // the characteristic length is the radius of the insphere of
      // the wedge
      edge_length[0] = fabs(coords1[j] - coords_avg[j]);
      edge_length[1] = fabs(coords2[j] - coords_avg[j]);
      edge_length[2] = fabs(coords3[j] - coords_avg[j]);
      edge_length[3] = fabs(coords4[j] - coords_avg[j]);
      edge_length[4] = fabs(coords5[j] - coords_avg[j]);
      edge_length[5] = fabs(coords6[j] - coords_avg[j]);

      smallest_distance = edge_length[0];
      largest_distance  = edge_length[0];

      for (size_t k = 1; k < N_VERTICES; k ++) {
      smallest_distance = std::min(smallest_distance, edge_length[k]);
      largest_distance  = std::max(largest_distance, edge_length[k]);
      }
      
    }

    subsurface_cells[i].set_characteristic_length(smallest_distance);
    subsurface_cells[i].set_maximum_length(largest_distance);
    
    for (size_t k = 0; k < DIM; k ++) {

      // update lowest and highest corner of mesh
      high_coordinates[k] = std::max(coords_avg[k], high_coordinates[k]);
      low_coordinates[k] = std::min(coords_avg[k], low_coordinates[k]);

      subsurface_cells[i].set_coordinate(k, coords_avg[k]);
      
    }
    
    subsurface_cells[i].set_cell_id((int) i);
    
  }
  
}

void Mesh::get_mesh_metadata(H5::H5File fp, int flag)
{

  H5::DataSet dataset_e = fp.openDataSet(MESH_MIXED_ELEMENTS);
  H5::DataType datatype_e = dataset_e.getDataType();
  H5::DataSpace dataspace_e = dataset_e.getSpace();

  int rank = dataspace_e.getSimpleExtentNdims();

  hsize_t *dims_e = new hsize_t[rank];
  dataspace_e.getSimpleExtentDims(dims_e);

#if VERBOSE
  for (int i = 0; i < rank; i ++)
      std::cout << "[~~~>] Dimension in " << i << " of element node map: " << dims_e[i] << std::endl;
#endif

  int *data_e = new int[dims_e[0]*dims_e[1]];
  dataset_e.read(data_e, H5::PredType::STD_I32LE);

  size_t nx_e = flag;
  size_t ny_e = dims_e[0] * dims_e[1] / nx_e;
  
  element_node_indices.resize(ny_e);

  for (size_t i = 0; i < ny_e; i ++) {
    element_node_indices[i].resize(nx_e);
    for (size_t j = 0; j < nx_e; j ++)
      element_node_indices[i][j] = data_e[i * nx_e + j];
  }
  
  if (flag == SUBSURFACE_NCOL) {
    n_cells_subsurface = ny_e;
    subsurface_cells.resize(ny_e);
  } else if (flag == SURFACE_NCOL) {
    n_cells_surface = ny_e;
    surface_cells.resize(ny_e);
  }
  
}

void Mesh::read_velocity(H5::H5File fp,
			 std::string groupname,
			 std::string datasetname,
			 std::string datasetname_next,
			 int flag,
			 int index)
{

  // read in the current time step data
  
  H5::Group group		= fp.openGroup(groupname.c_str());
  H5::DataSet dataset		= group.openDataSet(datasetname.c_str());
  H5::DataType datatype		= dataset.getDataType();
  H5::DataSpace dataspace	= dataset.getSpace();
  
  int rank = dataspace.getSimpleExtentNdims();

  hsize_t *dims = new hsize_t[rank];
  dataspace.getSimpleExtentDims(dims);
  
  double *data = new double[dims[0]*dims[1]];
  dataset.read(data, H5::PredType::IEEE_F64LE);

  if (flag == SUBSURFACE_FLAG) {
    for (int i = 0; i < n_cells_subsurface; i ++)
      subsurface_cells[i].set_velocity(index, data[i]);
  } else if (flag == SURFACE_FLAG) {
    for (int i = 0; i < n_cells_surface; i ++)
      surface_cells[i].set_velocity(index, data[i]);
  }

  // read in the next time step data for linear interpolation
  
  H5::DataSet dataset_next	= group.openDataSet(datasetname_next.c_str());
  H5::DataType datatype_next	= dataset_next.getDataType();
  H5::DataSpace dataspace_next	= dataset_next.getSpace();

  rank = dataspace_next.getSimpleExtentNdims();
  
  dims = new hsize_t[rank];
  dataspace_next.getSimpleExtentDims(dims);

  data = new double[dims[0]*dims[1]];
  dataset_next.read(data, H5::PredType::IEEE_F64LE);

  if (flag == SUBSURFACE_FLAG) {
    for (int i = 0; i < n_cells_subsurface; i ++)
      subsurface_cells[i].set_velocity_next(index, data[i]);
  } else if (flag == SURFACE_FLAG) {
    for (int i = 0; i < n_cells_surface; i ++)
      surface_cells[i].set_velocity_next(index, data[i]);
  }

}

void Mesh::read_water_depth(H5::H5File fp,
			    std::string datasetname,
			    std::string datasetname_next)
{

  // read in current time step
  
  H5::Group group		= fp.openGroup("surface-ponded_depth.cell.0");
  H5::DataSet dataset		= group.openDataSet(datasetname.c_str());
  H5::DataType datatype		= dataset.getDataType();
  H5::DataSpace dataspace	= dataset.getSpace();
  
  int rank = dataspace.getSimpleExtentNdims();

  hsize_t *dims = new hsize_t[rank];
  dataspace.getSimpleExtentDims(dims);
  
  double *data = new double[dims[0]*dims[1]];
  dataset.read(data, H5::PredType::IEEE_F64LE);

  for (int i = 0; i < n_cells_surface; i ++) {
    surface_cells[i].set_water_depth(data[i]);
  }

  // read in next time step

  H5::DataSet dataset_next	= group.openDataSet(datasetname_next.c_str());
  H5::DataType datatype_next	= dataset_next.getDataType();
  H5::DataSpace dataspace_next	= dataset_next.getSpace();
  
  rank = dataspace_next.getSimpleExtentNdims();

  dims = new hsize_t[rank];
  dataspace_next.getSimpleExtentDims(dims);
  
  data = new double[dims[0]*dims[1]];
  dataset_next.read(data, H5::PredType::IEEE_F64LE);

  for (int i = 0; i < n_cells_surface; i ++) {
    surface_cells[i].set_water_depth_next(data[i]);
  }
  
}

void Mesh::read_elevation(H5::H5File fp, std::string datasetname)
{

  H5::Group group		= fp.openGroup("surface-elevation.cell.0");
  H5::DataSet dataset		= group.openDataSet(datasetname.c_str());
  H5::DataType datatype		= dataset.getDataType();
  H5::DataSpace dataspace	= dataset.getSpace();
  
  int rank = dataspace.getSimpleExtentNdims();

  hsize_t *dims = new hsize_t[rank];
  dataspace.getSimpleExtentDims(dims);
  
  double *data = new double[dims[0]*dims[1]];
  dataset.read(data, H5::PredType::IEEE_F64LE);
  
  for (int i = 0; i < n_cells_surface; i ++) {
    surface_cells[i].set_coordinate(2, data[i]);
  }
  
}