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

  Copyright (C) 2003-2020 Uwe Schulzweida, <uwe.schulzweida AT mpimet.mpg.de>
  See COPYING file for copying and redistribution conditions.

  This program is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; version 2 of the License.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.
*/

#include <cdi.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"
#include "timebase.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;
  cdoPrint("%s Wrote %.1f GB of %d bit floats to %s %s, %.1f MVal/s", sinfo, data_size, memtype == MemType::Float ? 32 : 64,
           cdi_datatype2str(datatype), cdi_filetype2str(filetype), rout);

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

  rout = (tw > 0) ? 1024 * file_size / tw : -1;
  cdoPrint("%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;
  const char *const defaultgrid = "global_.2";
  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, tw0, t0, twsum = 0;

  cdoInitialize(process);

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

  if (operatorArgc() > 5) cdoAbort("Too many arguments!");

  const char *gridfile = defaultgrid;
  if (operatorArgc() > 0) nruns = parameter2int(cdoOperatorArgv(0));
  if (operatorArgc() > 1) gridfile = cdoOperatorArgv(1).c_str();
  if (operatorArgc() > 2) nlevs = parameter2int(cdoOperatorArgv(2));
  if (operatorArgc() > 3) ntimesteps = parameter2int(cdoOperatorArgv(3));
  if (operatorArgc() > 4) nvars = parameter2int(cdoOperatorArgv(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 = cdoDefineGrid(gridfile);
  auto gridsize = gridInqSize(gridID);

  if (nlevs == 1)
    zaxisID = zaxisFromName("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)
    {
      cdoPrint("nruns      : %d", nruns);
      cdoPrint("gridsize   : %zu", gridsize);
      cdoPrint("nlevs      : %d", nlevs);
      cdoPrint("ntimesteps : %d", ntimesteps);
      cdoPrint("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)
    {
      int varID = vlistDefVar(vlistID, gridID, zaxisID, TIME_VARYING);
      vlistDefVarParam(vlistID, varID, cdiEncodeParam(varID + 1, 255, 255));
    }

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

  // vlistDefNtsteps(vlistID, 1);

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

      const auto streamID = cdoOpenWrite(0);
      cdoDefVlist(streamID, vlistID);

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

      const auto julday = date_to_julday(CALENDAR_PROLEPTIC, 19870101);

      t0 = timer_val(timer_write);

      for (int tsID = 0; tsID < ntimesteps; tsID++)
        {
          const auto vdate = julday_to_date(CALENDAR_PROLEPTIC, julday + tsID);
          const auto vtime = 0;
          taxisDefVdate(taxisID, vdate);
          taxisDefVtime(taxisID, vtime);
          cdoDefTimestep(streamID, tsID);

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

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

      cdoStreamClose(streamID);

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

      file_size = (double) fileSize(cdoGetStreamName(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);

  cdoFinish();

  return nullptr;
}
