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

  Author: Uwe Schulzweida

*/

#include <cdi.h>
#include "julian_date.h"

#include "cdo_options.h"
#include "cdo_cdi_wrapper.h"
#include "cdo_zaxis.h"
#include "process_int.h"
#include "param_conversion.h"
#include <mpim_grid.h>
#include "griddes.h"
#include "timer.h"
#include "util_files.h"

static void
print_stat(const char *sinfo, MemType memtype, int datatype, int filetype, off_t nvalues, double data_size, double file_size,
           double tw)
{
  nvalues /= 1000000;
  data_size /= 1024. * 1024. * 1024.;

  double rout = (tw > 0) ? nvalues / tw : -1;
  cdo_print("%s Wrote %.1f GB of %d bit floats to %s %s, %.1f MVal/s", sinfo, data_size, (memtype == MemType::Float) ? 32 : 64,
            cdi_datatype_to_str(datatype), cdi_filetype_to_str(filetype), rout);

  file_size /= 1024. * 1024. * 1024.;

  rout = (tw > 0) ? 1024 * file_size / tw : -1;
  cdo_print("%s Wrote %.1f GB in %.1f seconds, total %.1f MB/s", sinfo, file_size, tw, rout);
}

void *
CDIwrite(void *process)
{
  const auto memtype = Options::CDO_Memtype;
  int nvars = 10, nlevs = 0, ntimesteps = 30;
  int zaxisID;
  int filetype = -1, datatype = -1;
  int nruns = 1;
  char sinfo[64] = { 0 };
  off_t nvalues = 0;
  double file_size = 0, data_size = 0;
  double tw, twsum = 0;

  cdo_initialize(process);

  if (Options::cdoVerbose) cdo_print("parameter: <nruns, <grid, <nlevs, <ntimesteps, <nvars>>>>>");

  if (cdo_operator_argc() > 5) cdo_abort("Too many arguments!");

  std::string defaultgrid = "global_.2";
  auto gridfile = defaultgrid;
  if (cdo_operator_argc() > 0) nruns = parameter_to_int(cdo_operator_argv(0));
  if (cdo_operator_argc() > 1) gridfile = cdo_operator_argv(1);
  if (cdo_operator_argc() > 2) nlevs = parameter_to_int(cdo_operator_argv(2));
  if (cdo_operator_argc() > 3) ntimesteps = parameter_to_int(cdo_operator_argv(3));
  if (cdo_operator_argc() > 4) nvars = parameter_to_int(cdo_operator_argv(4));

  nruns = std::min(std::max(nruns, 0), 9999);
  nlevs = std::min(std::max(nlevs, 1), 255);
  nvars = std::max(nvars, 1);
  ntimesteps = std::max(ntimesteps, 1);

  const auto gridID = cdo_define_grid(gridfile);
  auto gridsize = gridInqSize(gridID);

  if (nlevs == 1)
    zaxisID = zaxis_from_name("surface");
  else
    {
      Varray<double> levels(nlevs);
      for (int i = 0; i < nlevs; ++i) levels[i] = 100 * i;
      zaxisID = zaxisCreate(ZAXIS_HEIGHT, nlevs);
      zaxisDefLevels(zaxisID, &levels[0]);
    }

  if (Options::cdoVerbose)
    {
      cdo_print("nruns      : %d", nruns);
      cdo_print("gridsize   : %zu", gridsize);
      cdo_print("nlevs      : %d", nlevs);
      cdo_print("ntimesteps : %d", ntimesteps);
      cdo_print("nvars      : %d", nvars);
    }

  Varray<double> array(gridsize), xvals(gridsize), yvals(gridsize);

  auto gridID2 = gridID;
  if (gridInqType(gridID) == GRID_GME) gridID2 = gridToUnstructured(gridID, 0);

  if (gridInqType(gridID) != GRID_UNSTRUCTURED && gridInqType(gridID) != GRID_CURVILINEAR) gridID2 = gridToCurvilinear(gridID, 0);

  gridInqXvals(gridID2, &xvals[0]);
  gridInqYvals(gridID2, &yvals[0]);

  // Convert lat/lon units if required
  cdo_grid_to_radian(gridID2, CDI_XAXIS, gridsize, &xvals[0], "grid center lon");
  cdo_grid_to_radian(gridID2, CDI_YAXIS, gridsize, &yvals[0], "grid center lat");

  for (size_t i = 0; i < gridsize; ++i) array[i] = 2 - std::cos(std::acos(std::cos(xvals[i]) * std::cos(yvals[i])) / 1.2);

  Varray3D<double> vars(nvars);
  for (int varID = 0; varID < nvars; ++varID)
    {
      vars[varID].resize(nlevs);
      for (int levelID = 0; levelID < nlevs; ++levelID)
        {
          vars[varID][levelID].resize(gridsize);
          for (size_t i = 0; i < gridsize; ++i) vars[varID][levelID][i] = varID + array[i] * (levelID + 1);
        }
    }

  std::vector<float> farray;
  if (memtype == MemType::Float) farray.resize(gridsize);

  const auto vlistID = vlistCreate();

  for (int i = 0; i < nvars; ++i)
    {
      const auto varID = vlistDefVar(vlistID, gridID, zaxisID, TIME_VARYING);
      vlistDefVarParam(vlistID, varID, cdiEncodeParam(varID + 1, 255, 255));
    }

  const auto taxisID = cdo_taxis_create(TAXIS_RELATIVE);
  vlistDefTaxis(vlistID, taxisID);

  // vlistDefNtsteps(vlistID, 1);

  for (int irun = 0; irun < nruns; ++irun)
    {
      auto tw0 = timer_val(timer_write);
      data_size = 0;
      nvalues = 0;

      const auto streamID = cdo_open_write(0);
      cdo_def_vlist(streamID, vlistID);

      filetype = cdo_inq_filetype(streamID);
      datatype = vlistInqVarDatatype(vlistID, 0);
      if (datatype == CDI_UNDEFID) datatype = CDI_DATATYPE_FLT32;

      const auto julday = date_to_julday(CALENDAR_PROLEPTIC, 19870101);

      auto t0 = timer_val(timer_write);

      for (int tsID = 0; tsID < ntimesteps; ++tsID)
        {
          CdiDateTime vDateTime = { };
          vDateTime.date = cdiDate_set(julday_to_date(CALENDAR_PROLEPTIC, julday + tsID));
          taxisDefVdatetime(taxisID, vDateTime);
          cdo_def_timestep(streamID, tsID);

          for (int varID = 0; varID < nvars; ++varID)
            {
              for (int levelID = 0; levelID < nlevs; ++levelID)
                {
                  nvalues += gridsize;
                  cdo_def_record(streamID, varID, levelID);
                  if (memtype == MemType::Float)
                    {
                      for (size_t i = 0; i < gridsize; ++i) farray[i] = vars[varID][levelID][i];
                      cdo_write_record_f(streamID, &farray[0], 0);
                      data_size += gridsize * 4;
                    }
                  else
                    {
                      cdo_write_record(streamID, &vars[varID][levelID][0], 0);
                      data_size += gridsize * 8;
                    }
                }
            }

          if (Options::cdoVerbose)
            {
              tw = timer_val(timer_write) - t0;
              t0 = timer_val(timer_write);
              cdo_print("Timestep %d: %.2f seconds", tsID + 1, tw);
            }
        }

      cdo_stream_close(streamID);

      tw = timer_val(timer_write) - tw0;
      twsum += tw;

      file_size = (double) FileUtils::size(cdo_get_stream_name(0));

      if (nruns > 1) snprintf(sinfo, sizeof(sinfo), "(run %d)", irun + 1);

      print_stat(sinfo, memtype, datatype, filetype, nvalues, data_size, file_size, tw);
    }

  if (nruns > 1) print_stat("(mean)", memtype, datatype, filetype, nvalues, data_size, file_size, twsum / nruns);

  vlistDestroy(vlistID);

  cdo_finish();

  return nullptr;
}
