"""
An experimental support for curvilinear grid.
"""

# TODO :
# see if tick_iterator method can be simplified by reusing the parent method.

import functools

import numpy as np

import matplotlib.patches as mpatches
from matplotlib.path import Path
from matplotlib.transforms import IdentityTransform
import matplotlib.axes as maxes

from mpl_toolkits.axes_grid1.parasite_axes import host_axes_class_factory

from . import axislines, grid_helper_curvelinear
from .axis_artist import AxisArtist
from .grid_finder import ExtremeFinderSimple


class FloatingAxisArtistHelper(
        grid_helper_curvelinear.FloatingAxisArtistHelper):
    pass


class FixedAxisArtistHelper(grid_helper_curvelinear.FloatingAxisArtistHelper):

    def __init__(self, grid_helper, side, nth_coord_ticks=None):
        """
        nth_coord = along which coordinate value varies.
         nth_coord = 0 ->  x axis, nth_coord = 1 -> y axis
        """
        value, nth_coord = grid_helper.get_data_boundary(side)
        super().__init__(grid_helper, nth_coord, value, axis_direction=side)
        if nth_coord_ticks is None:
            nth_coord_ticks = nth_coord
        self.nth_coord_ticks = nth_coord_ticks

        self.value = value
        self.grid_helper = grid_helper
        self._side = side

    def update_lim(self, axes):
        self.grid_helper.update_lim(axes)
        self.grid_info = self.grid_helper.grid_info

    def get_tick_iterators(self, axes):
        """tick_loc, tick_angle, tick_label, (optionally) tick_label"""

        grid_finder = self.grid_helper.grid_finder

        lat_levs, lat_n, lat_factor = self.grid_info["lat_info"]
        lon_levs, lon_n, lon_factor = self.grid_info["lon_info"]

        lon_levs, lat_levs = np.asarray(lon_levs), np.asarray(lat_levs)
        if lat_factor is not None:
            yy0 = lat_levs / lat_factor
            dy = 0.001 / lat_factor
        else:
            yy0 = lat_levs
            dy = 0.001

        if lon_factor is not None:
            xx0 = lon_levs / lon_factor
            dx = 0.001 / lon_factor
        else:
            xx0 = lon_levs
            dx = 0.001

        extremes = self.grid_helper._extremes
        xmin, xmax = sorted(extremes[:2])
        ymin, ymax = sorted(extremes[2:])

        def transform_xy(x, y):
            x1, y1 = grid_finder.transform_xy(x, y)
            x2, y2 = axes.transData.transform(np.array([x1, y1]).T).T
            return x2, y2

        if self.nth_coord == 0:
            mask = (ymin <= yy0) & (yy0 <= ymax)
            yy0 = yy0[mask]
            xx0 = np.full_like(yy0, self.value)
            xx1, yy1 = transform_xy(xx0, yy0)

            xx00 = xx0.astype(float, copy=True)
            xx00[xx0 + dx > xmax] -= dx
            xx1a, yy1a = transform_xy(xx00, yy0)
            xx1b, yy1b = transform_xy(xx00 + dx, yy0)

            yy00 = yy0.astype(float, copy=True)
            yy00[yy0 + dy > ymax] -= dy
            xx2a, yy2a = transform_xy(xx0, yy00)
            xx2b, yy2b = transform_xy(xx0, yy00 + dy)

            labels = self.grid_info["lat_labels"]
            labels = [l for l, m in zip(labels, mask) if m]

        elif self.nth_coord == 1:
            mask = (xmin <= xx0) & (xx0 <= xmax)
            xx0 = xx0[mask]
            yy0 = np.full_like(xx0, self.value)
            xx1, yy1 = transform_xy(xx0, yy0)

            yy00 = yy0.astype(float, copy=True)
            yy00[yy0 + dy > ymax] -= dy
            xx1a, yy1a = transform_xy(xx0, yy00)
            xx1b, yy1b = transform_xy(xx0, yy00 + dy)

            xx00 = xx0.astype(float, copy=True)
            xx00[xx0 + dx > xmax] -= dx
            xx2a, yy2a = transform_xy(xx00, yy0)
            xx2b, yy2b = transform_xy(xx00 + dx, yy0)

            labels = self.grid_info["lon_labels"]
            labels = [l for l, m in zip(labels, mask) if m]

        def f1():
            dd = np.arctan2(yy1b - yy1a, xx1b - xx1a)  # angle normal
            dd2 = np.arctan2(yy2b - yy2a, xx2b - xx2a)  # angle tangent
            mm = (yy1b - yy1a == 0) & (xx1b - xx1a == 0)  # mask not defined dd
            dd[mm] = dd2[mm] + np.pi / 2

            tick_to_axes = self.get_tick_transform(axes) - axes.transAxes
            for x, y, d, d2, lab in zip(xx1, yy1, dd, dd2, labels):
                c2 = tick_to_axes.transform((x, y))
                delta = 0.00001
                if 0-delta <= c2[0] <= 1+delta and 0-delta <= c2[1] <= 1+delta:
                    d1, d2 = np.rad2deg([d, d2])
                    yield [x, y], d1, d2, lab

        return f1(), iter([])

    def get_line(self, axes):
        self.update_lim(axes)
        k, v = dict(left=("lon_lines0", 0),
                    right=("lon_lines0", 1),
                    bottom=("lat_lines0", 0),
                    top=("lat_lines0", 1))[self._side]
        xx, yy = self.grid_info[k][v]
        return Path(np.column_stack([xx, yy]))


class ExtremeFinderFixed(ExtremeFinderSimple):
    def __init__(self, extremes):
        self._extremes = extremes

    def __call__(self, transform_xy, x1, y1, x2, y2):
        """
        get extreme values.

        x1, y1, x2, y2 in image coordinates (0-based)
        nx, ny : number of division in each axis
        """
        return self._extremes


class GridHelperCurveLinear(grid_helper_curvelinear.GridHelperCurveLinear):

    def __init__(self, aux_trans, extremes,
                 grid_locator1=None,
                 grid_locator2=None,
                 tick_formatter1=None,
                 tick_formatter2=None):
        # docstring inherited
        self._extremes = extremes
        extreme_finder = ExtremeFinderFixed(extremes)
        super().__init__(aux_trans,
                         extreme_finder,
                         grid_locator1=grid_locator1,
                         grid_locator2=grid_locator2,
                         tick_formatter1=tick_formatter1,
                         tick_formatter2=tick_formatter2)

    def get_data_boundary(self, side):
        """
        return v=0, nth=1
        """
        lon1, lon2, lat1, lat2 = self._extremes
        return dict(left=(lon1, 0),
                    right=(lon2, 0),
                    bottom=(lat1, 1),
                    top=(lat2, 1))[side]

    def new_fixed_axis(self, loc,
                       nth_coord=None,
                       axis_direction=None,
                       offset=None,
                       axes=None):
        if axes is None:
            axes = self.axes
        if axis_direction is None:
            axis_direction = loc
        # This is not the same as the FixedAxisArtistHelper class used by
        # grid_helper_curvelinear.GridHelperCurveLinear.new_fixed_axis!
        _helper = FixedAxisArtistHelper(
            self, loc, nth_coord_ticks=nth_coord)
        axisline = AxisArtist(axes, _helper, axis_direction=axis_direction)
        # Perhaps should be moved to the base class?
        axisline.line.set_clip_on(True)
        axisline.line.set_clip_box(axisline.axes.bbox)
        return axisline

    # new_floating_axis will inherit the grid_helper's extremes.

    # def new_floating_axis(self, nth_coord,
    #                       value,
    #                       axes=None,
    #                       axis_direction="bottom"
    #                       ):

    #     axis = super(GridHelperCurveLinear,
    #                  self).new_floating_axis(nth_coord,
    #                                          value, axes=axes,
    #                                          axis_direction=axis_direction)

    #     # set extreme values of the axis helper
    #     if nth_coord == 1:
    #         axis.get_helper().set_extremes(*self._extremes[:2])
    #     elif nth_coord == 0:
    #         axis.get_helper().set_extremes(*self._extremes[2:])

    #     return axis

    def _update_grid(self, x1, y1, x2, y2):
        if self.grid_info is None:
            self.grid_info = dict()

        grid_info = self.grid_info

        grid_finder = self.grid_finder
        extremes = grid_finder.extreme_finder(grid_finder.inv_transform_xy,
                                              x1, y1, x2, y2)

        lon_min, lon_max = sorted(extremes[:2])
        lat_min, lat_max = sorted(extremes[2:])
        lon_levs, lon_n, lon_factor = \
                  grid_finder.grid_locator1(lon_min, lon_max)
        lat_levs, lat_n, lat_factor = \
                  grid_finder.grid_locator2(lat_min, lat_max)
        grid_info["extremes"] = lon_min, lon_max, lat_min, lat_max  # extremes

        grid_info["lon_info"] = lon_levs, lon_n, lon_factor
        grid_info["lat_info"] = lat_levs, lat_n, lat_factor

        grid_info["lon_labels"] = grid_finder.tick_formatter1("bottom",
                                                              lon_factor,
                                                              lon_levs)

        grid_info["lat_labels"] = grid_finder.tick_formatter2("bottom",
                                                              lat_factor,
                                                              lat_levs)

        if lon_factor is None:
            lon_values = np.asarray(lon_levs[:lon_n])
        else:
            lon_values = np.asarray(lon_levs[:lon_n]/lon_factor)
        if lat_factor is None:
            lat_values = np.asarray(lat_levs[:lat_n])
        else:
            lat_values = np.asarray(lat_levs[:lat_n]/lat_factor)

        lon_lines, lat_lines = grid_finder._get_raw_grid_lines(
            lon_values[(lon_min < lon_values) & (lon_values < lon_max)],
            lat_values[(lat_min < lat_values) & (lat_values < lat_max)],
            lon_min, lon_max, lat_min, lat_max)

        grid_info["lon_lines"] = lon_lines
        grid_info["lat_lines"] = lat_lines

        lon_lines, lat_lines = grid_finder._get_raw_grid_lines(
            # lon_min, lon_max, lat_min, lat_max)
            extremes[:2], extremes[2:], *extremes)

        grid_info["lon_lines0"] = lon_lines
        grid_info["lat_lines0"] = lat_lines

    def get_gridlines(self, which="major", axis="both"):
        grid_lines = []
        if axis in ["both", "x"]:
            grid_lines.extend(self.grid_info["lon_lines"])
        if axis in ["both", "y"]:
            grid_lines.extend(self.grid_info["lat_lines"])
        return grid_lines

    def get_boundary(self):
        """
        Return (N, 2) array of (x, y) coordinate of the boundary.
        """
        x0, x1, y0, y1 = self._extremes
        tr = self._aux_trans

        xx = np.linspace(x0, x1, 100)
        yy0 = np.full_like(xx, y0)
        yy1 = np.full_like(xx, y1)
        yy = np.linspace(y0, y1, 100)
        xx0 = np.full_like(yy, x0)
        xx1 = np.full_like(yy, x1)

        xxx = np.concatenate([xx[:-1], xx1[:-1], xx[-1:0:-1], xx0])
        yyy = np.concatenate([yy0[:-1], yy[:-1], yy1[:-1], yy[::-1]])
        t = tr.transform(np.array([xxx, yyy]).transpose())

        return t


class FloatingAxesBase:

    def __init__(self, *args, **kwargs):
        grid_helper = kwargs.get("grid_helper", None)
        if grid_helper is None:
            raise ValueError("FloatingAxes requires grid_helper argument")
        if not hasattr(grid_helper, "get_boundary"):
            raise ValueError("grid_helper must implement get_boundary method")

        self._axes_class_floating.__init__(self, *args, **kwargs)

        self.set_aspect(1.)
        self.adjust_axes_lim()

    def _gen_axes_patch(self):
        """
        Returns the patch used to draw the background of the axes.  It
        is also used as the clipping path for any data elements on the
        axes.

        In the standard axes, this is a rectangle, but in other
        projections it may not be.

        .. note::
            Intended to be overridden by new projection types.
        """
        grid_helper = self.get_grid_helper()
        t = grid_helper.get_boundary()
        return mpatches.Polygon(t)

    def cla(self):
        self._axes_class_floating.cla(self)
        # HostAxes.cla(self)
        self.patch.set_transform(self.transData)

        patch = self._axes_class_floating._gen_axes_patch(self)
        patch.set_figure(self.figure)
        patch.set_visible(False)
        patch.set_transform(self.transAxes)

        self.patch.set_clip_path(patch)
        self.gridlines.set_clip_path(patch)

        self._original_patch = patch

    def adjust_axes_lim(self):
        grid_helper = self.get_grid_helper()
        t = grid_helper.get_boundary()
        x, y = t[:, 0], t[:, 1]

        xmin, xmax = min(x), max(x)
        ymin, ymax = min(y), max(y)

        dx = (xmax-xmin) / 100
        dy = (ymax-ymin) / 100

        self.set_xlim(xmin-dx, xmax+dx)
        self.set_ylim(ymin-dy, ymax+dy)


@functools.lru_cache(None)
def floatingaxes_class_factory(axes_class):
    return type("Floating %s" % axes_class.__name__,
                (FloatingAxesBase, axes_class),
                {'_axes_class_floating': axes_class})


FloatingAxes = floatingaxes_class_factory(
    host_axes_class_factory(axislines.Axes))
FloatingSubplot = maxes.subplot_class_factory(FloatingAxes)