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

#include <algorithm>
#include <cdi.h>

#include "cdo_options.h"
#include "functs.h"
#include "process_int.h"
#include "array.h"

static int
vfarsetnmiss(const size_t len, Varray<double> &array, const double missval)
{
  size_t nmiss = 0;

  if (std::isnan(missval))
    {
      for (size_t i = 0; i < len; i++)
        if (DBL_IS_EQUAL(array[i], missval) || array[i] < 0)
          {
            array[i] = missval;
            nmiss++;
          }
    }
  else
    {
      for (size_t i = 0; i < len; i++)
        if (IS_EQUAL(array[i], missval) || array[i] < 0)
          {
            array[i] = missval;
            nmiss++;
          }
    }

  return nmiss;
}

void
vfarcpy(Field &field1, const Field &field2)
{
  auto size1 = field1.size;
  auto size2 = field2.size;
  auto &array1 = field1.vec_d;
  const auto &array2 = field2.vec_d;
  const auto &array2f = field2.vec_f;

  if (size1 == 0) size1 = gridInqSize(field1.grid);
  if (size2 == 0) size2 = gridInqSize(field2.grid);

  const auto len = size1;
  if (len != size2) cdoAbort("Fields have different size (%s)", __func__);

  if (field2.memType == MemType::Float)
    for (size_t i = 0; i < len; i++) array1[i] = array2f[i];
  else
    for (size_t i = 0; i < len; i++) array1[i] = array2[i];
}

void
vfaradd(Field &field1, const Field &field2)
{
  const auto missval1 = field1.missval;
  const auto missval2 = field2.missval;
  auto &array1 = field1.vec_d;
  const auto &array2 = field2.vec_d;
  const auto &array2f = field2.vec_f;

  const auto len = field1.size;
  if (len != field2.size) cdoAbort("Fields have different size (%s)", __func__);

  if (field1.nmiss || field2.nmiss)
    {
      for (size_t i = 0; i < len; i++) array1[i] = ADDMN(array1[i], array2[i]);
      field1.nmiss = fieldNumMiss(field1);
    }
  else
    {
      if (field2.memType == MemType::Float)
        {
          for (size_t i = 0; i < len; i++) array1[i] += array2f[i];
        }
      else
        {
          arrayAddArray(len, array1.data(), array2.data());
        }
    }
}

void
vfarsum(Field &field1, const Field &field2)
{
  auto size1 = field1.size;
  auto size2 = field2.size;
  const auto missval1 = field1.missval;
  const auto missval2 = field2.missval;
  auto &array1 = field1.vec_d;
  const auto &array2 = field2.vec_d;
  const auto &array2f = field2.vec_f;

  if (size1 == 0) size1 = gridInqSize(field1.grid);
  if (size2 == 0) size2 = gridInqSize(field2.grid);

  const auto len = size1;

  if (len != size2) cdoAbort("Fields have different size (%s)", __func__);

  if (field1.nmiss || field2.nmiss)
    {
      if (field2.memType == MemType::Float)
        {
          for (size_t i = 0; i < len; i++)
            if (!DBL_IS_EQUAL(array2f[i], missval2))
              {
                if (!DBL_IS_EQUAL(array1[i], missval1))
                  array1[i] += array2f[i];
                else
                  array1[i] = array2f[i];
              }
        }
      else
        {
          arrayAddArrayMV(len, array1.data(), array2.data(), missval2);
        }

      field1.nmiss = fieldNumMiss(field1);
    }
  else
    {
      if (field2.memType == MemType::Float)
        {
          for (size_t i = 0; i < len; i++) array1[i] += array2f[i];
        }
      else
        {
          arrayAddArray(len, array1.data(), array2.data());
        }
    }
}

void
vfarsumw(Field &field1, const Field &field2, double w)
{
  const auto missval1 = field1.missval;
  const auto missval2 = field2.missval;
  auto &array1 = field1.vec_d;
  const auto &array2 = field2.vec_d;

  auto len = field1.size;
  if (len != field2.size) cdoAbort("Fields have different size (%s)", __func__);

  if (field1.nmiss || field2.nmiss)
    {
      for (size_t i = 0; i < len; i++)
        if (!DBL_IS_EQUAL(array2[i], missval2))
          {
            if (!DBL_IS_EQUAL(array1[i], missval1))
              array1[i] += w * array2[i];
            else
              array1[i] = w * array2[i];
          }

      field1.nmiss = fieldNumMiss(field1);
    }
  else
    {
#ifdef _OPENMP
#pragma omp parallel for default(none) shared(len, array1, array2, w)
#endif
      for (size_t i = 0; i < len; i++) array1[i] += w * array2[i];
    }
}

/*
 * Compute the occurrence of values in field, if they do not equal refval.
 * This can be used to compute the lengths of multiple periods in a timeseries.
 * Missing field values are handled like refval, i.e. they stop a running period.
 * If there is missing data in the occurence field, missing fields values do not
 * change anything (they do not start a non-period by setting occurrence to zero).
 */
void
vfarsumtr(Field &occur, const Field &field, const double refval)
{
  auto omissval = occur.missval;
  auto fmissval = field.missval;
  auto &oarray = occur.vec_d;
  const auto &farray = field.vec_d;

  auto len = occur.size;
  if (len != field.size) cdoAbort("Fields have different size (%s)", __func__);

  if (occur.nmiss || field.nmiss)
    {
#ifdef _OPENMP
#pragma omp parallel for default(shared) schedule(static)
#endif
      for (size_t i = 0; i < len; i++)
        if (!DBL_IS_EQUAL(farray[i], fmissval))
          {
            if (!DBL_IS_EQUAL(oarray[i], omissval))
              oarray[i] = (DBL_IS_EQUAL(farray[i], refval)) ? 0.0 : oarray[i] + 1.0;
            else
              oarray[i] = (DBL_IS_EQUAL(farray[i], refval)) ? 0.0 : 1.0;
          }
        else
          {
            if (!DBL_IS_EQUAL(oarray[i], omissval)) oarray[i] = 0.0;
          }

      occur.nmiss = fieldNumMiss(occur);
    }
  else
    {
#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
      for (size_t i = 0; i < len; i++) oarray[i] = (DBL_IS_EQUAL(farray[i], refval)) ? 0.0 : oarray[i] + 1.0;
    }
}

void
vfarsumq(Field &field1, const Field &field2)
{
  const auto missval1 = field1.missval;
  const auto missval2 = field2.missval;
  auto &array1 = field1.vec_d;
  const auto &array2 = field2.vec_d;
  const auto &array2f = field2.vec_f;

  const auto len = field1.size;
  if (len != field2.size) cdoAbort("Fields have different size (%s)", __func__);

  if (field1.nmiss || field2.nmiss)
    {
      for (size_t i = 0; i < len; i++)
        if (!DBL_IS_EQUAL(array2[i], missval2))
          {
            if (!DBL_IS_EQUAL(array1[i], missval1))
              array1[i] += array2[i] * array2[i];
            else
              array1[i] = array2[i] * array2[i];
          }

      field1.nmiss = fieldNumMiss(field1);
    }
  else
    {
      if (field2.memType == MemType::Float)
        {
          for (size_t i = 0; i < len; i++) array1[i] += ((double) array2f[i]) * array2f[i];
        }
      else
        {
          for (size_t i = 0; i < len; i++) array1[i] += array2[i] * array2[i];
        }
    }
}

void
vfarsumqw(Field &field1, const Field &field2, const double w)
{
  const auto missval1 = field1.missval;
  const auto missval2 = field2.missval;
  auto &array1 = field1.vec_d;
  const auto &array2 = field2.vec_d;

  const auto len = field1.size;
  if (len != field2.size) cdoAbort("Fields have different size (%s)", __func__);

  if (field1.nmiss || field2.nmiss)
    {
      for (size_t i = 0; i < len; i++)
        if (!DBL_IS_EQUAL(array2[i], missval2))
          {
            if (!DBL_IS_EQUAL(array1[i], missval1))
              array1[i] += w * array2[i] * array2[i];
            else
              array1[i] = w * array2[i] * array2[i];
          }

      field1.nmiss = fieldNumMiss(field1);
    }
  else
    {
      for (size_t i = 0; i < len; i++) array1[i] += w * array2[i] * array2[i];
    }
}

void
vfarsub(Field &field1, const Field &field2)
{
  const auto missval1 = field1.missval;
  const auto missval2 = field2.missval;
  auto &array1 = field1.vec_d;
  const auto &array2 = field2.vec_d;

  const auto len = field1.size;
  if (len != field2.size) cdoAbort("Fields have different size (%s)", __func__);

  if (field1.nmiss || field2.nmiss)
    {
      for (size_t i = 0; i < len; i++) array1[i] = SUBMN(array1[i], array2[i]);

      field1.nmiss = fieldNumMiss(field1);
    }
  else
    {
      for (size_t i = 0; i < len; i++) array1[i] -= array2[i];
    }
}

void
vfarmul(Field &field1, const Field &field2)
{
  const auto missval1 = field1.missval;
  const auto missval2 = field2.missval;
  auto &array1 = field1.vec_d;
  const auto &array2 = field2.vec_d;

  const auto len = field1.size;
  if (len != field2.size) cdoAbort("Fields have different size (%s)", __func__);

  if (field1.nmiss || field2.nmiss)
    {
      for (size_t i = 0; i < len; i++) array1[i] = MULMN(array1[i], array2[i]);

      field1.nmiss = fieldNumMiss(field1);
    }
  else
    {
      for (size_t i = 0; i < len; i++) array1[i] *= array2[i];
    }
}

void
vfardiv(Field &field1, const Field &field2)
{
  const auto missval1 = field1.missval;
  const auto missval2 = field2.missval;
  auto &array1 = field1.vec_d;
  const auto &array2 = field2.vec_d;

  const auto len = field1.size;
  if (len != field2.size) cdoAbort("Fields have different size (%s)", __func__);

  for (size_t i = 0; i < len; i++) array1[i] = DIVMN(array1[i], array2[i]);

  field1.nmiss = fieldNumMiss(field1);
}

void
vfaratan2(Field &field1, const Field &field2)
{
  const auto missval1 = field1.missval;
  const auto missval2 = field2.missval;
  auto &array1 = field1.vec_d;
  const auto &array2 = field2.vec_d;

  const auto len = field1.size;
  if (len != field2.size) cdoAbort("Fields have different size (%s)", __func__);

  for (size_t i = 0; i < len; i++)
    array1[i] = DBL_IS_EQUAL(array1[i], missval1) || DBL_IS_EQUAL(array2[i], missval2) ? missval1 : atan2(array1[i], array2[i]);

  field1.nmiss = fieldNumMiss(field1);
}

void
vfarsetmiss(Field &field1, const Field &field2)
{
  const auto missval1 = field1.missval;
  auto &array1 = field1.vec_d;
  const auto &array2 = field2.vec_d;

  const auto len = field1.size;
  if (len != field2.size) cdoAbort("Fields have different size (%s)", __func__);

  if (field1.nmiss)
    {
      for (size_t i = 0; i < len; i++) array1[i] = DBL_IS_EQUAL(array1[i], missval1) ? array2[i] : array1[i];

      field1.nmiss = fieldNumMiss(field1);
    }
}

void
vfarmin(Field &field1, const Field &field2)
{
  const auto missval1 = field1.missval;
  const auto missval2 = field2.missval;
  auto &array1 = field1.vec_d;
  const auto &array2 = field2.vec_d;

  const auto len = field1.size;
  if (len != field2.size) cdoAbort("Fields have different size (%s)", __func__);

  if (field1.nmiss || field2.nmiss)
    {
      for (size_t i = 0; i < len; i++)
        {
          array1[i] = DBL_IS_EQUAL(array2[i], missval2)
                          ? array1[i]
                          : DBL_IS_EQUAL(array1[i], missval1) ? array2[i] : std::min(array1[i], array2[i]);
        }

      field1.nmiss = fieldNumMiss(field1);
    }
  else
    {
      for (size_t i = 0; i < len; i++) array1[i] = std::min(array1[i], array2[i]);
    }
}

void
vfarmax(Field &field1, const Field &field2)
{
  const auto missval1 = field1.missval;
  const auto missval2 = field2.missval;
  auto &array1 = field1.vec_d;
  const auto &array2 = field2.vec_d;

  const auto len = field1.size;
  if (len != field2.size) cdoAbort("Fields have different size (%s)", __func__);

  if (field1.nmiss || field2.nmiss)
    {
      for (size_t i = 0; i < len; i++)
        {
          array1[i] = DBL_IS_EQUAL(array2[i], missval2)
                          ? array1[i]
                          : DBL_IS_EQUAL(array1[i], missval1) ? array2[i] : std::max(array1[i], array2[i]);
        }

      field1.nmiss = fieldNumMiss(field1);
    }
  else
    {
      for (size_t i = 0; i < len; i++) array1[i] = std::max(array1[i], array2[i]);
    }
}

void
vfarminidx(Field &field1, Field &field2, const Field &field3, const int idx)
{
  const auto missval1 = field1.missval;
  const auto missval2 = field2.missval;
  const auto missval3 = field3.missval;
  auto &array1 = field1.vec_d;
  auto &array2 = field2.vec_d;
  const auto &array3 = field3.vec_d;

  const auto len = field3.size;
  if (len != field3.size) cdoAbort("Fields have different size (%s)", __func__);

  if (field2.nmiss || field3.nmiss)
    {
      for (size_t i = 0; i < len; i++)
        {
          if (DBL_IS_EQUAL(array3[i], missval3))
            {
              if (DBL_IS_EQUAL(array2[i], missval2)) array1[i] = missval1;
            }
          else
            {
              if (DBL_IS_EQUAL(array2[i], missval2))
                {
                  array2[i] = array3[i];
                  array1[i] = idx;
                }
              else if (array3[i] < array2[i])
                {
                  array2[i] = array3[i];
                  array1[i] = idx;
                }
            }
        }

      field1.nmiss = fieldNumMiss(field1);
      field2.nmiss = fieldNumMiss(field2);
    }
  else
    {
      for (size_t i = 0; i < len; i++)
        {
          if (array3[i] < array2[i])
            {
              array2[i] = array3[i];
              array1[i] = idx;
            }
        }
    }
}

void
vfarmaxidx(Field &field1, Field &field2, const Field &field3, const int idx)
{
  const auto missval1 = field1.missval;
  const auto missval2 = field2.missval;
  const auto missval3 = field3.missval;
  auto &array1 = field1.vec_d;
  auto &array2 = field2.vec_d;
  const auto &array3 = field3.vec_d;

  const auto len = field2.size;
  if (len != field3.size) cdoAbort("Fields have different size (%s)", __func__);

  if (field2.nmiss || field3.nmiss)
    {
      for (size_t i = 0; i < len; i++)
        {
          if (DBL_IS_EQUAL(array3[i], missval3))
            {
              if (DBL_IS_EQUAL(array2[i], missval2)) array1[i] = missval1;
            }
          else
            {
              if (DBL_IS_EQUAL(array2[i], missval2))
                {
                  array2[i] = array3[i];
                  array1[i] = idx;
                }
              else if (array3[i] > array2[i])
                {
                  array2[i] = array3[i];
                  array1[i] = idx;
                }
            }
        }

      field1.nmiss = fieldNumMiss(field1);
      field2.nmiss = fieldNumMiss(field2);
    }
  else
    {
      for (size_t i = 0; i < len; i++)
        {
          if (array3[i] > array2[i])
            {
              array2[i] = array3[i];
              array1[i] = idx;
            }
        }
    }
}

void
vfarvar(Field &field1, const Field &field2, const Field &field3, const int divisor)
{
  const auto missval1 = field1.missval;
  const auto missval2 = field2.missval;
  auto &array1 = field1.vec_d;
  const auto &array2 = field2.vec_d;
  const auto &array3 = field3.vec_d;
  double temp;

  const auto len = field1.size;
  if (len != field2.size) cdoAbort("Fields have different size (%s)", __func__);

  if (field1.nmiss || field2.nmiss)
    {
      for (size_t i = 0; i < len; i++)
        {
          temp = DIVMN(MULMN(array1[i], array1[i]), array3[i]);
          array1[i] = DIVMN(SUBMN(array2[i], temp), array3[i] - divisor);
          if (array1[i] < 0 && array1[i] > -1.e-5) array1[i] = 0;
        }
    }
  else
    {
      for (size_t i = 0; i < len; i++)
        {
          temp = DIV(MUL(array1[i], array1[i]), array3[i]);
          array1[i] = DIV(SUB(array2[i], temp), array3[i] - divisor);
          if (array1[i] < 0 && array1[i] > -1.e-5) array1[i] = 0;
        }
    }

  field1.nmiss = vfarsetnmiss(len, array1, missval1);
}

void
vfarstd(Field &field1, const Field &field2, const Field &field3, const int divisor)
{
  const auto missval1 = field1.missval;
  auto &array1 = field1.vec_d;

  const auto len = field1.size;
  if (len != field2.size) cdoAbort("Fields have different size (%s)", __func__);

  vfarvar(field1, field2, field3, divisor);

  size_t nmiss = 0;
  for (size_t i = 0; i < len; i++)
    if (DBL_IS_EQUAL(array1[i], missval1) || array1[i] < 0)
      {
        array1[i] = missval1;
        nmiss++;
      }
    else
      {
        array1[i] = IS_NOT_EQUAL(array1[i], 0) ? std::sqrt(array1[i]) : 0;
      }
  field1.nmiss = nmiss;
}

void
vfarcvar(Field &field1, const Field &field2, const int nsets, const int divisor)
{
  const int nsetx = nsets - divisor;
  const auto missval1 = field1.missval;
  const auto missval2 = field2.missval;
  auto &array1 = field1.vec_d;
  const auto &array2 = field2.vec_d;
  double temp;

  const auto len = field1.size;
  if (len != field2.size) cdoAbort("Fields have different size (%s)", __func__);

  if (nsetx == 0)
    {
      for (size_t i = 0; i < len; i++) array1[i] = missval1;
    }
  else if (field1.nmiss || field2.nmiss)
    {
      for (size_t i = 0; i < len; i++)
        {
          temp = MULMN(array1[i], array1[i]) / nsets;
          array1[i] = SUBMN(array2[i], temp) / nsetx;
          if (array1[i] < 0 && array1[i] > -1.e-5) array1[i] = 0;
        }
    }
  else
    {
      for (size_t i = 0; i < len; i++)
        {
          temp = MUL(array1[i], array1[i]) / nsets;
          array1[i] = SUB(array2[i], temp) / nsetx;
          if (array1[i] < 0 && array1[i] > -1.e-5) array1[i] = 0;
        }
    }

  field1.nmiss = vfarsetnmiss(len, array1, missval1);
}

void
vfarcstd(Field &field1, const Field &field2, const int nsets, const int divisor)
{
  const auto missval1 = field1.missval;
  auto &array1 = field1.vec_d;

  const auto len = field1.size;
  if (len != field2.size) cdoAbort("Fields have different size (%s)", __func__);

  vfarcvar(field1, field2, nsets, divisor);

  size_t nmiss = 0;
  for (size_t i = 0; i < len; i++)
    if (DBL_IS_EQUAL(array1[i], missval1) || array1[i] < 0)
      {
        array1[i] = missval1;
        nmiss++;
      }
    else
      {
        array1[i] = IS_NOT_EQUAL(array1[i], 0) ? std::sqrt(array1[i]) : 0;
      }

  field1.nmiss = nmiss;
}

void
vfarmoq(Field &field1, const Field &field2)
{
  const auto missval1 = field1.missval;
  const auto missval2 = field2.missval;
  auto &array1 = field1.vec_d;
  const auto &array2 = field2.vec_d;

  const auto len = field1.size;
  if (len != field2.size) cdoAbort("Fields have different size (%s)", __func__);

  if (field2.nmiss)
    {
      for (size_t i = 0; i < len; i++)
        if (!DBL_IS_EQUAL(array2[i], missval2))
          array1[i] = array2[i] * array2[i];
        else
          array1[i] = missval1;

      field1.nmiss = fieldNumMiss(field1);
    }
  else
    {
      for (size_t i = 0; i < len; i++) array1[i] = array2[i] * array2[i];
    }
}

void
vfarmoqw(Field &field1, const Field &field2, const double w)
{
  const auto missval1 = field1.missval;
  const auto missval2 = field2.missval;
  auto &array1 = field1.vec_d;
  const auto &array2 = field2.vec_d;

  const auto len = field1.size;
  if (len != field2.size) cdoAbort("Fields have different size (%s)", __func__);

  if (field2.nmiss)
    {
      for (size_t i = 0; i < len; i++)
        if (!DBL_IS_EQUAL(array2[i], missval2))
          array1[i] = w * array2[i] * array2[i];
        else
          array1[i] = missval1;

      field1.nmiss = fieldNumMiss(field1);
    }
  else
    {
      for (size_t i = 0; i < len; i++) array1[i] = w * array2[i] * array2[i];
    }
}

/**
 * Counts the number of nonmissing values. The result of the operation
 * is computed according to the following rules:
 *
 * field1  field2  result
 * a       b       a + 1
 * a       miss    a
 * miss    b       1
 * miss    miss    miss
 *
 * @param field1 the 1st input field, also holds the result
 * @param field2 the 2nd input field
 */
void
vfarcount(Field &field1, const Field &field2)
{
  const auto missval1 = field1.missval;
  const auto missval2 = field2.missval;
  auto &array1 = field1.vec_d;
  const auto &array2 = field2.vec_d;

  const auto len = field1.size;
  if (len != field2.size) cdoAbort("Fields have different size (%s)", __func__);

  if (field1.nmiss || field2.nmiss)
    {
      for (size_t i = 0; i < len; i++)
        if (!DBL_IS_EQUAL(array2[i], missval2))
          {
            if (!DBL_IS_EQUAL(array1[i], missval1))
              array1[i] += 1.0;
            else
              array1[i] = 1.0;
          }

      field1.nmiss = fieldNumMiss(field1);
    }
  else
    {
      for (size_t i = 0; i < len; i++) array1[i] += 1.0;
    }
}

void
vfarfun(Field &field1, const Field &field2, int function)
{
  // clang-format off
  switch (function)
    {
    case func_add:     vfaradd(field1, field2);   break;
    case func_min:     vfarmin(field1, field2);   break;
    case func_max:     vfarmax(field1, field2);   break;
    case func_sum:     vfarsum(field1, field2);   break;
    case func_mean:    vfarsum(field1, field2);   break;
    case func_avg:     vfaradd(field1, field2);   break;
    case func_sub:     vfarsub(field1, field2);   break;
    case func_mul:     vfarmul(field1, field2);   break;
    case func_div:     vfardiv(field1, field2);   break;
    case func_atan2:   vfaratan2(field1, field2); break;
    case func_setmiss: vfarsetmiss(field1, field2); break;
    default: cdoAbort("%s: function %d not implemented!", __func__, function);
    }
  // clang-format on
}
