"""Functions for polygon geometries."""

import geopandas as gpd
import networkx as nx
import numpy as np
import pandas as pd
from geopandas import GeoDataFrame, GeoSeries
from shapely import (
    area,
    get_exterior_ring,
    get_interior_ring,
    get_num_interior_rings,
    get_parts,
    polygons,
)
from shapely.ops import unary_union

from .general import _push_geom_col, to_lines
from .neighbors import get_neighbor_indices
from .overlay import clean_overlay


def eliminate_by_longest(
    gdf: GeoDataFrame,
    min_area: int | float,
    ignore_index: bool = False,
    aggfunc: str | dict | list = "first",
    **kwargs,
) -> GeoDataFrame:
    """Dissolves small polygons with the longest bordering neighbor polygon.

    Eliminates small geometries by dissolving them with the neighboring
    polygon with the longest shared border. The index and column values of the
    large polygons will be kept, unless else is specified.

    Args:
        gdf: GeoDataFrame with polygon geometries.
        min_area: minimum area for the polygons to be eliminated.
        ignore_index: If False (default), the resulting GeoDataFrame will keep the
            index of the large polygons. If True, the resulting axis will be labeled
            0, 1, …, n - 1.
        aggfunc: Aggregation function(s) to use when dissolving. Defaults to 'first',
            meaning the column values of the large polygons are kept.
        kwargs: Keyword arguments passed to the dissolve method.

    Returns:
        The GeoDataFrame with the small polygons dissolved into the large polygons.
    """
    if not ignore_index:
        idx_mapper = {i: idx for i, idx in enumerate(gdf.index)}
        idx_name = gdf.index.name

    gdf = gdf.reset_index(drop=True)

    small = gdf.loc[gdf.area <= min_area].assign(small_idx=lambda x: x.index)
    large = gdf.loc[gdf.area > min_area].assign(large_idx=lambda x: x.index)

    lines = to_lines(small[["small_idx", "geometry"]], large[["large_idx", "geometry"]])
    lines = lines[lines["small_idx"].notna()]
    lines["length__"] = lines.length

    longest = lines.sort_values("length__", ascending=False).drop_duplicates(
        "small_idx"
    )

    small_to_large = longest.set_index("small_idx")["large_idx"]
    small["dissolve_idx"] = small["small_idx"].map(small_to_large)
    large["dissolve_idx"] = large["large_idx"]

    kwargs.pop("as_index", None)
    eliminated = (
        pd.concat([large, small])
        .dissolve("dissolve_idx", aggfunc=aggfunc, **kwargs)
        .drop(
            ["length__", "small_idx", "large_idx"],
            axis=1,
            errors="ignore",
        )
    )

    if ignore_index:
        return eliminated.reset_index(drop=True)

    eliminated.index = eliminated.index.map(idx_mapper)
    eliminated.index.name = idx_name

    return eliminated


def eliminate_by_largest(
    gdf: GeoDataFrame,
    min_area: int | float,
    ignore_index: bool = False,
    aggfunc: str | dict | list = "first",
    **kwargs,
) -> GeoDataFrame:
    """Dissolves small polygons with the largest neighbor polygon.

    Eliminates small geometries by dissolving them with the neighboring
    polygon with the largest area. The index and column values of the
    large polygons will be kept, unless else is specified.

    Args:
        gdf: GeoDataFrame with polygon geometries.
        min_area: minimum area for the polygons to be eliminated.
        ignore_index: If False (default), the resulting GeoDataFrame will keep the
            index of the large polygons. If True, the resulting axis will be labeled
            0, 1, …, n - 1.
        aggfunc: Aggregation function(s) to use when dissolving. Defaults to 'first',
            meaning the column values of the large polygons are kept.
        kwargs: Keyword arguments passed to the dissolve method.

    Returns:
        The GeoDataFrame with the small polygons dissolved into the large polygons.
    """
    return _eliminate_by_area(
        gdf,
        min_area=min_area,
        ignore_index=ignore_index,
        sort_ascending=False,
        aggfunc=aggfunc,
        **kwargs,
    )


def eliminate_by_smallest(
    gdf: GeoDataFrame,
    min_area: int | float,
    ignore_index: bool = False,
    aggfunc: str | dict | list = "first",
    **kwargs,
) -> GeoDataFrame:
    return _eliminate_by_area(
        gdf,
        min_area=min_area,
        ignore_index=ignore_index,
        sort_ascending=True,
        aggfunc=aggfunc,
        **kwargs,
    )


def _eliminate_by_area(
    gdf: GeoDataFrame,
    min_area: int | float,
    sort_ascending: bool,
    ignore_index: bool = False,
    aggfunc="first",
    **kwargs,
) -> GeoDataFrame:
    if not ignore_index:
        idx_mapper = {i: idx for i, idx in enumerate(gdf.index)}
        idx_name = gdf.index.name

    gdf = gdf.reset_index(drop=True)

    small = gdf.loc[gdf.area <= min_area]
    large = gdf.loc[gdf.area > min_area]
    large["area__"] = large.area

    joined = small.sjoin(
        large[["area__", "geometry"]], predicate="touches"
    ).sort_values("area__", ascending=sort_ascending)

    largest = joined[~joined.index.duplicated()]

    large = large.assign(index_right=lambda x: x.index)

    kwargs.pop("as_index", None)
    eliminated = (
        pd.concat([large, largest])
        .dissolve("index_right", aggfunc=aggfunc, **kwargs)
        .drop(["area__"], axis=1, errors="ignore")
    )

    if ignore_index:
        return eliminated.reset_index(drop=True)

    eliminated.index = eliminated.index.map(idx_mapper)
    eliminated.index.name = idx_name

    return eliminated


def get_polygon_clusters(
    *gdfs: GeoDataFrame | GeoSeries,
    cluster_col: str = "cluster",
    allow_multipart: bool = False,
) -> GeoDataFrame | tuple[GeoDataFrame]:
    """Find which polygons overlap without dissolving.

    Devides polygons into clusters in a fast and precice manner by using spatial join
    and networkx to find the connected components, i.e. overlapping geometries.
    If multiple GeoDataFrames are given, the clusters will be based on all
    combined.

    This can be used instead of dissolve+explode, or before dissolving by the cluster
    column. This has been tested to be a lot faster if there are many
    non-overlapping polygons, but somewhat slower than dissolve+explode if most
    polygons overlap.

    Args:
        gdfs: One or more GeoDataFrames of polygons.
        cluster_col: Name of the resulting cluster column.
        allow_multipart: Whether to allow mutipart geometries in the gdfs.
            Defaults to False to avoid confusing results.

    Returns:
        One or more GeoDataFrames (same amount as was given) with a new cluster column.

    Examples
    --------

    Create geometries with three clusters of overlapping polygons.

    >>> import sgis as sg
    >>> gdf = sg.to_gdf([(0, 0), (1, 1), (0, 1), (4, 4), (4, 3), (7, 7)])
    >>> buffered = sg.buff(gdf, 1)
    >>> gdf
                                                geometry
    0  POLYGON ((1.00000 0.00000, 0.99951 -0.03141, 0...
    1  POLYGON ((2.00000 1.00000, 1.99951 0.96859, 1....
    2  POLYGON ((1.00000 1.00000, 0.99951 0.96859, 0....
    3  POLYGON ((5.00000 4.00000, 4.99951 3.96859, 4....
    4  POLYGON ((5.00000 3.00000, 4.99951 2.96859, 4....
    5  POLYGON ((8.00000 7.00000, 7.99951 6.96859, 7....

    Add a cluster column to the GeoDataFrame:

    >>> gdf = sg.get_polygon_clusters(gdf, cluster_col="cluster")
    >>> gdf
       cluster                                           geometry
    0        0  POLYGON ((1.00000 0.00000, 0.99951 -0.03141, 0...
    1        0  POLYGON ((2.00000 1.00000, 1.99951 0.96859, 1....
    2        0  POLYGON ((1.00000 1.00000, 0.99951 0.96859, 0....
    3        1  POLYGON ((5.00000 4.00000, 4.99951 3.96859, 4....
    4        1  POLYGON ((5.00000 3.00000, 4.99951 2.96859, 4....
    5        2  POLYGON ((8.00000 7.00000, 7.99951 6.96859, 7....

    If multiple GeoDataFrames are given, all are returned with common
    cluster values.

    >>> gdf2 = sg.to_gdf([(0, 0), (7, 7)])
    >>> gdf, gdf2 = sg.get_polygon_clusters(gdf, gdf2, cluster_col="cluster")
    >>> gdf2
       cluster                 geometry
    0        0  POINT (0.00000 0.00000)
    1        2  POINT (7.00000 7.00000)
    >>> gdf
       cluster                                           geometry
    0        0  POLYGON ((1.00000 0.00000, 0.99951 -0.03141, 0...
    1        0  POLYGON ((2.00000 1.00000, 1.99951 0.96859, 1....
    2        0  POLYGON ((1.00000 1.00000, 0.99951 0.96859, 0....
    3        1  POLYGON ((5.00000 4.00000, 4.99951 3.96859, 4....
    4        1  POLYGON ((5.00000 3.00000, 4.99951 2.96859, 4....
    5        2  POLYGON ((8.00000 7.00000, 7.99951 6.96859, 7....

    Dissolving 'by' the cluster column will make the dissolve much
    faster if there are a lot of non-overlapping polygons.

    >>> dissolved = gdf.dissolve(by="cluster", as_index=False)
    >>> dissolved
       cluster                                           geometry
    0        0  POLYGON ((0.99951 -0.03141, 0.99803 -0.06279, ...
    1        1  POLYGON ((4.99951 2.96859, 4.99803 2.93721, 4....
    2        2  POLYGON ((8.00000 7.00000, 7.99951 6.96859, 7....
    """
    if isinstance(gdfs[-1], str):
        *gdfs, cluster_col = gdfs

    concated = pd.DataFrame()
    orig_indices = ()
    for i, gdf in enumerate(gdfs):
        if isinstance(gdf, GeoSeries):
            gdf = gdf.to_frame()

        if not isinstance(gdf, GeoDataFrame):
            raise TypeError("'gdfs' should be one or more GeoDataFrames or GeoSeries.")

        if not allow_multipart and len(gdf) != len(gdf.explode(index_parts=False)):
            raise ValueError(
                "All geometries should be exploded to singlepart "
                "in order to get correct polygon clusters. "
                "To allow multipart geometries, set allow_multipart=True"
            )

        orig_indices = orig_indices + (gdf.index,)
        gdf["i__"] = i

        concated = pd.concat([concated, gdf], ignore_index=True)

    neighbors = get_neighbor_indices(concated, concated)

    edges = [(source, target) for source, target in neighbors.items()]

    graph = nx.Graph()
    graph.add_edges_from(edges)

    component_mapper = {
        j: i
        for i, component in enumerate(nx.connected_components(graph))
        for j in component
    }

    concated[cluster_col] = component_mapper

    concated = _push_geom_col(concated)

    n_gdfs = concated["i__"].unique()

    if len(n_gdfs) == 1:
        concated.index = orig_indices[0]
        return concated.drop(["i__"], axis=1)

    unconcated = ()
    for i in n_gdfs:
        gdf = concated[concated["i__"] == i]
        gdf.index = orig_indices[i]
        gdf = gdf.drop(["i__"], axis=1)
        unconcated = unconcated + (gdf,)

    return unconcated


def get_overlapping_polygons(
    gdf: GeoDataFrame | GeoSeries, ignore_index: bool = False
) -> GeoDataFrame | GeoSeries:
    """Find the areas that overlap.

    Does an intersection with itself and keeps only the duplicated geometries. The
    index of 'gdf' is preserved.

    Args:
        gdf: GeoDataFrame of polygons.
        ignore_index: If True, the resulting axis will be labeled 0, 1, …, n - 1.
            Defaults to False.

    Returns:
        A GeoDataFrame of the overlapping polygons.
    """
    if not gdf.index.is_unique:
        raise ValueError(
            "Index must be unique in order to correctly find "
            "overlapping polygon indices."
        )

    gdf = gdf.assign(overlap=gdf.index)

    intersected = clean_overlay(gdf, gdf[["geometry"]], how="intersection")

    points_joined = intersected.representative_point().to_frame().sjoin(intersected)

    duplicated_points = points_joined.loc[points_joined.index.duplicated()]

    duplicated_geoms = intersected.loc[intersected.index.isin(duplicated_points.index)]
    duplicated_geoms.index = duplicated_geoms["overlap"].values

    if ignore_index:
        duplicated_geoms = duplicated_geoms.reset_index(drop=True)

    return duplicated_geoms.drop("overlap", axis=1)


def get_overlapping_polygon_indices(gdf: GeoDataFrame | GeoSeries) -> pd.Index:
    if not gdf.index.is_unique:
        raise ValueError(
            "Index must be unique in order to correctly find "
            "overlapping polygon indices."
        )

    gdf = gdf.assign(overlap=gdf.index)

    intersected = clean_overlay(gdf, gdf[["geometry"]], how="intersection")

    intersected = intersected.set_index("overlap")

    points_joined = intersected.representative_point().to_frame().sjoin(intersected)

    duplicated_points = points_joined.loc[points_joined.index.duplicated()]

    return duplicated_points.index.unique()


def get_overlapping_polygon_product(gdf: GeoDataFrame | GeoSeries) -> pd.Index:
    if not gdf.index.is_unique:
        raise ValueError("Index must be unique to find overlapping polygon indices.")

    gdf = gdf.assign(overlap=gdf.index)

    intersected = clean_overlay(gdf, gdf[["geometry"]], how="intersection")

    intersected = intersected.set_index("overlap")

    points_joined = intersected.representative_point().to_frame().sjoin(intersected)

    duplicated_points = points_joined.loc[points_joined.index.duplicated()]

    unique = (
        duplicated_points.reset_index()
        .groupby(["overlap", "index_right"])
        .size()
        .reset_index()
    )

    unique = unique[unique.overlap != unique.index_right]
    series = unique.set_index("index_right").overlap
    series.index.name = None

    return series


def close_small_holes(
    gdf: GeoDataFrame | GeoSeries,
    max_area: int | float,
    *,
    copy: bool = True,
) -> GeoDataFrame | GeoSeries:
    """Closes holes in polygons if the area is less than the given maximum.

    It takes a GeoDataFrame or GeoSeries of polygons and
    fills the holes that are smaller than the specified area given in units of
    either square meters ('max_m2') or square kilometers ('max_km2').

    Args:
        gdf: GeoDataFrame or GeoSeries of polygons.
        max_area: The maximum area in the unit of the GeoDataFrame's crs.
        copy: if True (default), the input GeoDataFrame or GeoSeries is copied.
            Defaults to True.

    Returns:
        A GeoDataFrame or GeoSeries of polygons with closed holes in the geometry
        column.

    Raises:
        ValueError: If the coordinate reference system of the GeoDataFrame is not in
            meter units.
        ValueError: If both 'max_m2' and 'max_km2' is given.

    Examples
    --------

    Let's create a circle with a hole in it.

    >>> from sgis import close_small_holes, buff, to_gdf
    >>> point = to_gdf([260000, 6650000], crs=25833)
    >>> point
                            geometry
    0  POINT (260000.000 6650000.000)
    >>> circle = buff(point, 1000)
    >>> small_circle = buff(point, 500)
    >>> circle_with_hole = circle.overlay(small_circle, how="difference")
    >>> circle_with_hole.area
    0    2.355807e+06
    dtype: float64

    Close holes smaller than 1 square kilometer (1 million square meters).

    >>> holes_closed = close_small_holes(circle_with_hole, max_area=1_000_000)
    >>> holes_closed.area
    0    3.141076e+06
    dtype: float64

    The hole will not be closed if it is larger.

    >>> holes_closed = close_small_holes(circle_with_hole, max_area=1_000)
    >>> holes_closed.area
    0    2.355807e+06
    dtype: float64
    """
    if copy:
        gdf = gdf.copy()

    if isinstance(gdf, GeoDataFrame):
        gdf["geometry"] = gdf.geometry.map(
            lambda x: _close_small_holes_poly(x, max_area)
        )
        return gdf

    elif isinstance(gdf, gpd.GeoSeries):
        return gdf.map(lambda x: _close_small_holes_poly(x, max_area))

    else:
        raise ValueError(
            f"'gdf' should be of type GeoDataFrame or GeoSeries. Got {type(gdf)}"
        )


def close_all_holes(
    gdf: GeoDataFrame | GeoSeries,
    *,
    copy: bool = True,
) -> GeoDataFrame | GeoSeries:
    """Closes all holes in polygons.

    It takes a GeoDataFrame or GeoSeries of polygons and
    returns the outer circle.

    Args:
        gdf: GeoDataFrame or GeoSeries of polygons.
        copy: if True (default), the input GeoDataFrame or GeoSeries is copied.
            Defaults to True.

    Returns:
        A GeoDataFrame or GeoSeries of polygons with closed holes in the geometry
        column.

    Examples
    --------
    Let's create a circle with a hole in it.

    >>> from sgis import close_all_holes, buff, to_gdf
    >>> point = to_gdf([260000, 6650000], crs=25833)
    >>> point
                            geometry
    0  POINT (260000.000 6650000.000)
    >>> circle = buff(point, 1000)
    >>> small_circle = buff(point, 500)
    >>> circle_with_hole = circle.overlay(small_circle, how="difference")
    >>> circle_with_hole.area
    0    2.355807e+06
    dtype: float64

    Close the hole.

    >>> holes_closed = close_all_holes(circle_with_hole)
    >>> holes_closed.area
    0    3.141076e+06
    dtype: float64
    """
    if copy:
        gdf = gdf.copy()

    def close_all_holes_func(poly):
        return unary_union(polygons(get_exterior_ring(get_parts(poly))))

    close_all_holes_func = np.vectorize(close_all_holes_func)

    if isinstance(gdf, GeoDataFrame):
        gdf["geometry"] = close_all_holes_func(gdf.geometry)
        return gdf

    elif isinstance(gdf, gpd.GeoSeries):
        return close_all_holes_func(gdf)

    else:
        raise ValueError(
            f"'gdf' should be of type GeoDataFrame or GeoSeries. Got {type(gdf)}"
        )


def _close_small_holes_poly(poly, max_area):
    """Closes cmall holes within one shapely geometry of polygons."""

    # start with a list containing the polygon,
    # then append all holes smaller than 'max_km2' to the list.
    holes_closed = [poly]
    singlepart = get_parts(poly)
    for part in singlepart:
        n_interior_rings = get_num_interior_rings(part)

        if not (n_interior_rings):
            continue

        for n in range(n_interior_rings):
            hole = polygons(get_interior_ring(part, n))

            if area(hole) < max_area:
                holes_closed.append(hole)

    return unary_union(holes_closed)
