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

      Split      splitcode       Split codes
      Split      splitparam      Split parameters
      Split      splitname       Split variables
      Split      splitlevel      Split levels
      Split      splitgrid       Split grids
      Split      splitzaxis      Split zaxis
      Split      splittabnum     Split table numbers
*/

#include <cdi.h>

#include "process_int.h"
#include "cdo_history.h"
#include "util_files.h"
#include "cdo_zaxis.h"
#include "cdi_lockedIO.h"

#include <cassert>

static void
gen_filename(char *filename, bool swap_obase, std::string obase, const char *suffix)
{
  if (swap_obase) strcat(filename, obase.c_str());
  if (suffix[0]) strcat(filename, suffix);
}

static int
split_code(bool swap_obase, char *filesuffix, char *filename, int vlistID1, const VarList &varList1,
           std::vector<CdoStreamID> &streamIDs, std::vector<int> &vlistIDs)
{
  const auto nchars = strlen(filename);
  const auto nvars = vlistNvars(vlistID1);
  std::vector<int> codes(nvars);

  int nsplit = 0;
  for (int varID = 0; varID < nvars; varID++)
    {
      const auto code = varList1[varID].code;
      int index;
      for (index = 0; index < varID; index++)
        if (code == varList1[index].code) break;

      if (index == varID) codes[nsplit++] = code;
    }

  vlistIDs.resize(nsplit);
  streamIDs.resize(nsplit);

  for (int index = 0; index < nsplit; ++index)
    {
      vlistClearFlag(vlistID1);
      for (int varID = 0; varID < nvars; varID++)
        {
          const auto code = varList1[varID].code;
          if (codes[index] == code)
            {
              for (int levelID = 0; levelID < varList1[varID].nlevels; levelID++)
                {
                  vlistDefIndex(vlistID1, varID, levelID, index);
                  vlistDefFlag(vlistID1, varID, levelID, true);
                }
            }
        }

      auto vlistID2 = vlistCreate();
      cdo_vlist_copy_flag(vlistID2, vlistID1);
      vlistIDs[index] = vlistID2;

      const char *format = (codes[index] > 9999) ? "%05d" : ((codes[index] > 999) ? "%04d" : "%03d");
      sprintf(filename + nchars, format, codes[index]);
      gen_filename(filename, swap_obase, cdo_get_obase(), filesuffix);

      streamIDs[index] = cdo_open_write(filename);
    }

  return nsplit;
}

static int
split_param(bool swap_obase, char *filesuffix, char *filename, int vlistID1, const VarList &varList1,
            std::vector<CdoStreamID> &streamIDs, std::vector<int> &vlistIDs)
{
  const auto nchars = strlen(filename);
  char paramstr[32];
  const auto nvars = vlistNvars(vlistID1);
  std::vector<int> params(nvars);

  int nsplit = 0;
  for (int varID = 0; varID < nvars; varID++)
    {
      const auto param = varList1[varID].param;
      int index;
      for (index = 0; index < varID; index++)
        if (param == varList1[index].param) break;

      if (index == varID) params[nsplit++] = param;
    }

  vlistIDs.resize(nsplit);
  streamIDs.resize(nsplit);

  for (int index = 0; index < nsplit; ++index)
    {
      vlistClearFlag(vlistID1);
      for (int varID = 0; varID < nvars; varID++)
        {
          const auto param = varList1[varID].param;
          if (params[index] == param)
            {
              for (int levelID = 0; levelID < varList1[varID].nlevels; levelID++)
                {
                  vlistDefIndex(vlistID1, varID, levelID, index);
                  vlistDefFlag(vlistID1, varID, levelID, true);
                }
            }
        }

      auto vlistID2 = vlistCreate();
      cdo_vlist_copy_flag(vlistID2, vlistID1);
      vlistIDs[index] = vlistID2;

      cdiParamToString(params[index], paramstr, sizeof(paramstr));

      filename[nchars] = '\0';
      strcat(filename, paramstr);
      gen_filename(filename, swap_obase, cdo_get_obase(), filesuffix);

      streamIDs[index] = cdo_open_write(filename);
    }

  return nsplit;
}

static int
split_name(bool swap_obase, char *filesuffix, char *filename, int vlistID1, const VarList &varList1,
           std::vector<CdoStreamID> &streamIDs, std::vector<int> &vlistIDs)
{
  const auto nchars = strlen(filename);
  const auto nvars = vlistNvars(vlistID1);
  const auto nsplit = nvars;

  vlistIDs.resize(nsplit);
  streamIDs.resize(nsplit);

  for (int index = 0; index < nsplit; index++)
    {
      vlistClearFlag(vlistID1);
      int varID = index;
      for (int levelID = 0; levelID < varList1[varID].nlevels; levelID++)
        {
          vlistDefIndex(vlistID1, varID, levelID, index);
          vlistDefFlag(vlistID1, varID, levelID, true);
        }

      auto vlistID2 = vlistCreate();
      cdo_vlist_copy_flag(vlistID2, vlistID1);
      vlistIDs[index] = vlistID2;

      filename[nchars] = '\0';
      strcat(filename, varList1[varID].name);
      gen_filename(filename, swap_obase, cdo_get_obase(), filesuffix);

      streamIDs[index] = cdo_open_write(filename);
    }

  return nsplit;
}

static int
split_level(bool swap_obase, char *filesuffix, char *filename, int vlistID1, const VarList &varList1,
            std::vector<CdoStreamID> &streamIDs, std::vector<int> &vlistIDs)
{
  const auto nchars = strlen(filename);
  const auto nvars = vlistNvars(vlistID1);
  const auto nzaxis = vlistNzaxis(vlistID1);
  double ftmp[999];

  int nsplit = 0;
  for (int index = 0; index < nzaxis; index++)
    {
      const auto zaxisID = vlistZaxis(vlistID1, index);
      const auto nlevels = zaxisInqSize(zaxisID);
      for (int levelID = 0; levelID < nlevels; levelID++)
        {
          const auto level = cdo_zaxis_inq_level(zaxisID, levelID);
          int i;
          for (i = 0; i < nsplit; i++)
            if (IS_EQUAL(level, ftmp[i])) break;
          if (i == nsplit) ftmp[nsplit++] = level;
        }
    }

  vlistIDs.resize(nsplit);
  streamIDs.resize(nsplit);
  Varray<double> levels(nsplit);
  for (int index = 0; index < nsplit; ++index) levels[index] = ftmp[index];

  for (int index = 0; index < nsplit; ++index)
    {
      vlistClearFlag(vlistID1);
      for (int varID = 0; varID < nvars; varID++)
        {
          const auto zaxisID = varList1[varID].zaxisID;
          for (int levelID = 0; levelID < varList1[varID].nlevels; levelID++)
            {
              const auto level = cdo_zaxis_inq_level(zaxisID, levelID);
              if (IS_EQUAL(levels[index], level))
                {
                  vlistDefIndex(vlistID1, varID, levelID, index);
                  vlistDefFlag(vlistID1, varID, levelID, true);
                }
            }
        }

      auto vlistID2 = vlistCreate();
      cdo_vlist_copy_flag(vlistID2, vlistID1);
      vlistIDs[index] = vlistID2;

      sprintf(filename + nchars, "%06g", levels[index]);
      gen_filename(filename, swap_obase, cdo_get_obase(), filesuffix);

      streamIDs[index] = cdo_open_write(filename);
    }

  return nsplit;
}

static int
split_grid(bool swap_obase, char *filesuffix, char *filename, int vlistID1, const VarList &varList1,
           std::vector<CdoStreamID> &streamIDs, std::vector<int> &vlistIDs)
{
  const auto nchars = strlen(filename);
  const auto nsplit = vlistNgrids(vlistID1);
  const auto nvars = vlistNvars(vlistID1);

  vlistIDs.resize(nsplit);
  streamIDs.resize(nsplit);
  std::vector<int> gridIDs(nsplit);
  for (int index = 0; index < nsplit; ++index) gridIDs[index] = vlistGrid(vlistID1, index);

  for (int index = 0; index < nsplit; ++index)
    {
      vlistClearFlag(vlistID1);
      for (int varID = 0; varID < nvars; varID++)
        {
          const auto gridID = varList1[varID].gridID;
          if (gridIDs[index] == gridID)
            {
              for (int levelID = 0; levelID < varList1[varID].nlevels; levelID++)
                {
                  vlistDefIndex(vlistID1, varID, levelID, index);
                  vlistDefFlag(vlistID1, varID, levelID, true);
                }
            }
        }

      auto vlistID2 = vlistCreate();
      cdo_vlist_copy_flag(vlistID2, vlistID1);
      vlistIDs[index] = vlistID2;

      sprintf(filename + nchars, "%02d", vlistGridIndex(vlistID1, gridIDs[index]) + 1);
      gen_filename(filename, swap_obase, cdo_get_obase(), filesuffix);

      streamIDs[index] = cdo_open_write(filename);
    }

  return nsplit;
}

static int
split_zaxis(bool swap_obase, char *filesuffix, char *filename, int vlistID1, const VarList &varList1,
            std::vector<CdoStreamID> &streamIDs, std::vector<int> &vlistIDs)
{
  const auto nchars = strlen(filename);
  const auto nsplit = vlistNzaxis(vlistID1);
  const auto nvars = vlistNvars(vlistID1);

  vlistIDs.resize(nsplit);
  streamIDs.resize(nsplit);
  std::vector<int> zaxisIDs(nsplit);
  for (int index = 0; index < nsplit; ++index) zaxisIDs[index] = vlistZaxis(vlistID1, index);

  for (int index = 0; index < nsplit; ++index)
    {
      vlistClearFlag(vlistID1);
      for (int varID = 0; varID < nvars; varID++)
        {
          const auto zaxisID = varList1[varID].zaxisID;
          if (zaxisIDs[index] == zaxisID)
            {
              for (int levelID = 0; levelID < varList1[varID].nlevels; levelID++)
                {
                  vlistDefIndex(vlistID1, varID, levelID, index);
                  vlistDefFlag(vlistID1, varID, levelID, true);
                }
            }
        }

      auto vlistID2 = vlistCreate();
      cdo_vlist_copy_flag(vlistID2, vlistID1);
      vlistIDs[index] = vlistID2;

      sprintf(filename + nchars, "%02d", vlistZaxisIndex(vlistID1, zaxisIDs[index]) + 1);
      gen_filename(filename, swap_obase, cdo_get_obase(), filesuffix);

      streamIDs[index] = cdo_open_write(filename);
    }

  return nsplit;
}

static int
split_tabnum(bool swap_obase, char *filesuffix, char *filename, int vlistID1, const VarList &varList1,
             std::vector<CdoStreamID> &streamIDs, std::vector<int> &vlistIDs)
{
  const auto nchars = strlen(filename);
  const auto nvars = vlistNvars(vlistID1);
  std::vector<int> tabnums(nvars);

  int nsplit = 0;
  for (int varID = 0; varID < nvars; varID++)
    {
      const auto tabnum = tableInqNum(vlistInqVarTable(vlistID1, varID));
      int index;
      for (index = 0; index < varID; index++)
        if (tabnum == tableInqNum(vlistInqVarTable(vlistID1, index))) break;

      if (index == varID) tabnums[nsplit++] = tabnum;
    }

  vlistIDs.resize(nsplit);
  streamIDs.resize(nsplit);

  for (int index = 0; index < nsplit; ++index)
    {
      vlistClearFlag(vlistID1);
      for (int varID = 0; varID < nvars; varID++)
        {
          const auto tabnum = tableInqNum(vlistInqVarTable(vlistID1, varID));
          if (tabnums[index] == tabnum)
            {
              for (int levelID = 0; levelID < varList1[varID].nlevels; levelID++)
                {
                  vlistDefIndex(vlistID1, varID, levelID, index);
                  vlistDefFlag(vlistID1, varID, levelID, true);
                }
            }
        }

      auto vlistID2 = vlistCreate();
      cdo_vlist_copy_flag(vlistID2, vlistID1);
      vlistIDs[index] = vlistID2;

      sprintf(filename + nchars, "%03d", tabnums[index]);
      gen_filename(filename, swap_obase, cdo_get_obase(), filesuffix);

      streamIDs[index] = cdo_open_write(filename);
    }

  return nsplit;
}

void *
Split(void *process)
{
  cdo_initialize(process);

  if (process_self().m_ID != 0) cdo_abort("This operator can't be combined with other operators!");

  const auto dataIsUnchanged = data_is_unchanged();

  // clang-format off
  const auto SPLITCODE   = cdo_operator_add("splitcode",   0, 0, nullptr);
  const auto SPLITPARAM  = cdo_operator_add("splitparam",  0, 0, nullptr);
  const auto SPLITNAME   = cdo_operator_add("splitname",   0, 0, nullptr);
  const auto SPLITLEVEL  = cdo_operator_add("splitlevel",  0, 0, nullptr);
  const auto SPLITGRID   = cdo_operator_add("splitgrid",   0, 0, nullptr);
  const auto SPLITZAXIS  = cdo_operator_add("splitzaxis",  0, 0, nullptr);
  const auto SPLITTABNUM = cdo_operator_add("splittabnum", 0, 0, nullptr);
  // clang-format on

  const auto operatorID = cdo_operator_id();

  bool swap_obase = false;
  const char *uuid_attribute = nullptr;
  for (int i = 0; i < cdo_operator_argc(); ++i)
    {
      if (cdo_operator_argv(i) == "swap")
        swap_obase = true;
      else if (cdo_operator_argv(i).find("uuid=") == 0)
        uuid_attribute = &cdo_operator_argv(i)[0] + 5;
      else
        cdo_abort("Unknown parameter: >%s<", cdo_operator_argv(0).c_str());
    }

  const auto streamID1 = cdo_open_read(0);
  const auto vlistID1 = cdo_stream_inq_vlist(streamID1);

  VarList varList1;
  varListInit(varList1, vlistID1);

  char filename[8192] = { 0 };
  if (!swap_obase) strcpy(filename, cdo_get_obase().c_str());

  char filesuffix[32] = { 0 };
  FileUtils::gen_suffix(filesuffix, sizeof(filesuffix), cdo_inq_filetype(streamID1), vlistID1, cdo_get_stream_name(0));

  std::vector<int> vlistIDs;
  std::vector<CdoStreamID> streamIDs;

  int nsplit = 0;
  if (operatorID == SPLITCODE)
    {
      nsplit = split_code(swap_obase, filesuffix, filename, vlistID1, varList1, streamIDs, vlistIDs);
    }
  else if (operatorID == SPLITPARAM)
    {
      nsplit = split_param(swap_obase, filesuffix, filename, vlistID1, varList1, streamIDs, vlistIDs);
    }
  else if (operatorID == SPLITTABNUM)
    {
      nsplit = split_tabnum(swap_obase, filesuffix, filename, vlistID1, varList1, streamIDs, vlistIDs);
    }
  else if (operatorID == SPLITNAME)
    {
      nsplit = split_name(swap_obase, filesuffix, filename, vlistID1, varList1, streamIDs, vlistIDs);
    }
  else if (operatorID == SPLITLEVEL)
    {
      nsplit = split_level(swap_obase, filesuffix, filename, vlistID1, varList1, streamIDs, vlistIDs);
    }
  else if (operatorID == SPLITGRID)
    {
      nsplit = split_grid(swap_obase, filesuffix, filename, vlistID1, varList1, streamIDs, vlistIDs);
    }
  else if (operatorID == SPLITZAXIS)
    {
      nsplit = split_zaxis(swap_obase, filesuffix, filename, vlistID1, varList1, streamIDs, vlistIDs);
    }
  else
    {
      cdo_abort("not implemented!");
    }

  assert(nsplit > 0);

  for (int index = 0; index < nsplit; index++)
    {
      if (uuid_attribute) cdo_def_tracking_id(vlistIDs[index], uuid_attribute);

      cdo_def_vlist(streamIDs[index], vlistIDs[index]);
    }

  Field field;

  int tsID = 0;
  while (true)
    {
      const auto nrecs = cdo_stream_inq_timestep(streamID1, tsID);
      if (nrecs == 0) break;

      for (int index = 0; index < nsplit; index++) cdo_def_timestep(streamIDs[index], tsID);

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

          const auto index = vlistInqIndex(vlistID1, varID, levelID);
          const auto vlistID2 = vlistIDs[index];
          const auto varID2 = vlistFindVar(vlistID2, varID);
          const auto levelID2 = vlistFindLevel(vlistID2, varID, levelID);
          // printf("%d %d %d %d %d %d\n", index, vlistID2, varID, levelID, varID2, levelID2);

          cdo_def_record(streamIDs[index], varID2, levelID2);
          if (dataIsUnchanged)
            {
              cdo_copy_record(streamIDs[index], streamID1);
            }
          else
            {
              field.init(varList1[varID]);
              cdo_read_record(streamID1, field);
              cdo_write_record(streamIDs[index], field);
            }
        }

      tsID++;
    }

  cdo_stream_close(streamID1);

  for (auto &streamID : streamIDs) cdo_stream_close(streamID);
  for (auto &vlistID : vlistIDs) vlistDestroy(vlistID);

  cdo_finish();

  return nullptr;
}
