import argparse
import os
import re

module_start = """
// NOTE: THIS FILE IS AUTOMATICALLY GENERATED
// DO NOT EDIT BY HAND

// Re-run esum_cxx.py to update this file

// Fortran 2003 implementation of the msum routine
// provided by Raymond Hettinger:
// https://code.activestate.com/recipes/393090/
// This routine calculates the sum of N numbers
// exactly to within double precision arithmetic.
// Ported to C++.

// For performance reasons we implement a specialized
// version of esum for each possible value of N >= 3.

// Also for performance reasons, we explicitly unroll
// the outer loop of the msum method into groups of 3
// (and a group of 4 at the end, for even N). This seems
// to be significantly faster, but should still be exact
// to within the arithmetic because each one of the
// individual msums is (although this does not necessarily
// mean that the result is the same).

// This routine is called "esum" for generality
// because in principle we could add implementations
// other than msum that do exact arithmetic, without
// changing the interface as seen in the networks.

#ifndef ESUM_H
#define ESUM_H

#include <AMReX_REAL.H>
#include <AMReX_Array.H>
#include <ArrayUtilities.H>

using namespace amrex;

"""



module_end = """
#endif
"""


esum_base_template_start = """
template<int n, class T>
AMREX_GPU_HOST_DEVICE AMREX_INLINE
Real esum(T const& array)
{
    // return value
    Real sum = 0.0_rt;

    switch (n) {
    case 1:
        sum = array(1);
        break;
    case 2:
        sum = array(1) + array(2);
        break;
"""


esum_base_template_end = """
    default:
        sum = 0.0_rt;
        break;
    }
    return sum;
}
"""



esum_template_start = """
template<class T>
AMREX_GPU_HOST_DEVICE AMREX_INLINE
Real esum@NUM@(T const& array)
{
    // return value
    Real esum;
"""



esum_template_end = """
    return esum;
}
"""



sum_template = """
    esum = ArrayUtil::Math::sum(array, 1, @NUM@);
"""



kahan_template = """
    esum = array(1);
    Real x = 0._rt;
    for (int i = 2; i <= @NUM@; ++i) {
       Real y = array(i) - x;
       Real z = esum + y;
       x = (z - esum) - y;
       esum = z;
    }
"""



msum_template_start = """
    // Indices for tracking the partials array.
    // j keeps track of how many entries in partials are actually used.
    // The algorithm we model this off of, written in Python, simply
    // deletes array entries at the end of every outer loop iteration.
    // The Fortran equivalent to this might be to just zero them out,
    // but this results in a huge performance hit given how often
    // this routine is called during in a burn. So we opt instead to
    // just track how many of the values are meaningful, which j does
    // automatically, and ignore any data in the remaining slots.
    int i, j, k, km;

    // Note that for performance reasons we are not
    // initializing any unused values in this array.
    Array1D<Real, 0, @NUMPARTIALS@> partials;

    // Some temporary variables for holding intermediate data.
    Real x, y, z;

    // These temporary variables need to be explicitly
    // constructed for the algorithm to make sense.
    // If the compiler optimizes away the statement
    // lo = y - (hi - x), the approach fails. This could
    // be avoided with the volatile keyword, but at the
    // expense of forcing additional memory usage
    // which would slow down the calculation. Instead
    // we will rely on the compiler not to optimize
    // the statement away. This should be true for gcc
    // by default but is not necessarily true for all
    // compilers. In particular, Intel does not do this
    // by default, so you must use the -assume-protect-parens
    // flag for ifort.
    Real hi, lo;

    // The first partial is just the first term.
    esum = array(1);
"""

msum_template = """

    j = 0;
    partials(0) = esum;

    for (i = 2; i <= @NUM@; ++i) {

       km = j;
       j = 0;

       x = array(i+@START@);

       for (k = 0; k <= km; ++k) {
          y = partials(k);

          if (std::abs(x) < std::abs(y)) {
             // Swap x, y
             z = y;
             y = x;
             x = z;
          }

          hi = x + y;
          lo = y - (hi - x);

          if (lo != 0.0_rt) {
             partials(j) = lo;
             j++;
          }

          x = hi;

       }

       partials(j) = x;

    }

    esum = ArrayUtil::Math::sum(partials, 0, j);

"""



if __name__ == "__main__":

    sum_method = 0
    unroll = True

    parser = argparse.ArgumentParser()
    parser.add_argument('-s', help='summation method: -1 == sum(); 0 == msum; 1 == Kahan')
    parser.add_argument('--unroll', help='For msum, should we explicitly unroll the loop?')

    args = parser.parse_args()

    if args.s != None:
        sum_method = int(args.s)

    if args.unroll != None:
        if args.unroll == "True":
            unroll = True
        elif args.unroll == "False":
            unroll = False
        else:
            raise ValueError("--unroll can only be True or False.")

    with open("esum.H", "w") as ef:

        ef.write(module_start)

        first_esum = 3;
        last_esum = 30;

        for num in range(first_esum, last_esum + 1):

            ef.write(esum_template_start.replace("@NUM@", str(num)))

            if sum_method == -1:

                # ArrayUtil::Math::sum is just a sequential loop

                ef.write(sum_template.replace("@NUM@", str(num)))

            elif sum_method == 0:

                # msum

                if unroll:

                    ef.write(msum_template_start.replace("@NUM@", str(num)).replace("@NUMPARTIALS@", str(4)))

                    i = 1
                    while (i < num):
                        if (i == num - 3):
                            if (i > 0):
                                offset = i-1
                            else:
                                offset = 0
                            ef.write(msum_template.replace("@START@", str(offset)).replace("@NUM@", str(4)))
                            break
                        else:
                            if (i > 0):
                                offset = i-1
                            else:
                                offset = 0
                            ef.write(msum_template.replace("@START@", str(offset)).replace("@NUM@", str(3)))
                            i += 2

                else:

                    ef.write(msum_template_start.replace("@NUM@", str(num)).replace("@NUMPARTIALS@", str(num-1)))

                    ef.write(msum_template.replace("@START@", str(0)).replace("@NUM@", str(num)))

            elif sum_method == 1:

                # Kahan

                ef.write(kahan_template.replace("@NUM@", str(num)))

            else:

                raise ValueError("Unknown summation method.")

            ef.write(esum_template_end.replace("@NUM@", str(num)))
            ef.write("\n")

        # Now write out the base template that can call all the others

        ef.write(esum_base_template_start)

        for num in range(first_esum, last_esum + 1):
            ef.write("    case {}:\n".format(num))
            ef.write("        sum = esum{}(array);\n".format(num))
            ef.write("        break;\n")

        ef.write(esum_base_template_end)
        ef.write("\n")

        ef.write(module_end)
