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

      Vertwind    vertwind      Convert the vertical velocity to [m/s]
*/

#include <cdi.h>

#include "process_int.h"
#include "cdo_vlist.h"
#include "cdi_lockedIO.h"
#include "vertical_interp.h"
#include "util_string.h"
#include "cdo_zaxis.h"

#define R 287.07  /* spezielle Gaskonstante fuer Luft */
#define G 9.80665 /* Erdbeschleunigung */

void *
Vertwind(void *process)
{
  int varID, levelID;
  int nvct = 0;
  size_t nmiss;
  int tempID = -1, sqID = -1, psID = -1, omegaID = -1;
  char varname[CDI_MAX_NAME];
  Varray<double> vct;
  Varray<double> hpress, psProg;

  cdo_initialize(process);

  operator_check_argc(0);

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

  vlist_check_gridsize(vlistID1);

  int temp_code = 130;
  int sq_code = 133;
  int ps_code = 134;
  int omega_code = 135;

  int nvars = vlistNvars(vlistID1);
  for (varID = 0; varID < nvars; ++varID)
    {
      int code = vlistInqVarCode(vlistID1, varID);

      if (code <= 0)
        {
          vlistInqVarName(vlistID1, varID, varname);
          cstr_to_lower_case(varname);

          // clang-format off
          if      (strcmp(varname, "st") == 0) code = temp_code;
          else if (strcmp(varname, "sq") == 0) code = sq_code;
          else if (strcmp(varname, "aps") == 0) code = ps_code;
          else if (strcmp(varname, "omega") == 0) code = omega_code;
          // clang-format on
        }

      // clang-format off
      if      (code == temp_code)  tempID = varID;
      else if (code == sq_code)    sqID = varID;
      else if (code == ps_code)    psID = varID;
      else if (code == omega_code) omegaID = varID;
      // clang-format on
    }

  if (tempID == -1 || sqID == -1 || omegaID == -1)
    {
      if (tempID == -1) cdo_warning("Temperature (code 130) not found!");
      if (sqID == -1) cdo_warning("Specific humidity (code 133) not found!");
      if (omegaID == -1) cdo_warning("Vertical velocity (code 135) not found!");
      cdo_abort("Parameter not found!");
    }

  /* Get missing values */
  const auto missval_t = vlistInqVarMissval(vlistID1, tempID);
  const auto missval_sq = vlistInqVarMissval(vlistID1, sqID);
  const auto missval_wap = vlistInqVarMissval(vlistID1, omegaID);
  const auto missval_out = missval_wap;

  const auto gridID = vlistInqVarGrid(vlistID1, omegaID);
  const auto zaxisID = vlistInqVarZaxis(vlistID1, omegaID);

  if (psID == -1 && zaxisInqType(zaxisID) == ZAXIS_HYBRID) cdo_abort("Surface pressure (code 134) not found!");

  const auto gridsize = gridInqSize(gridID);
  const auto nlevel = zaxisInqSize(zaxisID);
  Varray<double> level(nlevel);
  cdo_zaxis_inq_levels(zaxisID, level.data());

  Varray<double> temp(gridsize * nlevel);
  Varray<double> sq(gridsize * nlevel);
  Varray<double> omega(gridsize * nlevel);
  Varray<double> wms(gridsize * nlevel);
  Varray<double> fpress(gridsize * nlevel);

  if (zaxisInqType(zaxisID) == ZAXIS_PRESSURE)
    {
      for (levelID = 0; levelID < nlevel; ++levelID)
        {
          const auto offset = (size_t) levelID * gridsize;
          for (size_t i = 0; i < gridsize; ++i) fpress[offset + i] = level[levelID];
        }
    }
  else if (zaxisInqType(zaxisID) == ZAXIS_HYBRID)
    {
      psProg.resize(gridsize);
      hpress.resize(gridsize * (nlevel + 1));

      nvct = zaxisInqVctSize(zaxisID);
      if (nlevel == (nvct / 2 - 1))
        {
          vct.resize(nvct);
          zaxisInqVct(zaxisID, vct.data());
        }
      else
        cdo_abort("Unsupported vertical coordinate table format!");
    }
  else
    cdo_abort("Unsupported Z-Axis type!");

  vlistClearFlag(vlistID1);
  for (levelID = 0; levelID < nlevel; ++levelID) vlistDefFlag(vlistID1, omegaID, levelID, true);

  const auto vlistID2 = vlistCreate();
  cdo_vlist_copy_flag(vlistID2, vlistID1);
  vlistDefNtsteps(vlistID2, vlistNtsteps(vlistID1));

  vlistDefVarCode(vlistID2, 0, 40);
  cdiDefKeyString(vlistID2, 0, CDI_KEY_NAME, "W");
  cdiDefKeyString(vlistID2, 0, CDI_KEY_LONGNAME, "Vertical velocity");
  cdiDefKeyString(vlistID2, 0, CDI_KEY_UNITS, "m/s");
  vlistDefVarMissval(vlistID2, 0, missval_out);

  const auto taxisID1 = vlistInqTaxis(vlistID1);
  const auto taxisID2 = taxisDuplicate(taxisID1);
  vlistDefTaxis(vlistID2, taxisID2);

  const auto streamID2 = cdo_open_write(1);

  cdo_def_vlist(streamID2, vlistID2);

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

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

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

          const auto offset = (size_t) levelID * gridsize;

          if (varID == tempID)
            cdo_read_record(streamID1, &temp[offset], &nmiss);
          else if (varID == sqID)
            cdo_read_record(streamID1, &sq[offset], &nmiss);
          else if (varID == omegaID)
            cdo_read_record(streamID1, &omega[offset], &nmiss);
          else if (varID == psID && zaxisInqType(zaxisID) == ZAXIS_HYBRID)
            cdo_read_record(streamID1, psProg.data(), &nmiss);
        }

      if (zaxisInqType(zaxisID) == ZAXIS_HYBRID)
        vct_to_hybrid_pressure(fpress.data(), hpress.data(), vct.data(), psProg.data(), nlevel, gridsize);

      for (levelID = 0; levelID < nlevel; ++levelID)
        {
          const auto offset = (size_t) levelID * gridsize;

          for (size_t i = 0; i < gridsize; ++i)
            {
              if (DBL_IS_EQUAL(temp[offset + i], missval_t) || DBL_IS_EQUAL(omega[offset + i], missval_wap)
                  || DBL_IS_EQUAL(sq[offset + i], missval_sq))
                {
                  wms[offset + i] = missval_out;
                }
              else
                {
                  // Virtuelle Temperatur bringt die Feuchteabhaengigkeit hinein
                  const auto tv = temp[offset + i] * (1. + 0.608 * sq[offset + i]);

                  // Die Dichte erhaelt man nun mit der Gasgleichung rho=p/(R*tv) Level in Pa!
                  const auto rho = fpress[offset + i] / (R * tv);
                  /*
                    Nun daraus die Vertikalgeschwindigkeit im m/s, indem man die Vertikalgeschwindigkeit
                    in Pa/s durch die Erdbeschleunigung und die Dichte teilt
                  */
                  wms[offset + i] = omega[offset + i] / (G * rho);
                }
            }
        }

      for (levelID = 0; levelID < nlevel; ++levelID)
        {
          const auto offset = (size_t) levelID * gridsize;

          size_t nmiss_out = 0;
          for (size_t i = 0; i < gridsize; i++)
            if (DBL_IS_EQUAL(wms[offset + i], missval_out)) nmiss_out++;

          cdo_def_record(streamID2, 0, levelID);
          cdo_write_record(streamID2, &wms[offset], nmiss_out);
        }

      tsID++;
    }

  cdo_stream_close(streamID2);
  cdo_stream_close(streamID1);

  vlistDestroy(vlistID2);

  cdo_finish();

  return nullptr;
}
