#ifndef CPU_KERNELS_HH
#define CPU_KERNELS_HH

#include <algorithm>
#include <cmath>
#include <cstring>
#include <stdexcept>
#include <string>
#include <type_traits>

// Ideally we'd use an alias declaration for a generic definition of
// *axpy. But Cython doesn't support alias declarations yet:
//
// https://github.com/cython/cython/issues/3272
//
// template <typename T>
// using axpy = void (*)(int N, T alpha, const T* X, int incX,
//                       T *Y, int incY);
//
// So, instead we'll do this the pre-C++11 way:

template <typename T>
struct axpy {
    typedef void (*ptr)(int N, T alpha, const T* X, int incX, T *Y, int incY);
};


// All elementwise functions, such as most activations, work in-place.


template <typename T, typename L>
struct argmax_result {
    T max;
    L max_idx;
};

template <typename T, typename L>
argmax_result<T, L> argmax(T const *arr, L len)
{
    static_assert(std::is_floating_point<T>::value,
        "Array should be floating point");
    static_assert(std::is_integral<L>::value, "Array length should be integral");

    argmax_result<T, L> r { arr[0],  0 };

    for (L i = 1; i < len; ++i) {
        if (arr[i] > r.max) {
            r.max = arr[i];
            r.max_idx = i;
        }
    }

    return r;
}

// The next two templates define argmax for a fixed number of elements.

template <typename T, typename L>
argmax_result<T, L> argmax(T a) {
    static_assert(std::is_floating_point<T>::value, "Argument should be floating point");
    argmax_result<T, L> acc { a, 0 };
    return acc;
}

template<typename T, typename L, typename... Args>
argmax_result<T, L> argmax(T a, Args... args) {
    static_assert(std::is_floating_point<T>::value, "Arguments should be floating point");

    auto acc = argmax<T, L>(args...);

    if (acc.max > a) {
        acc.max_idx += 1;
    } else {
        acc.max_idx = 0;
        acc.max = a;
    }

    return acc;
}


template <typename A, typename L>
void vec_add(A* X, const A* Y, A scale, L N)
{
    static_assert(std::is_floating_point<A>::value,
        "Array should be floating point");
    static_assert(std::is_integral<L>::value, "Array length should be integral");

    for (L i = 0; i < N; ++i)
        X[i] += scale * Y[i];
}

template <typename A, typename L>
void cpu_maxout(A* best__bo, L* which__bo, const A* cands__bop, L B, L O, L P)
{
    static_assert(std::is_floating_point<A>::value,
        "Array should be floating point");
    static_assert(std::is_integral<L>::value, "Array length should be integral");

    // For small inputs, we use an unrolled argmax.
    if (P == 2) {
        for (int i = 0; i < B * O; ++i) {
            A const *input = cands__bop + i * P;
            auto r = argmax<A, L>(input[0], input[1]);
            which__bo[i] = r.max_idx;
            best__bo[i] = r.max;
        }
    } else if (P == 3) {
        for (int i = 0; i < B * O; ++i) {
            A const *input = cands__bop + i * P;
            auto r = argmax<A, L>(input[0], input[1], input[2]);
            which__bo[i] = r.max_idx;
            best__bo[i] = r.max;
        }
    } else {
        for (int i = 0; i < B * O; ++i) {
            auto r = argmax<A, L>(cands__bop + i * P, P);
            which__bo[i] = r.max_idx;
            best__bo[i] = r.max;
        }
    }
}


template <typename A, typename L>
void cpu_backprop_maxout(A* dX__bop, const A* dX__bo, const L* which__bo,
    L B, L O, L P)
{
    static_assert(std::is_floating_point<A>::value,
        "Array should be floating point");
    static_assert(std::is_integral<L>::value, "Array length should be integral");

    for (L b = 0; b < B; ++b) {
        for (L o = 0; o < O; ++o) {
            if (*which__bo >= P) {
                throw std::out_of_range(std::string("index ") + std::to_string(*which__bo) + " is out of bounds for maxout with size " + std::to_string(P));
            }

            dX__bop[*which__bo] = *dX__bo;
            dX__bop += P;
            dX__bo += 1;
            which__bo += 1;
        }
    }
}

template <typename A, typename L>
void cpu_reduce_max(A* maxes__bo, L* which__bo, const A* X__to,
    const L* lengths__b, L B, L T, L O)
{
    static_assert(std::is_floating_point<A>::value,
        "Array should be floating point");
    static_assert(std::is_integral<L>::value, "Array length should be integral");

    for (const L* length = lengths__b; length < lengths__b + B; ++length) {
        if (*length <= 0)
            throw std::invalid_argument(std::string("all sequence lengths must be > 0, was: ") + std::to_string(*length));
        else if (*length > T) {
            throw std::out_of_range("lengths must sum up to the number of rows");
        }

        T -= *length;

        std::memcpy(maxes__bo, X__to, O * sizeof(*maxes__bo));
        X__to += O;
        for (L i = 1; i < *length; ++i) {
            for (L j = 0; j < O; ++j) {
                if (X__to[j] > maxes__bo[j]) {
                    maxes__bo[j] = X__to[j];
                    which__bo[j] = i;
                }
            }
            X__to += O;
        }

        maxes__bo += O;
        which__bo += O;
    }
}

template <typename A, typename L>
void cpu_backprop_reduce_max(A* dX__to, const A* d_maxes__bo, const L* which__bo,
    const L* lengths__b, L B, L T, L O)
{
    static_assert(std::is_floating_point<A>::value,
        "Array should be floating point");
    static_assert(std::is_integral<L>::value, "Array length should be integral");

    for (const L* length = lengths__b; length < lengths__b + B; ++length) {
        for (L i = 0; i < O; ++i) {
            L item = which__bo[i];
            if (item >= *length) {
                throw std::out_of_range(std::string("index ") + std::to_string(item) + " is out of bounds for maxout with length " + std::to_string(*length));
            }

            dX__to[item * O + i] = d_maxes__bo[i];
        }

        dX__to += *length * O;
        d_maxes__bo += O;
        which__bo += O;
    }
}

template <typename A, typename L>
void cpu_reduce_mean(A* means__bo, const A* X__to, const L* lengths__b,
    L B, L T, L O)
{
    static_assert(std::is_floating_point<A>::value,
        "Array should be floating point");
    static_assert(std::is_integral<L>::value, "Array length should be integral");

    for (const L* length = lengths__b; length < lengths__b + B; ++length) {
        if (*length < 0) {
            throw std::invalid_argument(std::string("all sequence lengths must be >= 0, was: ") + std::to_string(*length));
        }
        else if (length == 0) {
            means__bo += O;
            continue;
        }
        else if (*length > T) {
            throw std::out_of_range("lengths must sum up to the number of rows");
        }

        T -= *length;

        A scale = 1. / *length;
        for (L i = 0; i < *length; ++i) {
            vec_add(means__bo, X__to, scale, O);
            X__to += O;
        }

        means__bo += O;
    }
}

template <typename A, typename L>
void cpu_backprop_reduce_mean(A* dX__to, const A* d_means__bo, const L* lengths__b,
    L B, L T, L O)
{
    static_assert(std::is_floating_point<A>::value,
        "Array should be floating point");
    static_assert(std::is_integral<L>::value, "Array length should be integral");

    for (const L* length = lengths__b; length < lengths__b + B; ++length) {
        A scale = 1. / *length;
        for (L i = 0; i < *length; ++i) {
            vec_add(dX__to, d_means__bo, scale, O);
            dX__to += O;
        }

        d_means__bo += O;
    }
}

template <typename A, typename L>
void cpu_mish(A* Y, L N, A threshold)
{
    static_assert(std::is_floating_point<A>::value,
        "Array should be floating point");
    static_assert(std::is_integral<L>::value, "Array length should be integral");

    for (L i = 0; i < N; ++i) {
        if (Y[i] < threshold) {
            Y[i] *= std::tanh(std::log(1.0 + std::exp(Y[i])));
        }
    }
}

template <typename A, typename L>
void cpu_backprop_mish(A* dX, const A* X, L N, A threshold)
{
    static_assert(std::is_floating_point<A>::value,
        "Array should be floating point");
    static_assert(std::is_integral<L>::value, "Array length should be integral");

    for (L i = 0; i < N; ++i) {
        A x = X[i];

        if (x < threshold) {
            A exp_x = std::exp(x);
            A exp_2x = std::exp(2 * x);
            A exp_3x = std::exp(3 * x);
            A omega = (4. * (x + 1)) + (4 * exp_2x) + exp_3x + exp_x * (4. * x + 6);
            A delta = 2. * exp_x + exp_2x + 2.;
            dX[i] = dX[i] * ((exp_x * omega) / (delta * delta));
        }
    }
}

template <typename A, typename L>
void cpu_reduce_sum(A* sums__bo, const A* X__to, const L* lengths__b,
    L B, L T, L O)
{
    static_assert(std::is_floating_point<A>::value,
        "Array should be floating point");
    static_assert(std::is_integral<L>::value, "Array length should be integral");

    for (const L* length = lengths__b; length < lengths__b + B; ++length) {
        if (*length < 0) {
            throw std::invalid_argument(std::string("all sequence lengths must be >= 0, was: ") + std::to_string(*length));
        }
        else if (length == 0) {
            sums__bo += O;
            continue;
        }
        else if (*length > T) {
            throw std::out_of_range("lengths must sum up to the number of rows");
        }

        T -= *length;

        for (L i = 0; i < *length; ++i) {
            vec_add(sums__bo, X__to, static_cast<A>(1.0), O);
            X__to += O;
        }

        sums__bo += O;
    }
}

template <typename A, typename L>
void cpu_backprop_reduce_sum(A* dX__to, const A* d_sums__bo, const L* lengths__b,
    L B, L T, L O)
{
    static_assert(std::is_floating_point<A>::value,
        "Array should be floating point");
    static_assert(std::is_integral<L>::value, "Array length should be integral");

    for (const L* length = lengths__b; length < lengths__b + B; ++length) {
        for (L i = 0; i < *length; ++i) {
            vec_add(dX__to, d_sums__bo, static_cast<A>(1.0), O);
            dX__to += O;
        }

        d_sums__bo += O;
    }
}

template <typename A, typename L>
void cpu_relu(A* X, L N)
{
    static_assert(std::is_floating_point<A>::value,
        "Array should be floating point");
    static_assert(std::is_integral<L>::value, "Array length should be integral");

    for (L i = 0; i < N; ++i) {
        if (X[i] <= 0.0) {
            X[i] = 0.0;
        }
    }
}

template <typename A, typename L>
void seq2col(A* output, const A* X, const L* lengths, L nW, L B, L I, L nL)
{
    // Let's say nW is 1 (it usually is). Then we want to take:

    // 1a 1b 1c
    // 2a 2b 2c
    // 3a 3b 3c

    // And make

    // __ __ __ 1a 1b 1c 2a 2b 2c
    // 1a 1b 1c 2a 2b 2c 3a 3b 3c
    // 2a 2b 2c 3a 3b 3c __ __ __

    // Where __ is padding.

    // Now let's say nW is 2. Then we want to take:

    // 1a 1b 1c
    // 2a 2b 2c
    // 3a 3b 3c

    // And make

    // __ __ __ __ __ __ 1a 1b 1c 2a 2b 2c 3a 3b 3c
    // __ __ __ 1a 1b 1c 2a 2b 2c 3a 3b 3c __ __ __
    // 1a 1b 1c 2a 2b 2c 3a 3b 3c __ __ __ __ __ __

    // * x_start=-6, x_end=9 : (0-2) * 3, (0+2+1) * 3
    // * x_start=-3, x_end=13 : (1-2) * 3, (1+2+1) * 3
    // * x_start=0, x_end=16 : (2-2) * 3, (2+2+1) * 3

    // If lengths > 1, then the sequence lengths dictate
    // the boundaries/padding rather than the begin/end
    // of X.
    static_assert(std::is_floating_point<A>::value,
        "Array should be floating point");
    static_assert(std::is_integral<L>::value, "Array length should be integral");

    L nF = nW * 2 + 1;

    L seq_start = 0;
    for (L i = 0; i < nL; ++i) {
        // Calculate the bounds of the next sequence.
        L seq_end = seq_start + lengths[i];

        for (L j = seq_start; j < seq_end; ++j) {
            // Find the unconstrained window around b, which
            // may be out of the sequence bounds.
            L window_start = j - nW;
            L window_end = j + nW + 1;

            // Find the sequence-constrained window around b.
            L x_start = std::max(seq_start, window_start);
            L x_end = std::min(seq_end, window_end);
            L n_elems = x_end - x_start;

            L out_offset = x_start - window_start;

            std::memcpy(output + (j * nF * I) + (out_offset * I),
                X + (x_start * I),
                n_elems * I * sizeof(*output));
        }

        seq_start += lengths[i];
    }
}

template <typename A, typename L>
void backprop_seq2col(A* d_seqs, const A* d_cols, const L* lengths, L B, L I, L nW, L nL)
{
    // here's what we're doing, if we had 2d indexing.
    // for i in range(b):
    //     d_seq[i] += d_cols[i-2, 4]
    //     d_seq[i] += d_cols[i-1, 3]
    //     d_seq[i] += d_cols[i, 2]
    //     d_seq[i] += d_cols[i+1, 1]
    //     d_seq[i] += d_cols[i+2, 0]
    static_assert(std::is_floating_point<A>::value,
        "Array should be floating point");
    static_assert(std::is_integral<L>::value, "Array length should be integral");

    L nF = nW * 2 + 1;

    L seq_start = 0;
    for (L i = 0; i < nL; ++i) {
        // Calculate the bounds of the next sequence.
        L seq_end = seq_start + lengths[i];

        for (L j = seq_start; j < seq_end; ++j) {
            // Find the unconstrained window around b, which
            // may be out of the sequence bounds.
            L window_begin = j - nW;
            L window_end = j + nW + 1;

            // Find the sequence-constrained window around b.
            L d_seqs_begin = std::max(seq_start, window_begin);
            L d_seqs_end = std::min(seq_end, window_end);
            L n_elems = d_seqs_end - d_seqs_begin;

            // If the left window is cut short, we want to
            // start by the same amount in the output.
            L out_offset = d_seqs_begin - window_begin;

            vec_add(d_seqs + d_seqs_begin * I,
                d_cols + (j * nF * I) + (out_offset * I),
                static_cast<A>(1.), n_elems * I);
        }

        seq_start += lengths[i];
    }
}

template <typename F, typename I, typename L>
void cpu_gather_add(typename axpy<F>::ptr axpy, F* out_bo, const F* table_to, const I* indices_bk, L T, L O, L B, L K) {
     for (L b = 0; b < B; ++b) {
        for (L k = 0; k < K; ++k) {
            I idx = indices_bk[b * K + k];
            if (idx > T) {
                throw std::out_of_range("Embedding index out-of-bounds");
            }
            axpy(O, 1.0, table_to + idx * O, 1, out_bo + b * O, 1);
        }
    }
}


#endif // CPU_KERNELS_HH