import numpy

from amuse.units import units

from amuse.units.quantities import value_in, to_quantity

from amuse.datamodel import (
    UnstructuredGrid,
    StructuredBaseGrid,
    RegularBaseGrid,
)

try:
    import matplotlib
    from matplotlib import tri

    if not hasattr(tri, "LinearTriInterpolator"):
        raise Exception("LinearTriInterpolator not in matplotlib.tri")
    matplotlib_available = True
except:
    matplotlib_available = False

from warnings import warn


class interpolating_2D_remapper(object):
    def __init__(self, source, target, axes_names=None):
        """this class maps a source grid to a target grid using linear
        interpolation on a triangulation generated by adding a
        midpoint to every cell (source should be a structured grid)
        and thus generating 4 triangles for each cell. Values of the
        midpoints are averaged from the corners.
        """
        if len(source.shape) != 2:
            raise Exception("source grid is not 2D")
        if not isinstance(source, StructuredBaseGrid):
            raise Exception("source grid is not instance of StructuredBaseGrid")

        self.source = source
        self.target = target
        self._axes_names = list(axes_names or source.get_axes_names())
        self.generate_triangulation()

    def _generate_nodes(self, grid, attributes):

        Nx, Ny = grid.shape

        x, y = numpy.mgrid[0:Nx, 0:Ny]
        x1, y1 = numpy.mgrid[0:Nx - 1, 0:Ny - 1]

        x_ = x.flatten()
        y_ = y.flatten()
        x1_ = x1.flatten()
        y1_ = y1.flatten()

        l1 = Nx * Ny

        nodes = UnstructuredGrid(len(x_) + len(x1_))
        for name in attributes:
            values1 = getattr(grid, name)[x_, y_]
            values2 = (
                getattr(grid, name)[x1_, y1_]
                + getattr(grid, name)[x1_ + 1, y1_]
                + getattr(grid, name)[x1_, y1_ + 1]
                + getattr(grid, name)[x1_ + 1, y1_ + 1]
            )
            setattr(nodes[0], name, 0.0 * values1[0])
            setattr(nodes[:l1], name, 1.0 * values1)
            setattr(nodes[l1:], name, values2 / 4)
        return nodes

    def _generate_elements_and_boundaries(self, grid):
        Nx, Ny = grid.shape

        l1 = Nx * Ny

        i = numpy.arange(Nx * Ny).reshape((Nx, Ny))
        i1 = (numpy.arange((Nx - 1) * (Ny - 1)) + l1).reshape((Nx - 1, Ny - 1))

        e1 = numpy.zeros(((Nx - 1) * (Ny - 1), 3), dtype="i")
        e2 = numpy.zeros(((Nx - 1) * (Ny - 1), 3), dtype="i")
        e3 = numpy.zeros(((Nx - 1) * (Ny - 1), 3), dtype="i")
        e4 = numpy.zeros(((Nx - 1) * (Ny - 1), 3), dtype="i")

        e1[:, 0] = i[:-1, :-1].flatten()
        e1[:, 1] = i[1:, :-1].flatten()
        e1[:, 2] = i1[:, :].flatten()

        e2[:, 0] = i[1:, :-1].flatten()
        e2[:, 1] = i[1:, 1:].flatten()
        e2[:, 2] = i1[:, :].flatten()

        e3[:, 0] = i[1:, 1:].flatten()
        e3[:, 1] = i[:-1, 1:].flatten()
        e3[:, 2] = i1[:, :].flatten()

        e4[:, 0] = i[:-1, :-1].flatten()
        e4[:, 1] = i1[:, :].flatten()
        e4[:, 2] = i[:-1, 1:].flatten()

        elements = numpy.zeros((4 * (Nx - 1) * (Ny - 1), 3), dtype="i8")
        elements[0::4, :] = e1
        elements[1::4, :] = e2
        elements[2::4, :] = e3
        elements[3::4, :] = e4

        boundaries = [
            xx.flatten() for xx in [i[:, 0], i[-1, :], i[::-1, -1], i[0, ::-1]]
        ]

        elem = UnstructuredGrid(len(elements))
        elem.nodes = elements

        return elem, boundaries

    def convert_grid_to_nodes_and_elements(self, grid, attributes=None):

        if attributes is None:
            attributes = grid.get_attribute_names_defined_in_store()

        nodes = self._generate_nodes(grid, attributes)
        elements, boundaries = self._generate_elements_and_boundaries(grid)

        return nodes, elements, boundaries

    def generate_triangulation(self):

        nodes, elements, boundaries = self.convert_grid_to_nodes_and_elements(
            self.source, self._axes_names
        )

        xpos = to_quantity(getattr(nodes, self._axes_names[0]))
        ypos = to_quantity(getattr(nodes, self._axes_names[1]))

        self._xpos_unit = xpos.unit
        xpos = xpos.number
        self._ypos_unit = ypos.unit
        ypos = ypos.number

        n1 = elements.nodes[:, 0]
        n2 = elements.nodes[:, 1]
        n3 = elements.nodes[:, 2]
        elem = numpy.column_stack((n1, n2, n3))

        self._triangulation = tri.Triangulation(xpos, ypos, elem)

    def sample(self, values, xpos, ypos):
        interpolator = tri.LinearTriInterpolator(self._triangulation, values)
        return interpolator(xpos, ypos)

    def forward_mapping(self, attributes, target_names=None):
        if attributes is None:
            attributes = self.source.get_attribute_names_defined_in_store()
        if target_names is None:
            target_names = attributes

        source = self.source.empty_copy()
        channel1 = self.source.new_channel_to(source)
        target = self.target.empty_copy()
        channel2 = self.target.new_channel_to(target)
        channel3 = target.new_channel_to(self.target)

        channel1.copy_attributes(attributes)
        channel2.copy_attributes(self._axes_names)

        nodes = self._generate_nodes(source, attributes)

        xpos = value_in(getattr(target, self._axes_names[0]), self._xpos_unit)
        ypos = value_in(getattr(target, self._axes_names[1]), self._ypos_unit)

        for attribute, target_name in zip(attributes, target_names):
            values = to_quantity(getattr(nodes, attribute))
            unit = values.unit
            values = values.number
            samples = self.sample(values, xpos, ypos)
            setattr(
                target,
                target_name,
                (samples if unit is units.none else (samples | unit)),
            )

        channel3.copy_attributes(target_names)


class bilinear_2D_remapper(object):
    def __init__(
        self, source, target, check_inside=True, do_slices=False, x_periodic=False
    ):
        """this class maps a source grid to a target grid using bilinear
        interpolation. If check_inside=True, raise exception if any
        target point outside source grid. If the grids are 3 dimensional
        it can be used to do a 2D interpolation of each level if shapes
        are the same and positions are the same in the 3rd dimension.
        """
        if len(source.shape) != 2 and not do_slices:
            raise Exception(
                "source grid is not 2D, set do_slices=True remapping by slices"
            )
        if len(source.shape) != len(target.shape):
            raise Exception("incompatible shapes")
        if len(source.shape) != 2:
            if numpy.any(source.shape[2:] != target.shape[2:]):
                raise Exception("source and target need same shapes (after dim 2)")
        if not isinstance(source, RegularBaseGrid):
            raise Exception(
                f"source grid ({type(source)}) is not instance of RegularBaseGrid"
            )

        self.source = source
        self.target = target
        self._axes_names = source.get_axes_names()
        if len(source.shape) != 2:
            for x in self._axes_names[2:]:
                if numpy.any(getattr(source[0, 0], x) != getattr(target[0, 0], x)):
                    print(getattr(source[0, 0], x))
                    print(getattr(target[0, 0], x))
                    warn(f"positions (possibly) not the same on axes {x}")
        self.check_inside = check_inside
        self.x_periodic = x_periodic
        self._weights = None
        self._indices = None

    def _calculate_weights(self):
        x0 = to_quantity(getattr(self.source[0, 0], self._axes_names[0]))
        x1 = to_quantity(getattr(self.source[1, 1], self._axes_names[0]))
        y0 = to_quantity(getattr(self.source[0, 0], self._axes_names[1]))
        y1 = to_quantity(getattr(self.source[1, 1], self._axes_names[1]))

        # guaranteed by grid being RegularBaseGrid
        assert x0.max() == x0.min()
        assert x1.max() == x1.min()
        assert y0.max() == y0.min()
        assert y1.max() == y1.min()

        x0 = x0.min()
        x1 = x1.min()
        y0 = y0.min()
        y1 = y1.min()

        dx = x1 - x0
        dy = y1 - y0

        x = getattr(self.target, self._axes_names[0])
        y = getattr(self.target, self._axes_names[1])

        ix = numpy.floor((x - x0) / dx).astype(int)
        iy = numpy.floor((y - y0) / dy).astype(int)

        ix2 = ix + 1
        iy2 = iy + 1

        if self.x_periodic:
            numpy.place(ix, ix == -1, self.source.shape[0] - 1)
            numpy.place(ix, ix == self.source.shape[0], 0)
            numpy.place(ix2, ix2 == self.source.shape[0], 0)

        if self.check_inside:
            if (
                numpy.any(ix < 0)
                or numpy.any(ix > self.source.shape[0] - 1)
                or numpy.any(iy < 0)
                or numpy.any(iy > self.source.shape[1] - 1)
                or numpy.any(ix2 < 0)
                or numpy.any(ix2 > self.source.shape[0] - 1)
                or numpy.any(iy2 < 0)
                or numpy.any(iy2 > self.source.shape[1] - 1)
            ):
                raise Exception(
                    "target not fully inside (restricted) source grid as required"
                )

        if not self.x_periodic:
            ix = numpy.clip(ix, 0, self.source.shape[0] - 1)
            ix2 = numpy.clip(ix2, 0, self.source.shape[0] - 1)

        iy = numpy.clip(iy, 0, self.source.shape[1] - 1)
        iy2 = numpy.clip(iy2, 0, self.source.shape[1] - 1)

        wx = (x0 + (ix + 1) * dx - x) / dx
        wy = (y0 + (iy + 1) * dy - y) / dy
        wx = numpy.clip(wx, 0.0, 1.0)
        wy = numpy.clip(wy, 0.0, 1.0)

        self._weights = [wx, wy]

        while len(ix.shape) > 2:
            ix = numpy.amax(ix, axis=-1)
            iy = numpy.amax(iy, axis=-1)
            ix2 = numpy.amax(ix2, axis=-1)
            iy2 = numpy.amax(iy2, axis=-1)

        self._indices = [ix, ix2, iy, iy2]

    def _evaluate(self, values):
        ix, ix2, iy, iy2 = self._indices
        wx, wy = self._weights
        result = (
            wx * wy * values[ix, iy]
            + (1.0 - wx) * wy * values[ix2, iy]
            + wx * (1.0 - wy) * values[ix, iy2]
            + (1.0 - wx) * (1.0 - wy) * values[ix2, iy2]
        )
        return result

    def forward_mapping(self, attributes, target_names=None):
        if attributes is None:
            attributes = self.source.get_attribute_names_defined_in_store()
        if target_names is None:
            target_names = attributes
        if self._weights is None:
            self._calculate_weights()

        mapped_values = []
        for attribute, target_name in zip(attributes, target_names):
            values = getattr(self.source, attribute)
            samples = self._evaluate(values)
            mapped_values.append(samples)

        self.target.set_values_in_store(None, target_names, mapped_values)


class nearest_2D_remapper(object):
    def __init__(self, source, target, check_inside=True):
        """this class maps a source grid to a target grid getting closest
        grid value. If check_inside=True, raise exception if any
            target point outside source grid.
        """
        if len(source.shape) != 2:
            raise Exception("source grid is not 2D")
        if not isinstance(source, RegularBaseGrid):
            raise Exception("source grid is not instance of RegularBaseGrid")

        self.source = source
        self.target = target
        self._axes_names = source.get_axes_names()
        self.check_inside = check_inside
        self._indices = None

    def _calculate_weights(self):
        x = getattr(self.target, self._axes_names[0])
        y = getattr(self.target, self._axes_names[1])

        kwargs = {self._axes_names[0]: x, self._axes_names[1]: y}
        indices = self.source.get_index(**kwargs)

        ix = indices[..., 0]
        iy = indices[..., 1]
        if self.check_inside:
            if (
                numpy.any(ix < 0)
                or numpy.any(ix > self.source.shape[0] - 1)
                or numpy.any(iy < 0)
                or numpy.any(iy > self.source.shape[1] - 1)
            ):
                raise Exception("target not fully inside source grid as required")
        ix = numpy.clip(ix, 0, self.source.shape[0] - 1)
        iy = numpy.clip(iy, 0, self.source.shape[1] - 1)

        self._indices = [ix, iy]

    def _evaluate(self, values):
        return values[self._indices[0], self._indices[1]]

    def forward_mapping(self, attributes, target_names=None):
        if attributes is None:
            attributes = self.source.get_attribute_names_defined_in_store()
        if target_names is None:
            target_names = attributes
        if self._indices is None:
            self._calculate_weights()

        mapped_values = []
        for attribute, target_name in zip(attributes, target_names):
            values = getattr(self.source, attribute)
            samples = self._evaluate(values)
            mapped_values.append(samples)

        self.target.set_values_in_store(None, target_names, mapped_values)


def conservative_spherical_remapper(*args, **kwargs):
    raise Exception("conservative_spherical_remapper has moved to omuse.ext")
