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

  Author: Uwe Schulzweida

*/

/*
   This module contains the following operators:

      Sort sortcode  Sort by code number
*/

#include <cdi.h>

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

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

struct VarInfo
{
  int varID;
  int nlevels;
  int code;
  std::string param;
  std::string name;
  std::vector<LevInfo> levInfo;
};

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

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

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

  auto nlevels = varsInfo[vindex].nlevels;
  for (lindex = 0; lindex < nlevels; lindex++)
    if (varsInfo[vindex].levInfo[lindex].levelID == levelID) break;

  if (lindex == nlevels) cdo_abort("Internal problem; levelID not found!");

  varsInfo[vindex].levInfo[lindex].numMissVals = numMissVals;
}

class Sort : public Process
{
public:
  using Process::Process;
  inline static CdoModule module = {
    .name = "Sort",
    .operators = { { "sortcode" }, { "sortparam" }, { "sortname" }, { "sortlevel" } },
    .aliases = { { "sortvar", "sortname" } },
    .mode = EXPOSED,     // Module mode: 0:intern 1:extern
    .number = CDI_REAL,  // Allowed number type
    .constraints = { 1, 1, NoRestriction },
  };
  inline static RegisterEntry<Sort> registration = RegisterEntry<Sort>(module);

  int SORTCODE, SORTPARAM, SORTNAME, SORTLEVEL;
  bool compareLess = true;

  CdoStreamID streamID1;
  CdoStreamID streamID2;
  int taxisID1;
  int taxisID2;
  int nvars;

  VarList varList1;
  std::vector<VarInfo> varsInfo;
  Varray2D<double> vardata;
  int operatorID;

public:
  void
  init() override
  {
    SORTCODE = module.get_id("sortcode");
    SORTPARAM = module.get_id("sortparam");
    SORTNAME = module.get_id("sortname");
    SORTLEVEL = module.get_id("sortlevel");

    operatorID = cdo_operator_id();

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

    if (operatorID == SORTLEVEL && cdo_operator_argc() == 1)
      {
        auto iarg = parameter_to_int(cdo_operator_argv(0));
        if (iarg < 0) compareLess = false;
      }

    streamID1 = cdo_open_read(0);

    auto vlistID1 = cdo_stream_inq_vlist(streamID1);
    auto vlistID2 = vlistDuplicate(vlistID1);

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

    streamID2 = cdo_open_write(1);
    cdo_def_vlist(streamID2, vlistID2);

    varList1 = VarList(vlistID1);

    nvars = vlistNvars(vlistID1);

    varsInfo = std::vector<VarInfo>(nvars);
    for (int varID = 0; varID < nvars; ++varID)
      {
        const auto &var = varList1.vars[varID];
        varsInfo[varID].nlevels = var.nlevels;
        varsInfo[varID].levInfo.resize(var.nlevels);
      }

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

  void
  run() override
  {
    int tsID = 0;
    while (true)
      {
        auto numFields = cdo_stream_inq_timestep(streamID1, tsID);
        if (numFields == 0) break;

        cdo_taxis_copy_timestep(taxisID2, taxisID1);
        cdo_def_timestep(streamID2, tsID);

        for (int fieldID = 0; fieldID < numFields; ++fieldID)
          {
            auto [varID, levelID] = cdo_inq_field(streamID1);
            const auto &var = varList1.vars[varID];

            if (tsID == 0)
              {
                auto &varInfo = varsInfo[varID];
                varInfo.varID = varID;
                varInfo.code = var.code;
                varInfo.param = param_to_string(var.param);
                varInfo.name = var.name;
                varInfo.levInfo[levelID].levelID = levelID;
                varInfo.levInfo[levelID].level = cdo_zaxis_inq_level(var.zaxisID, levelID);
              }

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

            size_t numMissVals;
            cdo_read_field(streamID1, single, &numMissVals);

            setNmiss(varID, levelID, nvars, varsInfo, numMissVals);
            // varsInfo[varID].levInfo[levelID].numMissVals = numMissVals;
          }

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

            if (operatorID == SORTCODE)
              ranges::sort(varsInfo, {}, &VarInfo::code);
            else if (operatorID == SORTPARAM)
              ranges::sort(varsInfo, {}, &VarInfo::param);
            else if (operatorID == SORTNAME)
              ranges::sort(varsInfo, {}, &VarInfo::name);
            else if (operatorID == SORTLEVEL)
              {
                for (int vindex = 0; vindex < nvars; vindex++)
                  {
                    if (compareLess)
                      ranges::sort(varsInfo[vindex].levInfo, ranges::less(), &LevInfo::level);
                    else
                      ranges::sort(varsInfo[vindex].levInfo, ranges::greater(), &LevInfo::level);
                  }
              }

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

        for (int vindex = 0; vindex < nvars; vindex++)
          {
            const auto &varInfo = varsInfo[vindex];
            const auto &var = varList1.vars[varInfo.varID];
            for (int lindex = 0; lindex < var.nlevels; ++lindex)
              {
                auto levelID = varInfo.levInfo[lindex].levelID;
                auto numMissVals = varInfo.levInfo[lindex].numMissVals;

                if (tsID == 0 || !var.isConstant)
                  {
                    auto offset = var.gridsize * levelID;
                    auto single = &vardata[varInfo.varID][offset];

                    cdo_def_field(streamID2, varInfo.varID, levelID);
                    cdo_write_field(streamID2, single, numMissVals);
                  }
              }
          }

        tsID++;
      }
  }

  void
  close() override
  {
    cdo_stream_close(streamID1);
    cdo_stream_close(streamID2);
  }
};
