/*
  This file is part of CDO. CDO is a collection of Operators to manipulate and analyse Climate model Data.

  Author: Fabian Wachsmann

*/

#include <cdi.h>

#include "process_int.h"
#include "param_conversion.h"
#include "util_files.h"
#include "cdo_options.h"
#include "varray.h"
#include "dmemory.h"

#include <cassert>
#include <unistd.h>

#include "pmlist.h"
#include "merge_axis.h"
#include "listbuffer.h"

/** All Vars in the input file will be merged.
*** Therefore, they need to have the same structure i.e. axis sizes
*** Axissize contains the sizes of the first variable named in inputNames.
**/
void
MergeVarsOnAxis::check_axissize_consistency(std::vector<int> axissize)
{
  for (int i = 1; i < this->inputNames.nvalues; i++)
    {
      if (axissize[0] != (int) gridInqXsize(this->inputKeys[i].gridID))
        cdo_abort("ERROR (infile: '%s')! In merging variables to a variable with a character coordinate:\n          Size of x-axis: "
                 "'%d' of variable '%s'\n          differ from x-axis size of variable '%s': '%d'.",
                 cdo_get_stream_name(0), gridInqXsize(this->inputKeys[i].gridID), this->inputNames.values[i].c_str(),
                 this->inputNames.values[0].c_str(), axissize[0]);
      if (axissize[1] != (int) gridInqYsize(this->inputKeys[i].gridID))
        cdo_abort("ERROR (infile: '%s')! In merging variables to a variable with a character coordinate:\n          Size of y-axis: "
                 "'%d' of variable '%s'\n          differ from y-axis size of variable '%s': '%d'.",
                 cdo_get_stream_name(0), gridInqYsize(this->inputKeys[i].gridID), this->inputNames.values[i].c_str(),
                 this->inputNames.values[0].c_str(), axissize[1]);
      if (axissize[2] != zaxisInqSize(this->inputKeys[i].zaxisID))
        cdo_abort("ERROR (infile: '%s')! In merging variables to a variable with a character coordinate:\n          Size of z-axis: "
                 "'%d' of variable '%s'\n          differ from z-axis size of variable '%s': '%d'.",
                 cdo_get_stream_name(0), zaxisInqSize(this->inputKeys[i].zaxisID), this->inputNames.values[i].c_str(),
                 this->inputNames.values[0].c_str(), axissize[2]);
    }
}

/** The next function will define output.gridID, output.zaxisID
*** and reset axissize
***
*** The axis that has size=1 will be used as the merge axis
*/
std::vector<int>
MergeVarsOnAxis::define_new_axes(std::vector<int> axissize)
{
  check_axissize_consistency(axissize);

  double *xvals = nullptr, *yvals = nullptr, *zvals = nullptr, *xcell_bounds = nullptr, *ycell_bounds = nullptr;
  int nvertex = 0;
  double *subsvals = (double *) Malloc(this->inputNames.nvalues * sizeof(double));
  for (int i = 0; i < this->inputNames.nvalues; i++) subsvals[i] = i + 1;

  int gridType = gridInqType(this->inputKeys[0].gridID);
  if (axissize[0] == 1)
    {
      xvals = subsvals;
      yvals = (double *) Malloc(axissize[1] * sizeof(double));
      zvals = (double *) Malloc(axissize[2] * sizeof(double));
      gridInqYvals(this->inputKeys[0].gridID, yvals);
      zaxisInqLevels(this->inputKeys[0].zaxisID, zvals);
      axissize[0] = this->inputNames.nvalues;
      this->output.gridID = gridCreate(GRID_GENERIC, axissize[0] * axissize[1]);
      int zaxisType = zaxisInqType(this->inputKeys[0].zaxisID);
      this->output.zaxisID = zaxisCreate(zaxisType, axissize[2]);
    }
  else if (axissize[1] == 1)
    {
      xvals = (double *) Malloc(axissize[0] * sizeof(double));
      yvals = subsvals;
      zvals = (double *) Malloc(axissize[2] * sizeof(double));
      gridInqXvals(this->inputKeys[0].gridID, xvals);
      zaxisInqLevels(this->inputKeys[0].zaxisID, zvals);
      axissize[1] = this->inputNames.nvalues;
      this->output.gridID = gridCreate(GRID_GENERIC, axissize[0] * axissize[1]);
      int zaxisType = zaxisInqType(this->inputKeys[0].zaxisID);
      this->output.zaxisID = zaxisCreate(zaxisType, axissize[2]);
    }
  else if (axissize[2] == 1)
    {
      zvals = subsvals;
      axissize[2] = this->inputNames.nvalues;
      if (axissize[0] && axissize[1])
        {
          xvals = (double *) Malloc(axissize[0] * sizeof(double));
          yvals = (double *) Malloc(axissize[1] * sizeof(double));
          gridInqXvals(this->inputKeys[0].gridID, xvals);
          gridInqYvals(this->inputKeys[0].gridID, yvals);
          if ( gridType == GRID_UNSTRUCTURED )
            {
              this->output.gridID = gridCreate(gridType, axissize[0] );
              nvertex = gridInqNvertex(this->inputKeys[0].gridID);
              xcell_bounds = (double *) Malloc(nvertex * axissize[0] * sizeof(double));
              ycell_bounds = (double *) Malloc(nvertex * axissize[0] * sizeof(double));
              gridInqYbounds(this->inputKeys[0].gridID, ycell_bounds);
              gridInqXbounds(this->inputKeys[0].gridID, xcell_bounds);
            }
          else
/*            this->output.gridID = gridCreate(gridType, axissize[0]); */
            this->output.gridID = gridCreate(gridType, axissize[0] * axissize[1]); 
        }
      else
        this->output.gridID = gridCreate(gridType, 1);
      this->output.zaxisID = zaxisCreate(ZAXIS_GENERIC, axissize[2]);
    }

  gridDefXsize(this->output.gridID, axissize[0]);
  gridDefYsize(this->output.gridID, axissize[1]);
  if (axissize[0] == 0)
    axissize[0] = 1;
  else
    gridDefXvals(this->output.gridID, xvals);
  if (axissize[1] == 0)
    axissize[1] = 1;
  else
    gridDefYvals(this->output.gridID, yvals);
  zaxisDefLevels(this->output.zaxisID, zvals);

  if ( gridType == GRID_UNSTRUCTURED )
    {
      char unitstring[CDI_MAX_NAME];
      int length = CDI_MAX_NAME;
      cdiInqKeyString(this->inputKeys[0].gridID, CDI_YAXIS, CDI_KEY_UNITS, unitstring, &length);
      gridDefXunits(this->output.gridID, (const char *)unitstring);
      length = CDI_MAX_NAME;
      cdiInqKeyString(this->inputKeys[0].gridID, CDI_YAXIS, CDI_KEY_UNITS, unitstring, &length);
      gridDefYunits(this->output.gridID, (const char *)unitstring);
      gridDefNvertex(this->output.gridID, nvertex);
      gridDefXbounds(this->output.gridID, xcell_bounds);
      gridDefYbounds(this->output.gridID, ycell_bounds);
      axissize[1] = 1;
    }
  Free(xvals);
  Free(yvals);
  Free(zvals);
  std::vector<int> newaxissize = { axissize[0], axissize[1], axissize[2] };
  return newaxissize;
}

/**
*** This function will define
*** output.datatype, output.vlistID, output.varID
*** Therefore, a new var is created in output.vlistID
*** It will allocate output.data for ntsteps*axissizes
**/

void
MergeVarsOnAxis::define_var_structure(int vlistID, int ntsteps, std::vector<int> axissize)
{
  this->output.vlistID = vlistID;
  this->output.varID = vlistDefVar(this->output.vlistID, this->output.gridID, this->output.zaxisID, TIME_VARYING);
  int oldcode = vlistInqVarCode(this->output.vlistID, this->inputKeys[0].varID);
  vlistDefVarCode(this->output.vlistID, this->inputKeys[0].varID, 1);
  vlistDefVarCode(this->output.vlistID, this->output.varID, oldcode);
  cdiDefKeyString(this->output.vlistID, this->inputKeys[0].varID, CDI_KEY_NAME, "ChangedForMap");
  cdiDefKeyString(this->output.vlistID, this->output.varID, CDI_KEY_NAME, this->inputNames.values[0].c_str());
  auto datatype = vlistInqVarDatatype(this->output.vlistID, this->inputKeys[0].varID);
  vlistDefVarDatatype(this->output.vlistID, this->output.varID, datatype);
  const auto mv = vlistInqVarMissval(this->output.vlistID, this->inputKeys[0].varID);
  if (datatype == CDI_DATATYPE_FLT64)
    {
      this->output.datatype = 'd';
      vlistDefVarMissval(this->output.vlistID, this->output.varID, mv);
      this->data = Malloc(ntsteps * axissize[0] * axissize[1] * axissize[2] * sizeof(double));
    }
  else
    {
      this->output.datatype = 'f';
      vlistDefVarMissval(this->output.vlistID, this->output.varID, mv);
      this->data = Malloc(ntsteps * axissize[0] * axissize[1] * axissize[2] * sizeof(float));
    }
}
/**
*** This function reads data from the inputKeys.varID from a streamID
*** to data.
*** The index for data is defined for CMOR
**/

void
MergeVarsOnAxis::read_cmor_charvar(std::vector<int> axissize, int streamID, int oldgridsize)
{
  Varray<double> buffer_old(oldgridsize);

  int tsID = 0;
  while (true)
    {
      auto nrecs = streamInqTimestep(streamID, tsID);
      if (nrecs == 0) break;

      while (nrecs--)
        {
          int varIDrw, levelIDrw;
          size_t nmiss;
          streamInqRecord(streamID, &varIDrw, &levelIDrw);
          for (int i = 0; i < this->inputNames.nvalues; i++)
            if (varIDrw == this->inputKeys[i].varID)
              {
                streamReadRecord(streamID, buffer_old.data(), &nmiss);
                int newIndex;
                for (int j = 0; j < oldgridsize; j++)
                  {
                    // (lev x lat, basin )

                    /* (lev x lat, basin )
                                newIndex = j * levdim + levelID; */
                    /* Use this index because z-axis is registered at the end and rearranged by CMOR */
                    if (oldgridsize == axissize[0] * axissize[1])
                      newIndex = tsID * axissize[2] * axissize[0] * axissize[1] + j * axissize[2] + i;
                    else if (axissize[0] == this->inputNames.nvalues)
                      newIndex = tsID * axissize[2] * axissize[0] * axissize[1] + i * axissize[1] * axissize[2] + j * axissize[2]
                                 + levelIDrw;
                    else
                      newIndex = tsID * axissize[2] * axissize[0] * axissize[1] + levelIDrw * axissize[0] * axissize[1]
                                 + i * axissize[0] * axissize[1] / oldgridsize + j;
                    if (this->output.datatype == 'd')
                      ((double *) this->data)[newIndex] = (double) buffer_old[j];
                    else
                      ((float *) this->data)[newIndex] = (float) buffer_old[j];
                  }
              }
        }
      tsID++;
    }
}
