/*
  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.
*/

/*
   This module contains the following operators:

      Sort sortcode  Sort by code number
*/

#include <algorithm> // sort

#include <cdi.h>

#include "cdo_options.h"
#include "cdo_vlist.h"
#include "process_int.h"
#include "param_conversion.h"
#include "cdo_zaxis.h"

struct LevInfo
{
  int levelID;
  size_t nmiss;
  double level;
};

struct VarInfo
{
  int varID;
  int nlevs;
  int code;
  char param[CDI_MAX_NAME];
  char name[CDI_MAX_NAME];
  std::vector<LevInfo> levInfo;
};

static bool
cmpvarcode(const VarInfo &a, const VarInfo &b)
{
  return a.code < b.code;
}

static bool
cmpvarparam(const VarInfo &a, const VarInfo &b)
{
  return strcmp(a.param, b.param) < 0;
}

static bool
cmpvarname(const VarInfo &a, const VarInfo &b)
{
  return strcmp(a.name, b.name) < 0;
}

static bool
cmpvarlevel(const LevInfo &a, const LevInfo &b)
{
  return a.level < b.level;
}

static bool
cmpvarlevelrev(const LevInfo &a, const LevInfo &b)
{
  return a.level > b.level;
}

static void
setNmiss(int varID, int levelID, int nvars, std::vector<VarInfo> &varInfo, size_t nmiss)
{
  int vindex, lindex;

  for (vindex = 0; vindex < nvars; vindex++)
    if (varInfo[vindex].varID == varID) break;

  if (vindex == nvars) cdoAbort("Internal problem; varID not found!");

  const auto nlevs = varInfo[vindex].nlevs;
  for (lindex = 0; lindex < nlevs; lindex++)
    if (varInfo[vindex].levInfo[lindex].levelID == levelID) break;

  if (lindex == nlevs) cdoAbort("Internal problem; levelID not found!");

  varInfo[vindex].levInfo[lindex].nmiss = nmiss;
}

void *
Sort(void *process)
{
  int varID, levelID, zaxisID;
  int vindex, lindex;
  int nrecs, nlevs;
  size_t nmiss;
  bool (*cmpvarlev)(const LevInfo &a, const LevInfo &b) = cmpvarlevel;

  cdoInitialize(process);

  // clang-format off
  const auto SORTCODE  = cdoOperatorAdd("sortcode",  0, 0, nullptr);
  const auto SORTPARAM = cdoOperatorAdd("sortparam", 0, 0, nullptr);
  const auto SORTNAME  = cdoOperatorAdd("sortname",  0, 0, nullptr);
  const auto SORTLEVEL = cdoOperatorAdd("sortlevel", 0, 0, nullptr);
  // clang-format on

  const auto operatorID = cdoOperatorID();

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

  if (operatorID == SORTLEVEL && operatorArgc() == 1)
    {
      auto iarg = parameter2int(cdoOperatorArgv(0));
      if (iarg < 0) cmpvarlev = cmpvarlevelrev;
    }

  const auto streamID1 = cdoOpenRead(0);

  const auto vlistID1 = cdoStreamInqVlist(streamID1);
  const auto vlistID2 = vlistDuplicate(vlistID1);

  const auto taxisID1 = vlistInqTaxis(vlistID1);
  const auto taxisID2 = taxisDuplicate(taxisID1);
  vlistDefTaxis(vlistID2, taxisID2);
  /*
  if ( operatorID == SORTCODE )
      vlistSortCode(vlistID2);
   else if ( operatorID == SORTNAME )
      ;
   else if ( operatorID == SORTLEVEL )
      ;
  */

  const auto streamID2 = cdoOpenWrite(1);
  cdoDefVlist(streamID2, vlistID2);

  VarList varList1;
  varListInit(varList1, vlistID1);

  const auto nvars = vlistNvars(vlistID1);

  std::vector<VarInfo> varInfo(nvars);
  for (varID = 0; varID < nvars; ++varID)
    {
      nlevs = zaxisInqSize(vlistInqVarZaxis(vlistID1, varID));
      varInfo[varID].nlevs = nlevs;
      varInfo[varID].levInfo.resize(nlevs);
    }

  Varray2D<double> vardata(nvars);
  for (varID = 0; varID < nvars; varID++)
    {
      vardata[varID].resize(varList1[varID].gridsize * varList1[varID].nlevels);
    }

  int tsID = 0;
  while ((nrecs = cdoStreamInqTimestep(streamID1, tsID)))
    {
      taxisCopyTimestep(taxisID2, taxisID1);
      cdoDefTimestep(streamID2, tsID);

      for (int recID = 0; recID < nrecs; recID++)
        {
          cdoInqRecord(streamID1, &varID, &levelID);

          if (tsID == 0)
            {
              varInfo[varID].varID = varID;
              varInfo[varID].code = vlistInqVarCode(vlistID1, varID);
              const auto iparam = vlistInqVarParam(vlistID1, varID);
              paramToString(iparam, varInfo[varID].param, sizeof(varInfo[varID].param));
              vlistInqVarName(vlistID1, varID, varInfo[varID].name);
              zaxisID = vlistInqVarZaxis(vlistID1, varID);
              varInfo[varID].levInfo[levelID].levelID = levelID;
              varInfo[varID].levInfo[levelID].level = cdoZaxisInqLevel(zaxisID, levelID);
            }

          const auto offset = varList1[varID].gridsize * levelID;
          auto single = &vardata[varID][offset];

          cdoReadRecord(streamID1, single, &nmiss);

          setNmiss(varID, levelID, nvars, varInfo, nmiss);
          // varInfo[varID].levInfo[levelID].nmiss = nmiss;
        }

      if (tsID == 0)
        {
          if (Options::cdoVerbose)
            for (vindex = 0; vindex < nvars; vindex++)
              {
                nlevs = varInfo[vindex].nlevs;
                for (lindex = 0; lindex < nlevs; ++lindex)
                  printf("sort in: %d %s %d %d %g\n", vindex, varInfo[vindex].name, varInfo[vindex].code, varInfo[vindex].nlevs,
                         varInfo[vindex].levInfo[lindex].level);
              }

          if (operatorID == SORTCODE)
            std::sort(varInfo.begin(), varInfo.end(), cmpvarcode);
          else if (operatorID == SORTPARAM)
            std::sort(varInfo.begin(), varInfo.end(), cmpvarparam);
          else if (operatorID == SORTNAME)
            std::sort(varInfo.begin(), varInfo.end(), cmpvarname);
          else if (operatorID == SORTLEVEL)
            {
              for (vindex = 0; vindex < nvars; vindex++)
                std::sort(varInfo[vindex].levInfo.begin(), varInfo[vindex].levInfo.end(), cmpvarlev);
            }

          if (Options::cdoVerbose)
            for (vindex = 0; vindex < nvars; vindex++)
              {
                nlevs = varInfo[vindex].nlevs;
                for (lindex = 0; lindex < nlevs; ++lindex)
                  printf("sort out: %d %s %d %d %g\n", vindex, varInfo[vindex].name, varInfo[vindex].code, varInfo[vindex].nlevs,
                         varInfo[vindex].levInfo[lindex].level);
              }
        }

      for (vindex = 0; vindex < nvars; vindex++)
        {
          varID = varInfo[vindex].varID;
          nlevs = varInfo[vindex].nlevs;
          for (lindex = 0; lindex < nlevs; ++lindex)
            {
              levelID = varInfo[vindex].levInfo[lindex].levelID;
              nmiss = varInfo[vindex].levInfo[lindex].nmiss;

              if (tsID == 0 || vlistInqVarTimetype(vlistID1, varID) != TIME_CONSTANT)
                {
                  const auto offset = varList1[varID].gridsize * levelID;
                  auto single = &vardata[varID][offset];

                  cdoDefRecord(streamID2, varID, levelID);
                  cdoWriteRecord(streamID2, single, nmiss);
                }
            }
        }

      tsID++;
    }

  cdoStreamClose(streamID1);
  cdoStreamClose(streamID2);

  cdoFinish();

  return nullptr;
}
