#include <math.h>
#include <stdlib.h>
#include <stdio.h>
#include <stdint.h>
#include <string.h>

#include "mg2d.h"
#include "mg2d_boundary.h"
#include "mg2d_constants.h"

#define MAXITER 64
#define TOL 1e-9

static const double pde_coeffs[MG2D_DIFF_COEFF_NB] = {
    [MG2D_DIFF_COEFF_00] = 1.0,
    [MG2D_DIFF_COEFF_10] = 0.9,
    [MG2D_DIFF_COEFF_01] = 1.1,
    [MG2D_DIFF_COEFF_20] = 1.2,
    [MG2D_DIFF_COEFF_02] = 0.8,
    [MG2D_DIFF_COEFF_11] = 0.7,
};

#if 1
static double sol_00(double x, double y)
{
    return sin(M_PI * x) * sin(M_PI * y);
}
static double sol_10(double x, double y)
{
    return M_PI * cos(M_PI * x) * sin(M_PI * y);
}
static double sol_01(double x, double y)
{
    return M_PI * sin(M_PI * x) * cos(M_PI * y);
}
static double sol_20(double x, double y)
{
    return -M_PI * M_PI * sol_00(x, y);
}
static double sol_02(double x, double y)
{
    return -M_PI * M_PI * sol_00(x, y);
}
static double sol_11(double x, double y)
{
    return M_PI * M_PI * cos(M_PI * x) * cos(M_PI * y);
}
#define BC_TYPE MG2D_BC_TYPE_FIXVAL
#else
static double sol_00(double x, double y)
{
    return cos(M_PI * x) * cos(M_PI * y);
}
static double sol_10(double x, double y)
{
    return -M_PI * sin(M_PI * x) * cos(M_PI * y);
}
static double sol_01(double x, double y)
{
    return -M_PI * cos(M_PI * x) * sin(M_PI * y);
}
static double sol_20(double x, double y)
{
    return -M_PI * M_PI * sol_00(x, y);
}
static double sol_02(double x, double y)
{
    return -M_PI * M_PI * sol_00(x, y);
}
static double sol_11(double x, double y)
{
    return M_PI * M_PI * sin(M_PI * x) * sin(M_PI * y);
}
#define BC_TYPE MG2D_BC_TYPE_FIXDIFF
#endif

static double (*sol[MG2D_DIFF_COEFF_NB])(double x, double y) = {
    [MG2D_DIFF_COEFF_00] = sol_00,
    [MG2D_DIFF_COEFF_10] = sol_10,
    [MG2D_DIFF_COEFF_01] = sol_01,
    [MG2D_DIFF_COEFF_20] = sol_20,
    [MG2D_DIFF_COEFF_02] = sol_02,
    [MG2D_DIFF_COEFF_11] = sol_11,
};

int main(int argc, char **argv)
{
    MG2DContext *ctx;
    long int gridsize;
    int ret = 0;

    if (argc < 2) {
        fprintf(stderr, "Usage: %s <N>\n", argv[0]);
        return 1;
    }
    gridsize = strtol(argv[1], NULL, 0);
    if (gridsize <= 0) {
        fprintf(stderr, "Invalid parameters: %ld\n", gridsize);
        return 1;
    }

    ctx = mg2d_solver_alloc(gridsize);
    if (!ctx) {
        fprintf(stderr, "Error allocating the solver context\n");
        return 1;
    }

    ctx->step[0] = 1.0 / (gridsize - 1);
    ctx->step[1] = 1.0 / (gridsize - 1);

    ctx->fd_stencil = 2;

    ctx->maxiter = MAXITER;
    ctx->nb_relax_pre  = 2;
    ctx->nb_cycles     = 1;
    ctx->nb_relax_post = 2;
    ctx->tol = TOL;
    ctx->nb_threads = 1;
    ctx->log_level = MG2D_LOG_VERBOSE;

    {
        MG2DBoundary *bnd = ctx->boundaries[MG2D_BOUNDARY_0L];

        bnd->type = BC_TYPE;

        memset(bnd->val, 0, gridsize * sizeof(*bnd->val));

        if (bnd->type == MG2D_BC_TYPE_FIXVAL) {
            for (int j = 1; j < ctx->fd_stencil; j++) {
                double *dst = bnd->val + j * bnd->val_stride;

                for (ptrdiff_t k = -j; k < (ptrdiff_t)ctx->domain_size + j; k++)
                    dst[k] = sol[MG2D_DIFF_COEFF_00](-j * ctx->step[0], k * ctx->step[1]);
            }
        }
    }
    {
        MG2DBoundary *bnd = ctx->boundaries[MG2D_BOUNDARY_0U];

        bnd->type = BC_TYPE;

        memset(bnd->val, 0, gridsize * sizeof(*bnd->val));

        if (bnd->type == MG2D_BC_TYPE_FIXVAL) {
            for (int j = 1; j < ctx->fd_stencil; j++) {
                double *dst = bnd->val + j * bnd->val_stride;

                for (ptrdiff_t k = -j; k < (ptrdiff_t)ctx->domain_size + j; k++)
                    dst[k] = sol[MG2D_DIFF_COEFF_00]((gridsize - 1 + j) * ctx->step[0], k * ctx->step[1]);
            }
        }
    }
    {
        MG2DBoundary *bnd = ctx->boundaries[MG2D_BOUNDARY_1L];

        bnd->type = BC_TYPE;

        memset(bnd->val, 0, gridsize * sizeof(*bnd->val));

        if (bnd->type == MG2D_BC_TYPE_FIXVAL) {
            for (int j = 1; j < ctx->fd_stencil; j++) {
                double *dst = bnd->val + j * bnd->val_stride;

                for (ptrdiff_t k = -j; k < (ptrdiff_t)ctx->domain_size + j; k++)
                    dst[k] = sol[MG2D_DIFF_COEFF_00](k * ctx->step[0], -j * ctx->step[1]);
            }
        }
    }
    {
        MG2DBoundary *bnd = ctx->boundaries[MG2D_BOUNDARY_1U];

        bnd->type = BC_TYPE;

        memset(bnd->val, 0, gridsize * sizeof(*bnd->val));

        if (bnd->type == MG2D_BC_TYPE_FIXVAL) {
            for (int j = 1; j < ctx->fd_stencil; j++) {
                double *dst = bnd->val + j * bnd->val_stride;

                for (ptrdiff_t k = -j; k < (ptrdiff_t)ctx->domain_size + j; k++)
                    dst[k] = sol[MG2D_DIFF_COEFF_00](k * ctx->step[0], (gridsize - 1 + j) * ctx->step[1]);
            }
        }
    }


    for (size_t y = 0; y < ctx->domain_size; y++) {
        const double y_coord = y * ctx->step[1];

        memset(ctx->u + y * ctx->u_stride, 0, sizeof(*ctx->u) * ctx->domain_size);

        for (size_t x = 0; x < ctx->domain_size; x++) {
            const double x_coord = x * ctx->step[0];
            double rhs = 0.0;

            for (int i = 0; i < MG2D_DIFF_COEFF_NB; i++) {
                ctx->diff_coeffs[i][ctx->diff_coeffs_stride * y + x] = pde_coeffs[i];
                rhs += pde_coeffs[i] * sol[i](x_coord, y_coord);
            }

            ctx->rhs[y * ctx->rhs_stride + x] = rhs;
        }
    }

    ret = mg2d_solve(ctx);
    if (ret < 0) {
        fprintf(stderr, "Error solving the equation\n");
        ret = 1;
        goto fail;
    }

    mg2d_print_stats(ctx, NULL);

    {
        double max_err = 0.0;

        for (size_t y = 0; y < ctx->domain_size; y++) {
            const double y_coord = y * ctx->step[1];

            for (size_t x = 0; x < ctx->domain_size; x++) {
                const double x_coord = x * ctx->step[0];
                double err = fabs(ctx->u[y * ctx->u_stride + x] - sol[MG2D_DIFF_COEFF_00](x_coord, y_coord));
                if (err > max_err)
                    max_err = err;
            }
        }
        fprintf(stderr, "max(|solution - exact|): %g\n", max_err);
        fprintf(stdout, "%ld %g\n", gridsize, max_err);
    }

fail:
    mg2d_solver_free(&ctx);
    return ret;
}