#pragma once

#include <functional>

#include "Matrix.h"

/** Stencil */
struct Stencil {
    const double C; /*< Center */
    const double B; /*< Bottom */
    const double L; /*< Left */
    const double U; /*< Up */
    const double R; /*< Right */
};

class Solver {
public:
    Solver(const size_t nx, const size_t ny,
        const double xmin, const double xmax,
        const double ymin, const double ymax,
        const double h, const double k,
        std::function<double(double, double)> fun,
        std::function<double(double)> gN,
        std::function<double(double)> gE,
        std::function<double(double)> gS,
        std::function<double(double)> gW
    ) :
        _iter{0},
        _nx{nx}, _ny{ny},
        _xmin{xmin}, _xmax{xmax},
        _ymin{ymin}, _ymax{ymax},
        _stencil{ 4.0 / (h * h) + k * k,    /* Center */
                 -1.0 / (h * h),            /* Bottom */
                 -1.0 / (h * h),            /* Left */
                 -1.0 / (h * h),            /* Up */
                 -1.0 / (h * h)},           /* Right */
        _rhs(nx, ny, 0.),
        _sol(nx, ny, 0.),
        _tmp(nx, ny, 0.)
    {
        // Initialize Right Hand Size _rhs(x, y) = f(X(x), Y(y))
        for (size_t x = 0; x < nx; ++x) {
            for (size_t y = 0; y < ny; ++y) {
                _rhs(x, y) = fun(X(x), Y(y));
            }
        }

        // Set Dirichlet boundary conditions
        // North
        for (size_t x = 0; x < nx; ++x) { _sol.get(x, -1) = _tmp.get(x, -1) = gN(X(x)); }
        // East
        for (size_t y = 0; y < ny; ++y) { _sol.get(-1, y) = _tmp.get(-1, y) = gE(Y(y)); }
        // South
        for (size_t x = 0; x < nx; ++x) { _sol(x, 0) = _tmp(x, 0) = gS(X(x)); }
        // West
        for (size_t y = 0; y < ny; ++y) { _sol(0, y) = _tmp(0, y) = gW(Y(y)); }
    };

    /* Maps (i, j) indices to (x, y) position in the [xmin, xmax] x [ymin, ymax] domain */
    double X(size_t i) const {
        double ratio = static_cast<double>(i) / static_cast<double>(_nx - 1);
        return _xmin + ratio * (_xmax - _xmin);
    }
    double Y(size_t j) const {
        double ratio = static_cast<double>(j) / static_cast<double>(_ny - 1);
        return _ymin + ratio * (_ymax - _ymin);
    }

    enum class Dir {
        North, East, South, West,
        N = North, E = East, S = South, W = West
    };

    MatrixView<double> read_boundary(enum Dir dir) {
        switch (dir) {
            case Dir::North:
                return Col<double>(_sol, -2);
            case Dir::East:
                return Row<double>(_sol, -2);
            case Dir::South:
                return Col<double>(_sol, 1);
            default:    // Dir::West
                return Row<double>(_sol, 1);
        }
    };

    MatrixView<double> write_boundary(enum Dir dir) {
        switch (dir) {
            case Dir::North:
                return Col<double>(_sol, -1);
            case Dir::East:
                return Row<double>(_sol, -1);
            case Dir::South:
                return Col<double>(_sol, 0);
            default:    // Dir::West
                return Row<double>(_sol, 0);
        }
    };

    /** Solution getter */
    Matrix<double>& solution() { return _sol; }
    /** Right Hand Side getter (grid evaluated f(x, y)) */
    Matrix<double>& rhs() { return _rhs; }

    /** Performs a single Jacobi iteration */
    void iterate() {
        double s = 1.0 / _stencil.C;

        for (size_t y = 1; y < _ny - 1; ++y) {
            for (size_t x = 1; x < _nx - 1; ++x) {
                _tmp(x, y) = s * (_rhs(x, y) - (_stencil.B * _sol(x, y - 1) +
                                                _stencil.L * _sol(x - 1, y) +
                                                _stencil.U * _sol(x, y + 1) +
                                                _stencil.R * _sol(x + 1, y)));
            }
        }

        std::swap(_tmp, _sol);
        _iter++;
    };

    /** Computes the L2 and Inf norm of the residuals */
    template <enum Norm N = Norm::Frob>
    double resid_norm() {
        // Enforce valid Norm types
        static_assert(N == Norm::Frob || N == Norm::Max);

        // Norm accumulator
        double accum = 0.0;

        for (size_t y = 1; y < _ny - 1; ++y) {
            for (size_t x = 1; x < _nx - 1; ++x) {
                double resid = _rhs(x, y) - (_stencil.C * _sol(x, y) +
                                             _stencil.B * _sol(x, y - 1) +
                                             _stencil.L * _sol(x - 1, y) +
                                             _stencil.U * _sol(x, y + 1) +
                                             _stencil.R * _sol(x + 1, y));
                if constexpr (N == Norm::Frob) {
                    accum += resid * resid;
                } else {    // Maximum Norm
                    accum = std::max(std::abs(resid), accum);
                }
            }
        }

        if constexpr (N == Norm::Frob) {
            return std::sqrt(accum);
        } else {    // Maximum Norm
            return accum;
        }
    }

private:
    size_t _iter;           /*< Iteration count */
    const size_t _nx;       /*< Number of X-axis grid points */
    const size_t _ny;       /*< Number of Y-axis grid points */
    const double _xmin;     /*< Domain X-min (west border position) */
    const double _xmax;     /*< Domain X-max (east border position) */
    const double _ymin;     /*< Domain Y-min (south border position) */
    const double _ymax;     /*< Domain Y-max (north border position) */
    const Stencil _stencil; /*< Simple '+' shaped stencil */
    Matrix<double> _rhs;    /*< Grid evaluated RHS of the PDE, f(x, y) */
    Matrix<double> _sol;    /*< Solution after `_iter` iterations */
    Matrix<double> _tmp;    /*< Temp. data block, used in iterate() */
};