/*
 * Copyright 2019 Anton Khirnov <anton@khirnov.net>
 *
 * 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, either version 3 of the License, or
 * (at your option) any later version.
 *
 * 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.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include "config.h"

#include <errno.h>
#include <math.h>
#include <stddef.h>
#include <stdlib.h>
#include <string.h>

#include <threadpool.h>

#include "common.h"
#include "cpu.h"
#include "mg2d_boundary.h"
#include "mg2d_constants.h"
#include "residual_calc.h"

typedef void ResidualLineCalc(size_t linesize, double *dst, double *dst_max,
                              ptrdiff_t u_stride, const double *u, const double *rhs,
                              const double *diff_coeffs, ptrdiff_t diff_coeffs_offset,
                              double res_mult);
typedef void ResidualLineAdd (size_t linesize, double *dst, double *dst_max,
                              ptrdiff_t u_stride, const double *u, const double *rhs,
                              const double *diff_coeffs, ptrdiff_t diff_coeffs_offset,
                              double res_mult, double u_mult);

typedef struct ResidualCalcTask {
    size_t                size[2];

    double               *dst;
    ptrdiff_t             dst_stride;

    const double         *u;
    ptrdiff_t             u_stride;

    const double         *rhs;
    ptrdiff_t             rhs_stride;

    const double         *diff_coeffs;
    ptrdiff_t             diff_coeffs_stride;
    ptrdiff_t             diff_coeffs_offset;

    double u_mult;
    double res_mult;

    int    reflect;
    size_t reflect_dist;
} ResidualCalcTask;

struct ResidualCalcInternal {
    double *residual_max;
    size_t  residual_max_size;

    ResidualLineCalc *residual_line_calc;
    ResidualLineAdd  *residual_line_add;
    size_t calc_blocksize;

    ResidualCalcTask task;
};

#if HAVE_NASM
ResidualLineCalc mg2di_residual_line_calc_s1_avx2;
ResidualLineCalc mg2di_residual_line_calc_s2_avx2;
ResidualLineAdd  mg2di_residual_line_add_s1_avx2;
ResidualLineAdd  mg2di_residual_line_add_s2_avx2;
#endif

static void
derivatives_calc_s1(double *dst, const double *u, ptrdiff_t stride)
{
    dst[MG2D_DIFF_COEFF_00] = u[0];
    dst[MG2D_DIFF_COEFF_10] = (u[1]      - u[-1]);
    dst[MG2D_DIFF_COEFF_01] = (u[stride] - u[-stride]);

    dst[MG2D_DIFF_COEFF_20] = (u[1]      - 2.0 * u[0] + u[-1]);
    dst[MG2D_DIFF_COEFF_02] = (u[stride] - 2.0 * u[0] + u[-stride]);

    dst[MG2D_DIFF_COEFF_11] = (u[1 + stride] - u[stride - 1] - u[-stride + 1] + u[-stride - 1]);
}

static void
derivatives_calc_s2(double *dst, const double *u, ptrdiff_t stride)
{
    const double val    = u[0];

    const double valxp1 = u[ 1];
    const double valxp2 = u[ 2];
    const double valxm1 = u[-1];
    const double valxm2 = u[-2];
    const double valyp1 = u[ 1 * stride];
    const double valyp2 = u[ 2 * stride];
    const double valym1 = u[-1 * stride];
    const double valym2 = u[-2 * stride];

    const double valxp1yp1 = u[ 1 + 1 * stride];
    const double valxp1yp2 = u[ 1 + 2 * stride];
    const double valxp1ym1 = u[ 1 - 1 * stride];
    const double valxp1ym2 = u[ 1 - 2 * stride];

    const double valxp2yp1 = u[ 2 + 1 * stride];
    const double valxp2yp2 = u[ 2 + 2 * stride];
    const double valxp2ym1 = u[ 2 - 1 * stride];
    const double valxp2ym2 = u[ 2 - 2 * stride];

    const double valxm1yp1 = u[-1 + 1 * stride];
    const double valxm1yp2 = u[-1 + 2 * stride];
    const double valxm1ym1 = u[-1 - 1 * stride];
    const double valxm1ym2 = u[-1 - 2 * stride];

    const double valxm2yp1 = u[-2 + 1 * stride];
    const double valxm2yp2 = u[-2 + 2 * stride];
    const double valxm2ym1 = u[-2 - 1 * stride];
    const double valxm2ym2 = u[-2 - 2 * stride];

    dst[MG2D_DIFF_COEFF_00] = val;
    dst[MG2D_DIFF_COEFF_10] = (-1.0 * valxp2 + 8.0 * valxp1      - 8.0 * valxm1 + 1.0 * valxm2);
    dst[MG2D_DIFF_COEFF_01] = (-1.0 * valyp2 + 8.0 * valyp1      - 8.0 * valym1 + 1.0 * valym2);

    dst[MG2D_DIFF_COEFF_20] = (-1.0 * valxp2 + 16.0 * valxp1 - 30.0 * val + 16.0 * valxm1 - 1.0 * valxm2);
    dst[MG2D_DIFF_COEFF_02] = (-1.0 * valyp2 + 16.0 * valyp1 - 30.0 * val + 16.0 * valym1 - 1.0 * valym2);

    dst[MG2D_DIFF_COEFF_11] = ( 1.0 * valxp2yp2 -  8.0 * valxp2yp1 +  8.0 * valxp2ym1 - 1.0 * valxp2ym2
                               -8.0 * valxp1yp2 + 64.0 * valxp1yp1 - 64.0 * valxp1ym1 + 8.0 * valxp1ym2
                               +8.0 * valxm1yp2 - 64.0 * valxm1yp1 + 64.0 * valxm1ym1 - 8.0 * valxm1ym2
                               -1.0 * valxm2yp2 +  8.0 * valxm2yp1 -  8.0 * valxm2ym1 + 1.0 * valxm2ym2);
}

static void residual_calc_line_s1_c(size_t linesize, double *dst, double *dst_max,
                                    ptrdiff_t u_stride, const double *u, const double *rhs,
                                    const double *diff_coeffs, ptrdiff_t diff_coeffs_offset,
                                    double res_mult)
{
    double res_max = 0.0, res_abs;
    for (size_t i = 0; i < linesize; i++) {
        double u_vals[MG2D_DIFF_COEFF_NB];
        double res;

        derivatives_calc_s1(u_vals, u + i, u_stride);

        res = -rhs[i];
        for (int j = 0; j < ARRAY_ELEMS(u_vals); j++)
            res += u_vals[j] * diff_coeffs[j * diff_coeffs_offset + i];
        dst[i] = res_mult * res;

        res_abs = fabs(res);
        res_max = MAX(res_max, res_abs);
    }

    *dst_max = MAX(*dst_max, res_max);
}

static void residual_add_line_s1_c(size_t linesize, double *dst, double *dst_max,
                                   ptrdiff_t u_stride, const double *u, const double *rhs,
                                   const double *diff_coeffs, ptrdiff_t diff_coeffs_offset,
                                   double res_mult, double u_mult)
{
    double res_max = 0.0, res_abs;
    for (size_t i = 0; i < linesize; i++) {
        double u_vals[MG2D_DIFF_COEFF_NB];
        double res;

        derivatives_calc_s1(u_vals, u + i, u_stride);

        res = -rhs[i];
        for (int j = 0; j < ARRAY_ELEMS(u_vals); j++)
            res += u_vals[j] * diff_coeffs[j * diff_coeffs_offset + i];
        dst[i] = u_mult * u[i] + res_mult * res;

        res_abs = fabs(res);
        res_max = MAX(res_max, res_abs);
    }

    *dst_max = MAX(*dst_max, res_max);
}

static void residual_calc_line_s2_c(size_t linesize, double *dst, double *dst_max,
                                    ptrdiff_t u_stride, const double *u, const double *rhs,
                                    const double *diff_coeffs, ptrdiff_t diff_coeffs_offset,
                                    double res_mult)
{
    double res_max = 0.0, res_abs;
    for (size_t i = 0; i < linesize; i++) {
        double u_vals[MG2D_DIFF_COEFF_NB];
        double res;

        derivatives_calc_s2(u_vals, u + i, u_stride);

        res = -rhs[i];
        for (int j = 0; j < ARRAY_ELEMS(u_vals); j++)
            res += u_vals[j] * diff_coeffs[j * diff_coeffs_offset + i];
        dst[i] = res_mult * res;

        res_abs = fabs(res);
        res_max = MAX(res_max, res_abs);
    }

    *dst_max = MAX(*dst_max, res_max);
}

static void residual_add_line_s2_c(size_t linesize, double *dst, double *dst_max,
                                    ptrdiff_t u_stride, const double *u, const double *rhs,
                                    const double *diff_coeffs, ptrdiff_t diff_coeffs_offset,
                                    double res_mult, double u_mult)
{
    double res_max = 0.0, res_abs;
    for (size_t i = 0; i < linesize; i++) {
        double u_vals[MG2D_DIFF_COEFF_NB];
        double res;

        derivatives_calc_s2(u_vals, u + i, u_stride);

        res = -rhs[i];
        for (int j = 0; j < ARRAY_ELEMS(u_vals); j++)
            res += u_vals[j] * diff_coeffs[j * diff_coeffs_offset + i];
        dst[i] = u_mult * u[i] + res_mult * res;

        res_abs = fabs(res);
        res_max = MAX(res_max, res_abs);
    }

    *dst_max = MAX(*dst_max, res_max);
}

static int residual_calc_task(void *arg, unsigned int job_idx, unsigned int thread_idx)
{
    ResidualCalcInternal *priv = arg;
    ResidualCalcTask     *task = &priv->task;

    const double *diff_coeffs = task->diff_coeffs + job_idx * task->diff_coeffs_stride;
    double *dst = task->dst + job_idx * task->dst_stride;

    if (task->u_mult == 0.0) {
        priv->residual_line_calc(task->size[0], dst,
                                 priv->residual_max + thread_idx * priv->calc_blocksize,
                                 task->u_stride, task->u + job_idx * task->u_stride,
                                 task->rhs + job_idx * task->rhs_stride,
                                 diff_coeffs, task->diff_coeffs_offset, task->res_mult);
    } else {
        priv->residual_line_add(task->size[0], dst,
                                priv->residual_max + thread_idx * priv->calc_blocksize,
                                task->u_stride, task->u + job_idx * task->u_stride,
                                task->rhs + job_idx * task->rhs_stride,
                                diff_coeffs, task->diff_coeffs_offset, task->res_mult, task->u_mult);
    }

    if (task->reflect & (1 << MG2D_BOUNDARY_0L)) {
        for (int i = 1; i <= task->reflect_dist; i++)
            dst[-i] = dst[i];
    }
    if (task->reflect & (1 << MG2D_BOUNDARY_0U)) {
        for (int i = 1; i <= task->reflect_dist; i++)
            dst[task->size[0] - 1 + i] = dst[task->size[0] - 1 - i];
    }
    if ((task->reflect & (1 << MG2D_BOUNDARY_1L)) &&
        job_idx > 0 && job_idx <= task->reflect_dist) {
        memcpy(task->dst - job_idx * task->dst_stride, dst, sizeof(*dst) * task->size[0]);
    }
    if ((task->reflect & (1 << MG2D_BOUNDARY_1U)) &&
        job_idx >= task->size[1] - 1 - task->reflect_dist && job_idx < task->size[1] - 1) {
        memcpy(task->dst + (2 * (task->size[1] - 1) - job_idx) * task->dst_stride, dst, sizeof(*dst) * task->size[0]);
    }

    return 0;
}

int mg2di_residual_calc(ResidualCalcContext *ctx, size_t size[2],
                        double      *residual_max,
                        double               *dst, ptrdiff_t dst_stride,
                        const double           *u, ptrdiff_t u_stride,
                        const double         *rhs, ptrdiff_t rhs_stride,
                        const double *diff_coeffs, ptrdiff_t diff_coeffs_stride,
                                                   ptrdiff_t diff_coeffs_offset,
                        double u_mult, double res_mult,
                        int reflect, size_t reflect_dist)
{
    ResidualCalcInternal *priv = ctx->priv;
    ResidualCalcTask     *task = &priv->task;
    double res_max = 0.0;

    memset(priv->residual_max, 0, sizeof(*priv->residual_max) * priv->residual_max_size);

    task->size[0]            = size[0];
    task->size[1]            = size[1];
    task->dst                = dst;
    task->dst_stride         = dst_stride;
    task->u                  = u;
    task->u_stride           = u_stride;
    task->rhs                = rhs;
    task->rhs_stride         = rhs_stride;
    task->diff_coeffs        = diff_coeffs;
    task->diff_coeffs_stride = diff_coeffs_stride;
    task->diff_coeffs_offset = diff_coeffs_offset;
    task->u_mult             = u_mult;
    task->res_mult           = res_mult;
    task->reflect            = reflect;
    task->reflect_dist       = reflect_dist;

    tp_execute(ctx->tp, size[1], residual_calc_task, priv);

    for (size_t i = 0; i < priv->residual_max_size; i++)
        res_max = MAX(res_max, priv->residual_max[i]);
    *residual_max = res_max;

    return 0;
}

int mg2di_residual_calc_init(ResidualCalcContext *ctx)
{
    ResidualCalcInternal *priv = ctx->priv;
    double *tmp;

    priv->calc_blocksize = 1;
    switch (ctx->fd_stencil) {
    case 1:
        priv->residual_line_calc = residual_calc_line_s1_c;
        priv->residual_line_add  = residual_add_line_s1_c;
#if HAVE_NASM
        if (ctx->cpuflags & MG2DI_CPU_FLAG_AVX2) {
            priv->residual_line_calc = mg2di_residual_line_calc_s1_avx2;
            priv->residual_line_add  = mg2di_residual_line_add_s1_avx2;
            priv->calc_blocksize     = 4;
        }
#endif
        break;
    case 2:
        priv->residual_line_calc = residual_calc_line_s2_c;
        priv->residual_line_add  = residual_add_line_s2_c;
#if HAVE_NASM
        if (ctx->cpuflags & MG2DI_CPU_FLAG_AVX2) {
            priv->residual_line_calc = mg2di_residual_line_calc_s2_avx2;
            priv->residual_line_add  = mg2di_residual_line_add_s2_avx2;
            priv->calc_blocksize     = 4;
        }
#endif
        break;
    }

    priv->residual_max_size = tp_get_nb_threads(ctx->tp) * priv->calc_blocksize;
    tmp = realloc(priv->residual_max,
                  sizeof(*priv->residual_max) * priv->residual_max_size);
    if (!tmp) {
        priv->residual_max_size = 0;
        return -ENOMEM;
    }
    priv->residual_max = tmp;

    return 0;
}

ResidualCalcContext *mg2di_residual_calc_alloc(void)
{
    ResidualCalcContext *ctx;

    ctx = calloc(1, sizeof(*ctx));
    if (!ctx)
        return NULL;

    ctx->priv = calloc(1, sizeof(*ctx->priv));
    if (!ctx->priv) {
        free(ctx);
        return NULL;
    }

    return ctx;
}

void mg2di_residual_calc_free(ResidualCalcContext **pctx)
{
    ResidualCalcContext *ctx = *pctx;

    if (!ctx)
        return;

    free(ctx->priv->residual_max);
    free(ctx->priv);

    free(ctx);
    *pctx = NULL;
}