#include "nq_basic.h"
#include "nq_structure.h"

#include <netcdf.h>
#include <math.h>

#include <visu_tools.h>
#include <visu_elements.h>
#include <extraFunctions/scalarFields.h>
#include <coreTools/toolFileFormat.h>
#include <coreTools/toolMatrix.h>

/* Local methods */
static VisuDataLoader *nqLoader = NULL;

static gboolean loadNQETSF(VisuDataLoader *self, VisuDataLoadable *data,
                           guint type, guint nSet,
                           GCancellable *cancel, GError **error);

VisuDataLoader* visu_data_loader_ETSF_getStatic(void)
{
  const gchar *type[] = {"*.nc", "*-etsf.nc", (char*)0};

  if (nqLoader)
    return nqLoader;

  return nqLoader = visu_data_loader_new(_("ETSF file format"), type,
                                         FALSE, loadNQETSF, 5);
}

static gboolean loadNQETSF(VisuDataLoader *self _U_, VisuDataLoadable *data,
                           guint type, guint nSet _U_,
                           GCancellable *cancel _U_, GError **error)
{
  gboolean res;
  int netcdfId, status;
  guint i;
  int varId;
  size_t dimSize;
  char* varNbElement = "number_of_atom_species";
  guint nbEle;
  char* varNbNode = "number_of_atoms";
  guint nbNode;
  char* varAtomsNames = "atom_species_names";
  int varIdSymbols;
  size_t dimsSymbols[2];
  char* symbols;
  char* varNodeToEle = "atom_species";
  int varIdNodeToEle;
  size_t dimsNodeToEle;
  guint *nodeToEle;
  char* varCoord = "reduced_atom_positions";
  int varIdCoord;
  size_t dimsCoord[2];
  double* coord;
  char *varRprimd = "primitive_vectors";
  int varIdRprimd;
  size_t dimsRprimd[2];
  double rprimd[3][3];
  size_t start[] = {0, 0, 0};
  char *varTitle = "title";
  size_t sizeTitle;
  char title[256];
  VisuElement *ele;
  GArray *visuEle, *nbNodesPerEle, *gcoord_c, *gcoord_f;
  float vect[3];
  nc_type ncType;
  gchar *infoUTF8, *name;
  gboolean *flag;
  guint nbGrid, nbFGrid, grid[3], *cGrid, *nGrid;
  size_t dimsGrid[3], dimsNGrid[2];
  int varIdGrid, varIdNGrid;
  VisuBox *box;

  g_return_val_if_fail(error && *error == (GError*)0, FALSE);

  res = nqOpen_netcdfFile(visu_data_loadable_getFilename(data, type), &netcdfId, error);
  if (!res)
    return FALSE;

  /* From now on, the file is a NETCDF valid file, so
     we return TRUE, but errrors can still occurs. */

  /********************************/
  /* Retrieve the numeral values. */
  /********************************/
  /* Grep the number of elements. */
  if (!nqGetDim(netcdfId, error, varNbElement, &varId, &dimSize))
    {
      nqClose_netcdfFile(netcdfId);
      return TRUE;
    }
  nbEle = (int)dimSize;
  /* Grep the number of nodes. */
  if (!nqGetDim(netcdfId, error, varNbNode, &varId, &dimSize))
    {
      nqClose_netcdfFile(netcdfId);
      return TRUE;
    }
  nbNode = (int)dimSize;

  /**************************************************************/
  /* Check the conformance of the arrays to the specifications. */
  /**************************************************************/
  /* Check the names of elements. */
  dimsSymbols[0] = nbEle;
  dimsSymbols[1] = 80;
  if (!nqCheckVar(netcdfId, error, varAtomsNames, &varIdSymbols,
                  NC_CHAR, 2, dimsSymbols))
    {
      g_error_free(*error);
      *error = (GError*)0;
      dimsSymbols[1] = 2;
      if (!nqCheckVar(netcdfId, error, "chemical_symbols", &varIdSymbols,
                      NC_CHAR, 2, dimsSymbols))
        {
          nqClose_netcdfFile(netcdfId);
          return TRUE;
        }
    }
  /* Check the 'node to element'array. */
  dimsNodeToEle = nbNode;
  if (!nqCheckVar(netcdfId, error, varNodeToEle, &varIdNodeToEle,
		 NC_INT, 1, &dimsNodeToEle))
		{
      nqClose_netcdfFile(netcdfId);
      return TRUE;
    }
  /* Check the reduce coordinates array. */
  dimsCoord[0] = nbNode;
  dimsCoord[1] = 3;
  if (!nqCheckVar(netcdfId, error, varCoord, &varIdCoord,
		 NC_DOUBLE, 2, dimsCoord))
		{
      nqClose_netcdfFile(netcdfId);
      return TRUE;
    }
  /* Check the rprimd matrix. */
  dimsRprimd[0] = 3;
  dimsRprimd[1] = 3;
  if (!nqCheckVar(netcdfId, error, varRprimd, &varIdRprimd,
		 NC_DOUBLE, 2, dimsRprimd))
		{
      nqClose_netcdfFile(netcdfId);
      return TRUE;
    }

  /****************************/
  /* Retrieve now the arrays. */
  /****************************/
  /* Grep the names of elements. */
  dimsSymbols[0] = 1;
  symbols = g_malloc(sizeof(char) * (dimsSymbols[1] + 1) * nbEle);
  for (i = 0; i < nbEle; i++)
    {
      start[0] = i;
      status = nc_get_vara_text(netcdfId, varIdSymbols, start, dimsSymbols,
                                symbols + i * (dimsSymbols[1] + 1));
      if (status != NC_NOERR)
        {
          *error = g_error_new(NQ_ERROR, NQ_ERROR_FILE_FORMAT,
                               _("Retrieve value for variable '%s': %s."),
                               varAtomsNames, nc_strerror(status));
          nqClose_netcdfFile(netcdfId);
          g_free(symbols);
          return TRUE;
        }
      symbols[(i + 1) * (dimsSymbols[1] + 1) - 1] = '\0';
    }
  dimsSymbols[0] = nbEle;
  start[0] = 0;
  /* Grep table from 'node to element'. */
  nodeToEle = g_malloc(sizeof(int) * nbNode);
  status = nc_get_vara_int(netcdfId, varIdNodeToEle, start,
			   &dimsNodeToEle, (int*)nodeToEle);
  if (status != NC_NOERR)
    {
      *error = g_error_new(NQ_ERROR, NQ_ERROR_FILE_FORMAT,
			   _("Retrieve value for variable '%s': %s."),
			   varNodeToEle, nc_strerror(status));
      nqClose_netcdfFile(netcdfId);
      g_free(symbols);
      g_free(nodeToEle);
      return TRUE;
    }
  for (i = 0; i < nbNode; i++)
    {
      /* Specifications 1.3 says that the values range from 1.
         So we remove 1 to be coherent with C numbering. */
      nodeToEle[i] -= 1;
      if (nodeToEle[i] >= nbEle)
	{
	  *error = g_error_new(NQ_ERROR, NQ_ERROR_FILE_FORMAT,
			       _("Error in indexing array '%s', index out of bounds."),
			       varNodeToEle);
	  nqClose_netcdfFile(netcdfId);
	  g_free(symbols);
	  g_free(nodeToEle);
	  return TRUE;
	}
    }
  /* Grep the reduced coordinates. */
  coord = g_malloc(sizeof(double) * 3 * nbNode);
  status = nc_get_vara_double(netcdfId, varIdCoord, start, dimsCoord, coord);
  if (status != NC_NOERR)
    {
      *error = g_error_new(NQ_ERROR, NQ_ERROR_FILE_FORMAT,
			   _("Retrieve value for variable '%s': %s."),
			   varCoord, nc_strerror(status));
      nqClose_netcdfFile(netcdfId);
      g_free(symbols);
      g_free(nodeToEle);
      g_free(coord);
      return TRUE;
    }
  /* Grep the box definition. */
  status = nc_get_vara_double(netcdfId, varIdRprimd, start, dimsRprimd, &rprimd[0][0]);
  if (status != NC_NOERR)
    {
      *error = g_error_new(NQ_ERROR, NQ_ERROR_FILE_FORMAT,
			   _("Retrieve value for variable '%s': %s."),
			   varRprimd, nc_strerror(status));
      nqClose_netcdfFile(netcdfId);
      g_free(symbols);
      g_free(nodeToEle);
      g_free(coord);
      return TRUE;
    }

  /* Prepare the arrays for memory allocation in the VisuData object. */
  box = visu_box_new_full(rprimd, VISU_BOX_PERIODIC);
  if (visu_box_getGeometry(box, VISU_BOX_DXX) == G_MAXFLOAT)
    {
      *error = g_error_new(NQ_ERROR, NQ_ERROR_FILE_FORMAT,
                           _("Wrong ETSF format, basis-set is degenerated.\n"));
      g_object_unref(box);
      nqClose_netcdfFile(netcdfId);
      g_free(symbols);
      g_free(nodeToEle);
      g_free(coord);
      return TRUE;
    }
  visu_box_setUnit(box, TOOL_UNITS_BOHR);
  visu_boxed_setBox(VISU_BOXED(data), VISU_BOXED(box));
  g_object_unref(box);

  nbNodesPerEle = g_array_sized_new(FALSE, TRUE, sizeof(guint), nbEle);
  g_array_set_size(nbNodesPerEle, nbEle);
  for (i = 0; i < nbNode; i++)
    g_array_index(nbNodesPerEle, guint, nodeToEle[i]) += 1;
  visuEle = g_array_sized_new(FALSE, FALSE, sizeof(VisuElement*), nbEle);
  for (i = 0; i < nbEle; i++)
    {
      name = g_strstrip(symbols + i * (dimsSymbols[1] + 1));
      ele = visu_element_retrieveFromName(name, (gboolean*)0);
      g_array_append_val(visuEle, ele);
    }

  DBG_fprintf(stderr, "NQ structure : File '%s' parsed.\n", visu_data_loadable_getFilename(data, type));
  DBG_fprintf(stderr, " | box definition : ( %f %f %f )\n",
	      rprimd[0][0], rprimd[0][1], rprimd[0][2]);
  DBG_fprintf(stderr, " |                  ( %f %f %f )\n",
	      rprimd[1][0], rprimd[1][1], rprimd[1][2]);
  DBG_fprintf(stderr, " |                  ( %f %f %f )\n",
	      rprimd[2][0], rprimd[2][1], rprimd[2][2]);
  DBG_fprintf(stderr, " | number of nodes per element (%d):\n", nbEle);
  for (i = 0; i < nbEle ; i++)
    DBG_fprintf(stderr, " |  %d nodes for '%s' (%p)\n",
                g_array_index(nbNodesPerEle, guint, i),
                g_array_index(visuEle, VisuElement*, i)->name,
                g_array_index(visuEle, gpointer, i));
  g_free(symbols);

  visu_node_array_allocate(VISU_NODE_ARRAY(data), visuEle, nbNodesPerEle);
  for (i = 0; i < nbNode; i++)
    {
      vect[0] = coord[3 * i + 0];
      vect[1] = coord[3 * i + 1];
      vect[2] = coord[3 * i + 2];
      visu_data_addNodeFromElement(VISU_DATA(data),
                                   g_array_index(visuEle, VisuElement*, nodeToEle[i]),
                                   vect, TRUE);
    }

  g_free(nodeToEle);
  g_free(coord);
  g_array_free(nbNodesPerEle, TRUE);
  g_array_free(visuEle, TRUE);

  /**********************/
  /* Optional elements. */
  /**********************/
  status = nc_inq_att(netcdfId, NC_GLOBAL, varTitle, &ncType, &sizeTitle);
  if (status == NC_NOERR && ncType == NC_CHAR && sizeTitle < 255)
    {
  	  status = nc_get_att_text(netcdfId, NC_GLOBAL, varTitle, title);
  	  if (status == NC_NOERR)
  	    {
  	      title[sizeTitle] = '\0';
	      infoUTF8 = g_locale_to_utf8(title, -1, NULL, NULL, NULL);
	      if (infoUTF8)
		{
		  visu_data_setDescription(VISU_DATA(data), infoUTF8);
		  g_free(infoUTF8);
		}
	      else
		g_warning("Can't convert '%s' to UTF8.\n", title);
  	    }
    }
  /* Support for BigDFT grid points. */
  nbGrid = 0;
  grid[0] = 0;
  grid[1] = 0;
  grid[2] = 0;
  if (nqGetDim(netcdfId, (GError**)0, "max_number_of_basis_grid_points", &varId, &dimSize))
    {
      nbGrid = (guint)dimSize;
      if (nqGetDim(netcdfId, (GError**)0, "number_of_grid_points_vector1", &varId, &dimSize))
        grid[0] = (guint)dimSize;
      if (nqGetDim(netcdfId, (GError**)0, "number_of_grid_points_vector2", &varId, &dimSize))
        grid[1] = (guint)dimSize;
      if (nqGetDim(netcdfId, (GError**)0, "number_of_grid_points_vector3", &varId, &dimSize))
        grid[2] = (guint)dimSize;
      if (!grid[0] || !grid[1] || !grid[2])
        nbGrid = 0;
      DBG_fprintf(stderr, "NQ structure: grid is defined (%d) in %dx%dx%d.\n",
                  nbGrid, grid[0], grid[1], grid[2]);
      dimsGrid[0] = 1;
      dimsGrid[1] = nbGrid;
      dimsGrid[2] = 3;
      if (!nqCheckVar(netcdfId, (GError**)0, "coordinates_of_basis_grid_points",
                      &varIdGrid, NC_INT, 3, dimsGrid))
        nbGrid = 0;
      dimsNGrid[0] = 1;
      dimsNGrid[1] = nbGrid;
      if (!nqCheckVar(netcdfId, (GError**)0, "number_of_coefficients_per_grid_point",
                      &varIdNGrid, NC_INT, 2, dimsNGrid))
        nbGrid = 0;
      DBG_fprintf(stderr, "NQ structure: grid is present %d and valid.\n", nbGrid);
    }
  if (nbGrid)
    {
      nbFGrid = 0;
      cGrid = g_malloc(sizeof(guint) * nbGrid * 3);
      nGrid = (guint*)0;
      status = nc_get_vara_int(netcdfId, varIdGrid, start,
                               dimsGrid, (int*)cGrid);
      if (status != NC_NOERR)
        {
          g_free(cGrid);
          nbGrid = 0;
        }
      else
        {
          nGrid = g_malloc(sizeof(guint) * nbGrid);
          status = nc_get_vara_int(netcdfId, varIdNGrid, start,
                                   dimsNGrid, (int*)nGrid);
          if (status != NC_NOERR)
            {
              g_free(cGrid);
              g_free(nGrid);
              nbGrid = 0;
            }
          else
            {
              /* Calculate the number of fine grid points. */
              nbFGrid = 0;
              for (i = 0; i < nbGrid; i++)
                if (nGrid[i] > 1) nbFGrid += 1;
            }
        }
      DBG_fprintf(stderr, "NQ structure: grid has been read.\n");
      if (cGrid && nGrid)
        {
          /* Add the grid nodes. */
          gcoord_c = g_array_sized_new(FALSE, FALSE, sizeof(gfloat) * 3, nbGrid);
          gcoord_f = (GArray*)0;
          if (nbFGrid > 0)
            gcoord_f = g_array_sized_new(FALSE, FALSE, sizeof(gfloat) * 3, nbFGrid);
          for (i = 0; i < nbGrid; i++)
            {
              vect[0] = (float)cGrid[i * 3 + 0] / (float)grid[0];
              vect[1] = (float)cGrid[i * 3 + 1] / (float)grid[1];
              vect[2] = (float)cGrid[i * 3 + 2] / (float)grid[2];
              g_array_append_val(gcoord_c, vect);
              if (nGrid[i] > 1)
                g_array_append_val(gcoord_f, vect);
            }
          g_free(cGrid);
          g_free(nGrid);
          g_object_set_data_full(G_OBJECT(data), "BigDFT_coarse_grid",
                                 (gpointer)gcoord_c, (GDestroyNotify)g_array_unref);
          if (nbFGrid > 0)
            g_object_set_data_full(G_OBJECT(data), "BigDFT_fine_grid",
                                   (gpointer)gcoord_f, (GDestroyNotify)g_array_unref);
        }
    }

  nqClose_netcdfFile(netcdfId);
  
  /******************************************/
  /* Add some flags to the VisuData object. */
  /******************************************/
  flag = g_malloc(sizeof(gboolean));
  *flag = TRUE;
  g_object_set_data_full(G_OBJECT(data), VISU_SCALAR_FIELD_DEFINED_IN_STRUCT_FILE,
			 (gpointer)flag, g_free);

  return TRUE;
}
