/*
  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 <cmath>
#include <string.h>

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include "cdo_options.h"
#include "cimdOmp.h"
#include "constants.h"
#include "cimdOmp.h"
#include "mpmo.h"
#include "array.h"

extern "C" void gaussianLatitudes(double *latitudes, double *weights, size_t nlat);

static void
jspleg1(double *pleg, double plat, long ktrunc, double *work)
{
  /*
     jspleg1 - Routine to calculate legendre functions

     Purpose
     --------

     This routine calculates the legendre functions for one latitude.
     (but not their derivatives)


     Interface
     ----------

     jspleg1( pleg, plat, ktrunc)


     Input parameters
     ----------------

     plat      - Latitude in radians
     ktrunc    - Spectral truncation


     Output parameters
     -----------------

     pleg      - Array of legendre functions for one latitude.
                 The array must be at least (KTRUNC+1)*(KTRUNC+4)/2
                 words long.

     Method
     ------

     Recurrence relation with explicit relations for P(m,m) and
     P(m,m+1)


     AUTHOR
     ------

     J.D.Chambers         ECMWF        9 November 1993


     Modifications
     -------------

     None

  */
  long itout1, i1m, ilm, jm, jcn, im2;
  double zsin, zcos, zf1m, zre1, zf2m, znsqr, ze1, ze2;
  double zjmsqr;
  double *zhlp1, *zhlp2, *zhlp3;

  /* Initialization */

  itout1 = ktrunc + 1;
  /*  zsin   = std::sin(plat); */
  zsin = plat;
  zcos = std::sqrt(1. - zsin * zsin);

  zhlp1 = work;
  zhlp2 = work + itout1;
  zhlp3 = work + itout1 + itout1;

  /*  Step 1.        M = 0, N = 0 and N = 1 */

  ilm = 1;
  pleg[0] = 1.0;
  zf1m = std::sqrt(3.0);
  pleg[1] = zf1m * zsin;

  /*  Step 2.       Sum for M = 0 to T (T = truncation) */

  for (jm = 1; jm < itout1; jm++)
    {
      zhlp1[jm] = std::sqrt(2. * jm + 3.);
      zhlp2[jm] = 1. / std::sqrt(2. * jm);
    }

  zhlp1[0] = std::sqrt(3.);

  for (jm = 0; jm < itout1; jm++)
    {
      i1m = jm - 1;
      zre1 = zhlp1[jm];
      ze1 = 1. / zre1;

      /*   Step 3.       M > 0 only */

      if (i1m != -1)
        {
          zf2m = zf1m * zcos * zhlp2[jm];
          zf1m = zf2m * zre1;

          /*  Step 4.       N = M and N = M+1 */

          ilm = ilm + 1;
          pleg[ilm] = zf2m;
          ilm = ilm + 1;
          pleg[ilm] = zf1m * zsin;

          /* When output truncation is reached, return to calling program */

          if (jm == (itout1 - 1)) break;
        }

      /*  Step 5.       Sum for N = M+2 to T+1 */

      zjmsqr = jm * jm;
      im2 = i1m + 2;

      for (jcn = im2; jcn < itout1; jcn++)
        {
          znsqr = (jcn + 1) * (jcn + 1);
          zhlp3[jcn] = std::sqrt((4. * znsqr - 1.) / (znsqr - zjmsqr));
        }

      for (jcn = im2; jcn < itout1; jcn++)
        {
          ze2 = zhlp3[jcn];
          ilm = ilm + 1;
          pleg[ilm] = ze2 * (zsin * pleg[ilm - 1] - ze1 * pleg[ilm - 2]);
          ze1 = 1. / ze2;
        }
    }
}

/* ============================================= */
/* phcs - Compute values of Legendre polynomials */
/*        and their meridional derivatives       */
/* ============================================= */
static void
phcs(double *pnm, double *hnm, long waves, double pmu, double *ztemp1, double *ztemp2)
{
  long twowaves;

  long jk, jn, jm;

  double jnmjk;
  double zcos2;
  double lat;
  double zan;
  double zsinpar;
  double zcospar;
  double zsqp;
  double zcosfak;
  double zsinfak;
  double zq;
  double zwm2;
  double zw;
  double zwq;
  double zq2m1;
  double zwm2q2;
  double z2q2;
  double zcnm;
  double zdnm;
  double zenm;

  twowaves = waves << 1;

  zcos2 = std::sqrt(1.0 - pmu * pmu);
  lat = std::acos(pmu);
  zan = 1.0;

  ztemp1[0] = 0.5;

  for (jn = 1; jn < twowaves; jn++)
    {
      zsqp = 1.0 / std::sqrt((double) (jn + jn * jn));
      zan *= std::sqrt(1.0 - 1.0 / (4 * jn * jn));

      zcospar = std::cos(lat * jn);
      zsinpar = std::sin(lat * jn) * jn * zsqp;
      zcosfak = 1.0;

      for (jk = 2; jk < jn; jk += 2)
        {
          jnmjk = jn - jk;
          zcosfak *= (jk - 1.0) * (jn + jnmjk + 2.0) / (jk * (jn + jnmjk + 1.0));
          zsinfak = zcosfak * (jnmjk) *zsqp;
          zcospar += zcosfak * std::cos(lat * jnmjk);
          zsinpar += zsinfak * std::sin(lat * jnmjk);
        }

      /*  code for jk == jn */

      if ((jn & 1) == 0)
        {
          zcosfak *= (double) ((jn - 1) * (jn + 2)) / (double) (jn * (jn + 1));
          zcospar += zcosfak * 0.5;
        }
      ztemp1[jn] = zan * zcospar;
      ztemp2[jn - 1] = zan * zsinpar;
    }

  memcpy(pnm, ztemp1, waves * sizeof(double));
  pnm += waves;
  memcpy(pnm, ztemp2, waves * sizeof(double));
  pnm += waves;

  hnm[0] = 0.0;
  for (jn = 1; jn < waves; jn++) hnm[jn] = jn * (pmu * ztemp1[jn] - std::sqrt((jn + jn + 1.0) / (jn + jn - 1.0)) * ztemp1[jn - 1]);

  hnm += waves;

  hnm[0] = pmu * ztemp2[0];

  for (jn = 1; jn < waves; jn++)
    hnm[jn] = (jn + 1) * pmu * ztemp2[jn]
              - std::sqrt(((jn + jn + 3.0) * ((jn + 1) * (jn + 1) - 1.0)) / (jn + jn + 1.0)) * ztemp2[jn - 1];

  hnm += waves;

  for (jm = 2; jm < waves; jm++)
    {
      pnm[0] = std::sqrt(1.0 + 1.0 / (jm + jm)) * zcos2 * ztemp2[0];
      hnm[0] = jm * pmu * pnm[0];
#if defined(CRAY)
#pragma _CRI novector
#endif
#if defined(__uxp__)
#pragma loop scalar
#endif
      for (jn = 1; jn < twowaves - jm; jn++)
        {
          zq = jm + jm + jn - 1;
          zwm2 = zq + jn;
          zw = zwm2 + 2;
          zwq = zw * zq;
          zq2m1 = zq * zq - 1.;
          zwm2q2 = zwm2 * zq2m1;
          z2q2 = zq2m1 * 2;
          zcnm = std::sqrt((zwq * (zq - 2.)) / (zwm2q2 - z2q2));
          zdnm = std::sqrt((zwq * (jn + 1.)) / zwm2q2);
          zenm = std::sqrt(zw * jn / ((zq + 1.0) * zwm2));
          pnm[jn] = zcnm * ztemp1[jn] - pmu * (zdnm * ztemp1[jn + 1] - zenm * pnm[jn - 1]);
          hnm[jn] = (jm + jn) * pmu * pnm[jn] - std::sqrt(zw * jn * (zq + 1) / zwm2) * pnm[jn - 1];
        }
      memcpy(ztemp1, ztemp2, twowaves * sizeof(double));
      memcpy(ztemp2, pnm, twowaves * sizeof(double));
      pnm += waves;
      hnm += waves;
    }
}

void
after_legini_full(long ntr, long nlat, double *restrict poli, double *restrict pold, double *restrict pdev, double *restrict pol2,
                  double *restrict pol3, double *restrict coslat)
{
  long waves = ntr + 1; // used in omp loop
  long dimsp = (ntr + 1) * (ntr + 2); // used in omp loop

  Varray<double> gmu(nlat), gwt(nlat);
  gaussianLatitudes(gmu.data(), gwt.data(), nlat);

#ifndef _OPENMP
  Varray<double> pnm(dimsp), hnm(dimsp), ztemp1(waves << 1), ztemp2(waves << 1);
#endif

#ifdef _OPENMP
#pragma omp parallel for default(none) shared(nlat, waves, dimsp, poli, pold, pdev, pol2, pol3, coslat, gmu, gwt, PlanetRadius)
#endif
  for (long jgl = 0; jgl < nlat; ++jgl)
    {
#ifdef _OPENMP
      Varray<double> pnm(dimsp), hnm(dimsp), ztemp1(waves << 1), ztemp2(waves << 1);
#endif
      const double gmusq = 1.0 - gmu[jgl] * gmu[jgl];
      coslat[jgl] = std::sqrt(gmusq);

      phcs(pnm.data(), hnm.data(), waves, gmu[jgl], ztemp1.data(), ztemp2.data());

      const double zgwt = gwt[jgl];
      const double zrafgmusqr = 1. / (PlanetRadius * gmusq);
      const double zradsqrtgmusqr = 1. / (-PlanetRadius * std::sqrt(gmusq));

      const int lpold = pold != nullptr;
      const int lpdev = pdev != nullptr;
      const int lpol2 = pol2 != nullptr;
      const int lpol3 = pol3 != nullptr;

      long jsp = jgl;
      for (long jm = 0; jm < waves; ++jm)
        for (long jn = 0; jn < waves - jm; ++jn)
          {
            poli[jsp] = pnm[jm * waves + jn] * 2.0;
            if (lpold) pold[jsp] = pnm[jm * waves + jn] * zgwt;
            if (lpdev) pdev[jsp] = hnm[jm * waves + jn] * 2.0 * zradsqrtgmusqr;
            if (lpol2) pol2[jsp] = hnm[jm * waves + jn] * zgwt * zrafgmusqr;
            if (lpol3) pol3[jsp] = pnm[jm * waves + jn] * zgwt * jm * zrafgmusqr;
            jsp += nlat;
          }
    }
}

void
after_legini(long ntr, long nlat, double *restrict poli, double *restrict pold, double *restrict coslat)
{
  long waves = ntr + 1; // used in omp loop
  long dimpnm = (ntr + 1) * (ntr + 4) / 2; // used in omp loop

#ifdef _OPENMP
  Varray2D<double> pnm2(Threading::ompNumThreads);
  Varray2D<double> work2(Threading::ompNumThreads);
  for (long i = 0; i < Threading::ompNumThreads; ++i) pnm2[i].resize(dimpnm);
  for (long i = 0; i < Threading::ompNumThreads; ++i) work2[i].resize(3 * waves);
#else
  Varray<double> pnm(dimpnm), work(3 * waves);
#endif

  Varray<double> gmu(nlat), gwt(nlat);
  gaussianLatitudes(gmu.data(), gwt.data(), nlat);
  for (long jgl = 0; jgl < nlat; ++jgl) gwt[jgl] *= 0.5;

  for (long jgl = 0; jgl < nlat; ++jgl) coslat[jgl] = std::sqrt(1.0 - gmu[jgl] * gmu[jgl]);

#ifdef _OPENMP
#pragma omp parallel for default(none) shared(nlat, waves, ntr, pnm2, work2, gwt, gmu, poli, pold)
#endif
  for (long jgl = 0; jgl < nlat / 2; jgl++)
    {
#ifdef _OPENMP
      const auto ompthID = cdo_omp_get_thread_num();
      auto pnm = pnm2[ompthID].data();
      auto work = work2[ompthID].data();
#endif
      const double zgwt = gwt[jgl];

      jspleg1(&pnm[0], gmu[jgl], ntr, &work[0]);

      long latn = jgl;
      long lats;
      long isp = 0;
      double is;
      for (long jm = 0; jm < waves; ++jm)
        {
#ifdef HAVE_OPENMP4
#pragma omp simd
#endif
          for (long jn = 0; jn < waves - jm; ++jn)
            {
              is = (jn + 1) % 2 * 2 - 1;
              lats = latn - jgl + nlat - jgl - 1;
              poli[latn] = pnm[isp];
              pold[latn] = pnm[isp] * zgwt;
              poli[lats] = pnm[isp] * is;
              pold[lats] = pnm[isp] * zgwt * is;
              latn += nlat;
              isp++;
            }
          isp++;
        }
    }
}

static inline void
sp2fc_kernel(long nlat, const double *pol, const double *sal, double *restrict far, double *restrict fai)
{
  const auto sar = sal[0];
  const auto sai = sal[1];
  for (long lat = 0; lat < nlat; lat++)
    {
      far[lat] += pol[lat] * sar;
      fai[lat] += pol[lat] * sai;
    }
}

void
sp2fc(const double *sa, double *fa, const double *poli, long nlev, long nlat, long nfc, long nt)
{
  long ntp1 = nt + 1;
  long nsp2 = (nt + 1) * (nt + 2);

  std::vector<long> cumindex(ntp1);
  cumindex[0] = 0;
  for (long jmm = 1; jmm < ntp1; jmm++) cumindex[jmm] = cumindex[jmm - 1] + (ntp1 - jmm + 1);

#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
  for (long lev = 0; lev < nlev; lev++)
    {
      auto pol = poli;
      auto sal = sa + lev * nsp2;
      auto fal = fa + lev * nfc * nlat;
      memset(fal, 0, nfc * nlat * sizeof(double));

      for (long jmm = 0; jmm < ntp1; jmm++)
        {
          auto polt = pol + cumindex[jmm] * nlat;
          auto falt = fal + jmm * 2 * nlat;
          auto salt = sal + cumindex[jmm] * 2;
          for (long jfc = 0; jfc < ntp1 - jmm; jfc++)
            {
              sp2fc_kernel(nlat, polt + jfc * nlat, salt + jfc * 2, falt, falt + nlat);
            }
        }
    }
}

static inline void
fc2sp_kernel(long nlat, const double *pol, const double *far, const double *fai, double *sal)
{
  double sar = 0.0;
  double sai = 0.0;
#ifdef HAVE_OPENMP4
#pragma omp simd reduction(+ : sar) reduction(+ : sai)
#endif
  for (long lat = 0; lat < nlat; lat++)
    {
      sar += pol[lat] * far[lat];
      sai += pol[lat] * fai[lat];
    }
  sal[0] = sar;
  sal[1] = sai;
}

void
fc2sp(const double *fa, double *sa, const double *poli, long nlev, long nlat, long nfc, long nt)
{
  long ntp1 = nt + 1;
  long nsp2 = (nt + 1) * (nt + 2);

  std::vector<long> cumindex(ntp1);
  cumindex[0] = 0;
  for (long jmm = 1; jmm < ntp1; jmm++) cumindex[jmm] = cumindex[jmm - 1] + (ntp1 - jmm + 1);

#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
  for (long lev = 0; lev < nlev; lev++)
    {
      auto pol = poli;
      auto fal = fa + lev * nfc * nlat;
      auto sal = sa + lev * nsp2;

      for (long jmm = 0; jmm < ntp1; jmm++)
        {
          auto polt = pol + cumindex[jmm] * nlat;
          auto falt = fal + jmm * 2 * nlat;
          auto salt = sal + cumindex[jmm] * 2;
          for (long jfc = 0; jfc < ntp1 - jmm; jfc++)
            {
              fc2sp_kernel(nlat, polt + jfc * nlat, falt, falt + nlat, salt + jfc * 2);
            }
        }
    }
}

/* ======================================== */
/* Convert Spectral Array to new truncation */
/* ======================================== */

void
sp2sp(double *arrayIn, long truncIn, double *arrayOut, long truncOut)
{
  if (truncOut <= truncIn)
    {
      for (long n = 0; n <= truncOut; n++)
        {
          for (long m = n; m <= truncOut; m++)
            {
              *arrayOut++ = *arrayIn++;
              *arrayOut++ = *arrayIn++;
            }
          arrayIn += 2 * (truncIn - truncOut);
        }
    }
  else
    {
      for (long n = 0; n <= truncIn; n++)
        {
          for (long m = n; m <= truncIn; m++)
            {
              *arrayOut++ = *arrayIn++;
              *arrayOut++ = *arrayIn++;
            }
          for (long m = truncIn + 1; m <= truncOut; ++m)
            {
              *arrayOut++ = 0.0;
              *arrayOut++ = 0.0;
            }
        }
      for (long n = truncIn + 1; n <= truncOut; ++n)
        for (long m = n; m <= truncOut; ++m)
          {
            *arrayOut++ = 0.0;
            *arrayOut++ = 0.0;
          }
    }
}

/* ======================================== */
/* Cut spectral wave numbers                */
/* ======================================== */

void
spcut(double *arrayIn, double *arrayOut, long trunc, const int *waves)
{
  for (long n = 0; n <= trunc; n++)
    {
      for (long m = n; m <= trunc; m++)
        {
          if (waves[m])
            {
              *arrayOut++ = *arrayIn++;
              *arrayOut++ = *arrayIn++;
            }
          else
            {
              *arrayOut++ = 0.0;
              *arrayOut++ = 0.0;
              arrayIn++;
              arrayIn++;
            }
        }
    }
}
