from __future__ import annotations

from typing import List, Union

import numba as nb
import numpy as np
import xarray as xr
from numba_celltree.celltree_base import default_tolerance
from numba_celltree.constants import Point
from numba_celltree.geometry_utils import points_in_triangles

import xugrid
from xugrid.constants import FloatArray, IntArray, MissingOptionalModule

try:
    import shapely

except ImportError:
    shapely = MissingOptionalModule("shapely")

try:
    import mapbox_earcut

except ImportError:
    mapbox_earcut = MissingOptionalModule("mapbox_earcut")


@nb.njit(inline="always")
def in_bounds(p: Point, a: Point, b: Point) -> bool:
    """
    Check whether point p falls within the bounding box created by a and b
    (after we've checked the size of the cross product).

    However, we must take into account that a line may be either vertical
    (dx=0) or horizontal (dy=0) and only evaluate the non-zero value.

    If the area created by p, a, b is tiny AND p is within the bounds of a and
    b, the point lies very close to the edge.
    """
    # Already in numba_celltree, unreleased.
    dx = b.x - a.x
    dy = b.y - a.y
    if abs(dx) >= abs(dy):
        if dx > 0:
            return a.x <= p.x and p.x <= b.x
        return b.x <= p.x and p.x <= a.x
    else:
        if dy > 0:
            return a.y <= p.y and p.y <= b.y
        return b.y <= p.y and p.y <= a.y


def _locate_polygon(
    grid: "xu.Ugrid2d",  # type: ignore # noqa
    exterior: FloatArray,
    interiors: List[FloatArray],
    all_touched: bool,
) -> IntArray:
    """
    Locate a single polygon.

    This algorithm burns a polygon vector geometry in a 2d topology by:

    * Extracting the exterior and interiors (holes) coordinates from the
      polygon.
    * Breaking every polygon down into a triangles using an "earcut" algorithm.
    * Searching the grid for these triangles.

    Due to the use of the separating axes theorem, _locate_faces finds all
    intersecting triangles, including those who only touch the edge.
    To enable all_touched=False, we have to search the centroids of candidates
    in the intersecting triangles.

    Parameters
    ----------
    grid: Ugrid2d
    exterior: FloatArray
        Exterior of the polygon.
    interiors: List[FloatArray]
        Interior holes of the polygon.
    all_touched: bool
        Whether to include include cells whose centroid falls inside, of every
        cell that is touched.
    """

    import mapbox_earcut

    rings = np.cumsum([len(exterior)] + [len(interior) for interior in interiors])
    vertices = np.vstack([exterior] + interiors).astype(np.float64)
    triangles = mapbox_earcut.triangulate_float64(vertices, rings).reshape((-1, 3))
    triangle_indices, grid_indices = grid.celltree.locate_faces(vertices, triangles)
    if all_touched:
        return grid_indices
    else:
        tolerance = default_tolerance(grid.celltree.bb_distances[:, 2])
        centroids = grid.centroids[grid_indices]
        inside = points_in_triangles(
            points=centroids,
            face_indices=triangle_indices,
            faces=triangles,
            vertices=vertices,
            tolerance=tolerance,
        )
        return grid_indices[inside]


def _burn_polygons(
    polygons: "geopandas.GeoSeries",  # type: ignore # noqa
    like: "xugrid.Ugrid2d",
    values: np.ndarray,
    all_touched: bool,
    output: FloatArray,
) -> None:
    exterior_coordinates = [
        shapely.get_coordinates(exterior) for exterior in polygons.exterior
    ]
    interior_coordinates = [
        [shapely.get_coordinates(p_interior) for p_interior in p_interiors]
        for p_interiors in polygons.interiors
    ]
    to_burn = np.empty(like.n_face, dtype=bool)

    for exterior, interiors, value in zip(
        exterior_coordinates, interior_coordinates, values
    ):
        to_burn = _locate_polygon(like, exterior, interiors, all_touched)
        output[to_burn] = value

    return


def _burn_points(
    points: "geopandas.GeoSeries",  # type: ignore # noqa
    like: "xugrid.Ugrid2d",
    values: np.ndarray,
    output: FloatArray,
) -> None:
    """Simply searches the points in the ``like`` 2D topology."""
    xy = shapely.get_coordinates(points)
    to_burn = like.locate_points(xy)
    output[to_burn] = values
    return


def _burn_lines(
    lines: "geopandas.GeoSeries",  # type: ignore # noqa
    like: "xugrid.Ugrid2d",
    values: np.ndarray,
    output: FloatArray,
) -> None:
    """
    Burn the line values into the underlying faces.

    This algorithm breaks any linestring down into edges (two x, y points). We
    search and intersect every edge in the ``like`` grid, the intersections are
    discarded.
    """
    xy, index = shapely.get_coordinates(lines, return_index=True)
    # From the coordinates and the index, create the (n_edge, 2, 2) shape array
    # containing the edge_coordinates.
    linear_index = np.arange(index.size)
    segments = np.column_stack([linear_index[:-1], linear_index[1:]])
    # Only connections with vertices with the same index are valid.
    valid = np.diff(index) == 0
    segments = segments[valid]
    edges = xy[segments]
    # Now query the grid for these edges.
    edge_index, face_index, _ = like.intersect_edges(edges)
    # Find the associated values.
    line_index = index[1:][valid]
    value_index = line_index[edge_index]
    output[face_index] = values[value_index]
    return


def burn_vector_geometry(
    gdf: "geopandas.GeoDataframe",  # type: ignore # noqa
    like: Union["xugrid.Ugrid2d", "xugrid.UgridDataArray", "xugrid.UgridDataset"],
    column: str | None = None,
    fill: Union[int, float] = np.nan,
    all_touched: bool = False,
) -> xugrid.UgridDataArray:
    """
    Burn vector geometries (points, lines, polygons) into a Ugrid2d mesh.

    If no ``column`` argument is provided, a value of 1.0 will be burned in to
    the mesh.

    Parameters
    ----------
    gdf: geopandas.GeoDataFrame
        Polygons, points, and/or lines to be burned into the grid.
    like: UgridDataArray, UgridDataset, or Ugrid2d
        Grid to burn the vector data into.
    column: str, optional
        Name of the geodataframe column of which to the values to burn into
        grid.
    fill: int, float, optional, default value ``np.nan``.
        Fill value for nodata areas.
    all_touched: bool, optional, default value ``False``.
        All mesh faces (cells) touched by polygons will be updated, not just
        those whose center point is within the polygon.

    Returns
    -------
    burned: UgridDataArray
    """
    import geopandas as gpd

    POINT = shapely.GeometryType.POINT
    LINESTRING = shapely.GeometryType.LINESTRING
    LINEARRING = shapely.GeometryType.LINEARRING
    POLYGON = shapely.GeometryType.POLYGON
    GEOM_NAMES = {v: k for k, v in shapely.GeometryType.__members__.items()}

    if not isinstance(gdf, gpd.GeoDataFrame):
        raise TypeError(f"gdf must be GeoDataFrame, received: {type(gdf).__name__}")
    if isinstance(like, (xugrid.UgridDataArray, xugrid.UgridDataset)):
        like = like.ugrid.grid
    if not isinstance(like, xugrid.Ugrid2d):
        raise TypeError(
            "Like must be Ugrid2d, UgridDataArray, or UgridDataset;"
            f"received: {type(like).__name__}"
        )
    geometry_id = shapely.get_type_id(gdf.geometry)
    allowed_types = (POINT, LINESTRING, LINEARRING, POLYGON)
    if not np.isin(geometry_id, allowed_types).all():
        received = ", ".join(
            [GEOM_NAMES[geom_id] for geom_id in np.unique(geometry_id)]
        )
        raise TypeError(
            "GeoDataFrame contains unsupported geometry types. Can only burn "
            "Point, LineString, LinearRing, and Polygon geometries. Received: "
            f"{received}"
        )

    points = gdf.loc[geometry_id == POINT]
    lines = gdf.loc[(geometry_id == LINESTRING) | (geometry_id == LINEARRING)]
    polygons = gdf.loc[geometry_id == POLYGON]

    if column is None:
        point_values = np.ones(len(points), dtype=float)
        line_values = np.ones(len(lines), dtype=float)
        poly_values = np.ones(len(polygons), dtype=float)
    else:
        point_values = points[column].to_numpy()
        line_values = lines[column].to_numpy()
        poly_values = polygons[column].to_numpy()

    output = np.full(like.n_face, fill)
    if len(polygons) > 0:
        _burn_polygons(polygons.geometry, like, poly_values, all_touched, output)
    if len(lines) > 0:
        _burn_lines(lines.geometry, like, line_values, output)
    if len(points) > 0:
        _burn_points(points.geometry, like, point_values, output)

    return xugrid.UgridDataArray(
        obj=xr.DataArray(output, dims=[like.face_dimension], name=column),
        grid=like,
    )


def grid_from_earcut_polygons(
    polygons: "geopandas.GeoDataFrame",  # type: ignore # noqa
    return_index: bool = False,
):
    import geopandas as gpd

    if not isinstance(polygons, gpd.GeoDataFrame):
        raise TypeError(f"Expected GeoDataFrame, received: {type(polygons).__name__}")

    geometry = polygons.geometry
    POLYGON = shapely.GeometryType.POLYGON
    GEOM_NAMES = {v: k for k, v in shapely.GeometryType.__members__.items()}
    geometry_id = shapely.get_type_id(geometry)

    if not (geometry_id == POLYGON).all():
        GEOM_NAMES = {v: k for k, v in shapely.GeometryType.__members__.items()}
        received = ", ".join(
            [GEOM_NAMES[geom_id] for geom_id in np.unique(geometry_id)]
        )
        raise TypeError(
            "geometry contains unsupported geometry types. Can only triangulate "
            f"Polygon geometries. Received: {received}"
        )

    # Shapely does not provide a vectorized manner to get the interior rings
    # easily, an index argument is always required, which is a poor fit for
    # heterogeneous polygons (where some have no holes, and some may have
    # hundreds).
    # map_box_earcut is also not vectorized for polygons, so we simply loop
    # over the geometries here.
    exterior_coordinates = [
        shapely.get_coordinates(exterior) for exterior in geometry.exterior
    ]
    interior_coordinates = [
        [shapely.get_coordinates(p_interior) for p_interior in p_interiors]
        for p_interiors in geometry.interiors
    ]

    all_triangles = []
    offset = 0
    for exterior, interiors in zip(exterior_coordinates, interior_coordinates):
        rings = np.cumsum([len(exterior)] + [len(interior) for interior in interiors])
        vertices = np.vstack([exterior] + interiors).astype(np.float64)
        triangles = mapbox_earcut.triangulate_float64(vertices, rings).reshape((-1, 3))
        all_triangles.append(triangles + offset)
        offset += len(vertices)

    face_nodes = np.concatenate(all_triangles, dtype=int).reshape((-1, 3))
    all_vertices = shapely.get_coordinates(geometry)
    node_x = all_vertices[:, 0]
    node_y = all_vertices[:, 1]
    grid = xugrid.Ugrid2d(node_x, node_y, -1, face_nodes)

    if return_index:
        n_triangles = [len(triangles) for triangles in all_triangles]
        index = np.repeat(np.arange(len(geometry)), n_triangles)
        return grid, index
    else:
        return grid


def earcut_triangulate_polygons(
    polygons: "geopandas.GeoDataframe",  # type: ignore # noqa
    column: str | None = None,
) -> xugrid.UgridDataArray:
    """
    Break down polygons using mapbox_earcut, and create a mesh from the
    resulting triangles.

    If no ``column`` argument is provided, the polygon index will be assigned
    to the grid faces.

    Parameters
    ----------
    polygons: geopandas.GeoDataFrame
        Polygons to convert to triangles.
    column: str, optional
        Name of the geodataframe column of which to the values to assign
        to the grid faces.

    Returns
    -------
    triangulated: UgridDataArray
    """
    grid, index = grid_from_earcut_polygons(polygons, return_index=True)

    if column is not None:
        da = (
            polygons[column]
            .reset_index(drop=True)
            .to_xarray()
            .isel(index=index)
            .rename({"index": grid.face_dimension})
        )
    else:
        da = xr.DataArray(data=index, dims=(grid.face_dimension,))

    return xugrid.UgridDataArray(da, grid)
