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

      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"

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

void *
Split(void *process)
{
  int nchars = 0;
  int varID;
  int levelID, levID;
  int varID2, levelID2;
  int vlistID2;
  int itmp[999];
  double ftmp[999];
  int nsplit = 0;
  size_t nmiss;
  bool swap_obase = false;
  const char *uuid_attribute = nullptr;
  std::vector<int> vlistIDs;
  std::vector<CdoStreamID> streamIDs;

  cdoInitialize(process);

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

  const auto lcopy = unchangedRecord();

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

  const auto operatorID = cdoOperatorID();

  for (int i = 0; i < operatorArgc(); ++i)
    {
      if (cdoOperatorArgv(i) == "swap")
        swap_obase = true;
      else if (cdoOperatorArgv(i).find("uuid=") == 0)
        uuid_attribute = &cdoOperatorArgv(i)[0] + 5;
      else
        cdoAbort("Unknown parameter: >%s<", cdoOperatorArgv(0).c_str());
    }

  const auto streamID1 = cdoOpenRead(0);
  const auto vlistID1 = cdoStreamInqVlist(streamID1);

  const auto nvars = vlistNvars(vlistID1);

  char filename[8192];
  if (!swap_obase)
    {
      strcpy(filename, cdoGetObase());
      nchars = strlen(filename);
    }

  auto refname = cdoGetStreamName(0);
  char filesuffix[32] = { 0 };
  cdoGenFileSuffix(filesuffix, sizeof(filesuffix), cdoInqFiletype(streamID1), vlistID1, refname);

  if (operatorID == SPLITCODE)
    {
      nsplit = 0;
      for (varID = 0; varID < nvars; varID++)
        {
          const auto code = vlistInqVarCode(vlistID1, varID);
          int index;
          for (index = 0; index < varID; index++)
            if (code == vlistInqVarCode(vlistID1, index)) break;

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

      vlistIDs.resize(nsplit);
      streamIDs.resize(nsplit);
      std::vector<int> codes(nsplit);
      for (int index = 0; index < nsplit; ++index) codes[index] = itmp[index];

      for (int index = 0; index < nsplit; ++index)
        {
          vlistClearFlag(vlistID1);
          for (varID = 0; varID < nvars; varID++)
            {
              const auto code = vlistInqVarCode(vlistID1, varID);
              if (codes[index] == code)
                {
                  const auto nlevs = zaxisInqSize(vlistInqVarZaxis(vlistID1, varID));
                  for (levID = 0; levID < nlevs; levID++)
                    {
                      vlistDefIndex(vlistID1, varID, levID, index);
                      vlistDefFlag(vlistID1, varID, levID, true);
                    }
                }
            }

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

          if (codes[index] > 9999)
            {
              sprintf(filename + nchars, "%05d", codes[index]);
              gen_filename(filename, swap_obase, cdoGetObase(), filesuffix);
            }
          else if (codes[index] > 999)
            {
              sprintf(filename + nchars, "%04d", codes[index]);
              gen_filename(filename, swap_obase, cdoGetObase(), filesuffix);
            }
          else
            {
              sprintf(filename + nchars, "%03d", codes[index]);
              gen_filename(filename, swap_obase, cdoGetObase(), filesuffix);
            }

          streamIDs[index] = cdoOpenWrite(filename);
        }
    }
  else if (operatorID == SPLITPARAM)
    {
      char paramstr[32];
      nsplit = 0;
      for (varID = 0; varID < nvars; varID++)
        {
          const auto param = vlistInqVarParam(vlistID1, varID);
          int index;
          for (index = 0; index < varID; index++)
            if (param == vlistInqVarParam(vlistID1, index)) break;

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

      vlistIDs.resize(nsplit);
      streamIDs.resize(nsplit);
      std::vector<int> params(nsplit);
      for (int index = 0; index < nsplit; ++index) params[index] = itmp[index];

      for (int index = 0; index < nsplit; ++index)
        {
          vlistClearFlag(vlistID1);
          for (varID = 0; varID < nvars; varID++)
            {
              const auto param = vlistInqVarParam(vlistID1, varID);
              if (params[index] == param)
                {
                  const auto nlevs = zaxisInqSize(vlistInqVarZaxis(vlistID1, varID));
                  for (levID = 0; levID < nlevs; levID++)
                    {
                      vlistDefIndex(vlistID1, varID, levID, index);
                      vlistDefFlag(vlistID1, varID, levID, true);
                    }
                }
            }

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

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

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

          streamIDs[index] = cdoOpenWrite(filename);
        }
    }
  else if (operatorID == SPLITTABNUM)
    {
      nsplit = 0;
      for (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)
            {
              itmp[nsplit] = tabnum;
              nsplit++;
            }
        }

      vlistIDs.resize(nsplit);
      streamIDs.resize(nsplit);
      std::vector<int> tabnums(nsplit);
      for (int index = 0; index < nsplit; ++index) tabnums[index] = itmp[index];

      for (int index = 0; index < nsplit; ++index)
        {
          vlistClearFlag(vlistID1);
          for (varID = 0; varID < nvars; varID++)
            {
              const auto tabnum = tableInqNum(vlistInqVarTable(vlistID1, varID));
              if (tabnums[index] == tabnum)
                {
                  const auto nlevs = zaxisInqSize(vlistInqVarZaxis(vlistID1, varID));
                  for (levID = 0; levID < nlevs; levID++)
                    {
                      vlistDefIndex(vlistID1, varID, levID, index);
                      vlistDefFlag(vlistID1, varID, levID, true);
                    }
                }
            }
          vlistID2 = vlistCreate();
          cdoVlistCopyFlag(vlistID2, vlistID1);
          vlistIDs[index] = vlistID2;

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

          streamIDs[index] = cdoOpenWrite(filename);
        }
    }
  else if (operatorID == SPLITNAME)
    {
      char varname[CDI_MAX_NAME];
      nsplit = nvars;

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

      for (int index = 0; index < nsplit; index++)
        {
          vlistClearFlag(vlistID1);
          varID = index;
          const auto nlevs = zaxisInqSize(vlistInqVarZaxis(vlistID1, varID));
          for (levID = 0; levID < nlevs; levID++)
            {
              vlistDefIndex(vlistID1, varID, levID, index);
              vlistDefFlag(vlistID1, varID, levID, true);
            }

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

          filename[nchars] = '\0';
          vlistInqVarName(vlistID1, varID, varname);
          strcat(filename, varname);
          gen_filename(filename, swap_obase, cdoGetObase(), filesuffix);

          streamIDs[index] = cdoOpenWrite(filename);
        }
    }
  else if (operatorID == SPLITLEVEL)
    {
      const auto nzaxis = vlistNzaxis(vlistID1);
      nsplit = 0;
      for (int index = 0; index < nzaxis; index++)
        {
          const auto zaxisID = vlistZaxis(vlistID1, index);
          const auto nlevs = zaxisInqSize(zaxisID);
          for (levID = 0; levID < nlevs; levID++)
            {
              const auto level = cdoZaxisInqLevel(zaxisID, levID);
              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 (varID = 0; varID < nvars; varID++)
            {
              const auto zaxisID = vlistInqVarZaxis(vlistID1, varID);
              const auto nlevs = zaxisInqSize(zaxisID);
              for (levID = 0; levID < nlevs; levID++)
                {
                  const auto level = cdoZaxisInqLevel(zaxisID, levID);
                  if (IS_EQUAL(levels[index], level))
                    {
                      vlistDefIndex(vlistID1, varID, levID, index);
                      vlistDefFlag(vlistID1, varID, levID, true);
                    }
                }
            }
          vlistID2 = vlistCreate();
          cdoVlistCopyFlag(vlistID2, vlistID1);
          vlistIDs[index] = vlistID2;

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

          streamIDs[index] = cdoOpenWrite(filename);
        }
    }
  else if (operatorID == SPLITGRID)
    {
      int gridID;

      nsplit = vlistNgrids(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 (varID = 0; varID < nvars; varID++)
            {
              gridID = vlistInqVarGrid(vlistID1, varID);
              if (gridIDs[index] == gridID)
                {
                  const auto nlevs = zaxisInqSize(vlistInqVarZaxis(vlistID1, varID));
                  for (levID = 0; levID < nlevs; levID++)
                    {
                      vlistDefIndex(vlistID1, varID, levID, index);
                      vlistDefFlag(vlistID1, varID, levID, true);
                    }
                }
            }
          vlistID2 = vlistCreate();
          cdoVlistCopyFlag(vlistID2, vlistID1);
          vlistIDs[index] = vlistID2;

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

          streamIDs[index] = cdoOpenWrite(filename);
        }
    }
  else if (operatorID == SPLITZAXIS)
    {
      nsplit = vlistNzaxis(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 (varID = 0; varID < nvars; varID++)
            {
              const auto zaxisID = vlistInqVarZaxis(vlistID1, varID);
              if (zaxisIDs[index] == zaxisID)
                {
                  const auto nlevs = zaxisInqSize(zaxisID);
                  for (levID = 0; levID < nlevs; levID++)
                    {
                      vlistDefIndex(vlistID1, varID, levID, index);
                      vlistDefFlag(vlistID1, varID, levID, true);
                    }
                }
            }
          vlistID2 = vlistCreate();
          cdoVlistCopyFlag(vlistID2, vlistID1);
          vlistIDs[index] = vlistID2;

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

          streamIDs[index] = cdoOpenWrite(filename);
        }
    }
  else
    {
      cdoAbort("not implemented!");
    }

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

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

  Varray<double> array;
  if (!lcopy)
    {
      auto gridsizemax = vlistGridsizeMax(vlistID1);
      if (vlistNumber(vlistID1) != CDI_REAL) gridsizemax *= 2;
      array.resize(gridsizemax);
    }

  int nrecs;
  int tsID = 0;
  while ((nrecs = cdoStreamInqTimestep(streamID1, tsID)))
    {
      for (int index = 0; index < nsplit; index++) cdoDefTimestep(streamIDs[index], tsID);

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

          const auto index = vlistInqIndex(vlistID1, varID, levelID);
          vlistID2 = vlistIDs[index];
          varID2 = vlistFindVar(vlistID2, varID);
          levelID2 = vlistFindLevel(vlistID2, varID, levelID);
          // printf("%d %d %d %d %d %d\n", index, vlistID2, varID, levelID, varID2, levelID2);
          cdoDefRecord(streamIDs[index], varID2, levelID2);
          if (lcopy)
            {
              cdoCopyRecord(streamIDs[index], streamID1);
            }
          else
            {
              cdoReadRecord(streamID1, array.data(), &nmiss);
              cdoWriteRecord(streamIDs[index], array.data(), nmiss);
            }
        }

      tsID++;
    }

  cdoStreamClose(streamID1);

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

  cdoFinish();

  return nullptr;
}
