from typing import Tuple
import sympy as sp
from pystencils.astnodes import LoopOverCoordinate
from pystencils.cache import memorycache
from pystencils import Field
from pystencils.fd import Diff
from pystencils.transformations import generic_visit
from .derivative import diff_args
from .derivation import FiniteDifferenceStencilDerivation


def fd_stencils_standard(indices, dx, fa):
    order = len(indices)
    assert all(i >= 0 for i in indices), "Can only discretize objects with (integer) subscripts"
    if order == 1:
        idx = indices[0]
        return (fa.neighbor(idx, 1) - fa.neighbor(idx, -1)) / (2 * dx)
    elif order == 2:
        if indices[0] == indices[1]:
            return (-2 * fa + fa.neighbor(indices[0], -1) + fa.neighbor(indices[0], +1)) / (dx ** 2)
        else:
            offsets = [(1, 1), [-1, 1], [1, -1], [-1, -1]]
            return sum(o1 * o2 * fa.neighbor(indices[0], o1).neighbor(indices[1], o2)
                       for o1, o2 in offsets) / (4 * dx ** 2)
    raise NotImplementedError("Supports only derivatives up to order 2")


def fd_stencils_isotropic(indices, dx, fa):
    dim = fa.field.spatial_dimensions
    if dim == 1:
        return fd_stencils_standard(indices, dx, fa)
    elif dim == 2:
        order = len(indices)

        if order == 1:
            idx = indices[0]
            assert 0 <= idx < 2
            other_idx = 1 if indices[0] == 0 else 0
            weights = {-1: sp.Rational(1, 12) / dx,
                       0: sp.Rational(1, 3) / dx,
                       1: sp.Rational(1, 12) / dx}
            upper_terms = sum(fa.neighbor(idx, +1).neighbor(other_idx, off) * w for off, w in weights.items())
            lower_terms = sum(fa.neighbor(idx, -1).neighbor(other_idx, off) * w for off, w in weights.items())
            return upper_terms - lower_terms
        elif order == 2:
            if indices[0] == indices[1]:
                idx = indices[0]
                other_idx = 1 if idx == 0 else 0
                diagonals = sp.Rational(1, 12) * sum(fa.neighbor(0, i).neighbor(1, j) for i in (-1, 1) for j in (-1, 1))
                div_direction = sp.Rational(5, 6) * sum(fa.neighbor(idx, i) for i in (-1, 1))
                other_direction = - sp.Rational(1, 6) * sum(fa.neighbor(other_idx, i) for i in (-1, 1))
                center = - sp.Rational(5, 3) * fa
                return (diagonals + div_direction + other_direction + center) / (dx ** 2)
            else:
                return fd_stencils_standard(indices, dx, fa)
    raise NotImplementedError("Supports only derivatives up to order 2 for 1D and 2D setups")


def fd_stencils_forth_order_isotropic(indices, dx, fa):
    order = len(indices)
    if order != 1:
        raise NotImplementedError("Forth order finite difference discretization is "
                                  "currently only supported for first derivatives")
    dim = indices[0]
    if dim not in (0, 1):
        raise NotImplementedError("Forth order finite difference discretization is only implemented for 2D")

    stencils = forth_order_2d_derivation()
    return stencils[dim].apply(fa) / dx


def fd_stencils_isotropic_high_density_code(indices, dx, fa):
    dim = fa.field.spatial_dimensions
    if dim == 1:
        return fd_stencils_standard(indices, dx, fa)
    elif dim == 2:
        order = len(indices)

        if order == 1:
            idx = indices[0]
            assert 0 <= idx < 2
            other_idx = 1 if indices[0] == 0 else 0
            weights = {-1: sp.Rational(1, 12) / dx,
                       0: sp.Rational(1, 3) / dx,
                       1: sp.Rational(1, 12) / dx}
            upper_terms = sum(fa.neighbor(idx, +1).neighbor(other_idx, off) * w for off, w in weights.items())
            lower_terms = sum(fa.neighbor(idx, -1).neighbor(other_idx, off) * w for off, w in weights.items())
            return upper_terms - lower_terms
        elif order == 2:
            if indices[0] == indices[1]:
                idx = indices[0]
                diagonals = sp.Rational(1, 8) * sum(fa.neighbor(0, i).neighbor(1, j) for i in (-1, 1) for j in (-1, 1))
                div_direction = sp.Rational(1, 2) * sum(fa.neighbor(idx, i) for i in (-1, 1))
                center = - sp.Rational(3, 2) * fa
                return (diagonals + div_direction + center) / (dx ** 2)
            else:
                return fd_stencils_standard(indices, dx, fa)
    raise NotImplementedError("Supports only derivatives up to order 2 for 1D and 2D setups")


def discretize_spatial(expr, dx, stencil=fd_stencils_standard):
    if isinstance(stencil, str):
        if stencil == 'standard':
            stencil = fd_stencils_standard
        elif stencil == 'isotropic':
            stencil = fd_stencils_isotropic
        elif stencil == 'isotropic_hd':
            stencil = fd_stencils_isotropic_high_density_code
        else:
            raise ValueError("Unknown stencil. Supported 'standard' and 'isotropic'")

    def visitor(e):
        if isinstance(e, Diff):
            arg, *indices = diff_args(e)
            if not isinstance(arg, Field.Access):
                raise ValueError("Only derivatives with field or field accesses as arguments can be discretized")
            return stencil(indices, dx, arg)
        else:
            new_args = [discretize_spatial(a, dx, stencil) for a in e.args]
            return e.func(*new_args) if new_args else e

    return generic_visit(expr, visitor)


def discretize_spatial_staggered(expr, dx, stencil=fd_stencils_standard):
    def staggered_visitor(e, coordinate, sign):
        if isinstance(e, Diff):
            arg, *indices = diff_args(e)
            if len(indices) != 1:
                raise ValueError("Function supports only up to second derivatives")
            if not isinstance(arg, Field.Access):
                raise ValueError("Argument of inner derivative has to be field access")
            target = indices[0]
            if target == coordinate:
                assert sign in (-1, 1)
                return (arg.neighbor(coordinate, sign) - arg) / dx * sign
            else:
                return (stencil(indices, dx, arg.neighbor(coordinate, sign))
                        + stencil(indices, dx, arg)) / 2
        elif isinstance(e, Field.Access):
            return (e.neighbor(coordinate, sign) + e) / 2
        elif isinstance(e, sp.Symbol):
            loop_idx = LoopOverCoordinate.is_loop_counter_symbol(e)
            return e + sign / 2 if loop_idx == coordinate else e
        else:
            new_args = [staggered_visitor(a, coordinate, sign) for a in e.args]
            return e.func(*new_args) if new_args else e

    def visitor(e):
        if isinstance(e, Diff):
            arg, *indices = diff_args(e)
            if isinstance(arg, Field.Access):
                return stencil(indices, dx, arg)
            else:
                if not len(indices) == 1:
                    raise ValueError("This term is not support by the staggered discretization strategy")
                target = indices[0]
                return (staggered_visitor(arg, target, 1) - staggered_visitor(arg, target, -1)) / dx
        else:
            new_args = [visitor(a) for a in e.args]
            return e.func(*new_args) if new_args else e

    return generic_visit(expr, visitor)


# -------------------------------------- special stencils --------------------------------------------------------------


@memorycache(maxsize=1)
def forth_order_2d_derivation() -> Tuple[FiniteDifferenceStencilDerivation.Result, ...]:
    # Symmetry, isotropy and 4th order conditions are not enough to fully specify the stencil
    # one weight has to be specifically set to a somewhat arbitrary value
    second_neighbor_weight = sp.Rational(1, 10)
    second_neighbor_stencil = [(i, j)
                               for i in (-2, -1, 0, 1, 2)
                               for j in (-2, -1, 0, 1, 2)
                               ]
    x_diff = FiniteDifferenceStencilDerivation((0,), second_neighbor_stencil)
    x_diff.set_weight((2, 0), second_neighbor_weight)
    x_diff.assume_symmetric(0, anti_symmetric=True)
    x_diff.assume_symmetric(1)
    x_diff_stencil = x_diff.get_stencil(isotropic=True)

    y_diff = FiniteDifferenceStencilDerivation((1,), second_neighbor_stencil)
    y_diff.set_weight((0, 2), second_neighbor_weight)
    y_diff.assume_symmetric(1, anti_symmetric=True)
    y_diff.assume_symmetric(0)
    y_diff_stencil = y_diff.get_stencil(isotropic=True)

    return x_diff_stencil, y_diff_stencil