#!/usr/bin/env python3

# Combine shapefiles into a GeoPackage, grouped by geometry type.
#
# Searches recursively beneath:
#     cwd/input/
#
# Creates:
#     cwd/output/Antarctica_human_footprint.gpkg
#
# GeoPackage layers:
#     point_layers
#     line_layers
#     polygon_layers
#
# Processing:
#     - Keep plausible Antarctic point geometries unchanged
#     - Rebuild bad or implausible EPSG:4326 point geometries from
#       Lat_DD and Long_DD / Lon_DD where possible
#     - Drop bad point geometries that cannot be rebuilt
#     - Repair invalid polygon geometries using make_valid()
#     - Retain only polygonal parts from repaired polygon geometries
#     - Drop features with missing or unusable geometry
#     - Reproject all features to EPSG:3031
#     - Drop fields listed in DROP_FIELDS
#     - Add a source field containing the original shapefile name
#     - Merge differing schemas using the union of fields
#
# Install requirements:
#     pip install geopandas pyogrio pandas shapely pyproj


from pathlib import Path
import math

import geopandas as gpd
import pandas as pd

from pyproj import CRS
from shapely.geometry import Point, MultiPolygon


try:
    from shapely import make_valid
except ImportError:
    from shapely.validation import make_valid


# =============================================================================
# CONFIGURATION
# =============================================================================

INPUT_DIR = Path.cwd() / "input"
OUTPUT_DIR = Path.cwd() / "output"

OUTPUT_FILE = OUTPUT_DIR / "Antarctica_human_footprint.gpkg"

TARGET_CRS = "EPSG:3031"

# Northern limit for plausible Antarctic point coordinates.
# Points north of 60°S are treated as bad for this dataset.
ANTARCTIC_NORTH_LIMIT = -50.0

DROP_FIELDS = {
    "shape_leng",
    "shape_area",
    "lat_dd",
    "long_dd",
    "lon_dd",
    "point",
    "area",
    "id",
    "length",
}

LAYER_NAMES = {
    "point": "point_layers",
    "line": "line_layers",
    "polygon": "polygon_layers",
}


# =============================================================================
# FUNCTIONS
# =============================================================================

def geometry_family(gdf):
    """Return point, line or polygon based on the geometries present."""

    geom_types = set(
        gdf.geometry.geom_type.dropna().unique()
    )

    if not geom_types:
        return None

    if geom_types <= {"Point", "MultiPoint"}:
        return "point"

    if geom_types <= {"LineString", "MultiLineString", "LinearRing"}:
        return "line"

    if geom_types <= {"Polygon", "MultiPolygon"}:
        return "polygon"

    raise ValueError(
        f"Unsupported or mixed geometry types: {sorted(geom_types)}"
    )


def drop_unwanted_fields(gdf):
    """Drop configured fields case-insensitively."""

    drop_fields_lower = {
        field.lower()
        for field in DROP_FIELDS
    }

    fields_to_drop = [
        column
        for column in gdf.columns
        if column.lower() in drop_fields_lower
    ]

    return gdf.drop(columns=fields_to_drop)


# =============================================================================
# POINT REPAIR
# =============================================================================

def point_geometry_is_bad(geom):
    """
    Return True where a geographic point geometry is missing, unusable
    or implausible for this Antarctic dataset.

    Catches:
        - NULL / empty geometry
        - non-finite coordinates
        - extreme sentinel coordinates such as +/-1.79769313486232E+308
        - longitude outside -180 to 180
        - latitude outside -90 to ANTARCTIC_NORTH_LIMIT
        - invalid Z sentinel values
    """

    if geom is None or geom.is_empty:
        return True

    try:
        bounds = geom.bounds
    except Exception:
        return True

    if not all(
        math.isfinite(value)
        for value in bounds
    ):
        return True

    if any(
        abs(value) > 1e100
        for value in bounds
    ):
        return True

    minx, miny, maxx, maxy = bounds

    # Antarctic plausibility check.
    if (
        minx < -180
        or maxx > 180
        or miny < -90
        or maxy > ANTARCTIC_NORTH_LIMIT
    ):
        return True

    if geom.geom_type == "Point" and geom.has_z:

        try:
            z = geom.z

            if (
                not math.isfinite(z)
                or abs(z) > 1e100
            ):
                return True

        except Exception:
            return True

    return False


def repair_invalid_point_geometry(gdf, relative_path):
    """
    Repair bad EPSG:4326 point geometries from Lat_DD and
    Long_DD / Lon_DD where possible.

    Good, plausible Antarctic geometries are left unchanged.

    Bad geometries that cannot be reconstructed from plausible Antarctic
    latitude/longitude attributes are set to NULL and subsequently dropped.
    """

    if gdf.crs is None:
        return gdf

    source_crs = CRS.from_user_input(
        gdf.crs
    )

    # Lat_DD / Long_DD are geographic coordinates, so only use them
    # to rebuild geometry where the source layer is EPSG:4326.
    if not source_crs.equals(
        CRS.from_epsg(4326)
    ):
        return gdf

    columns = {
        column.lower(): column
        for column in gdf.columns
    }

    lat_field = columns.get(
        "lat_dd"
    )

    lon_field = (
        columns.get("long_dd")
        or columns.get("lon_dd")
    )

    repaired_indices = []
    unrecoverable_indices = []

    for idx, geom in gdf.geometry.items():

        # Good geometry: leave exactly as supplied.
        if not point_geometry_is_bad(geom):
            continue

        # Bad geometry: try to rebuild from lat/lon attributes.
        if (
            lat_field is not None
            and lon_field is not None
        ):

            lat = gdf.at[
                idx,
                lat_field,
            ]

            lon = gdf.at[
                idx,
                lon_field,
            ]

            if (
                not pd.isna(lat)
                and not pd.isna(lon)
            ):

                try:
                    lat = float(lat)
                    lon = float(lon)

                except (TypeError, ValueError):
                    lat = None
                    lon = None

                if (
                    lat is not None
                    and lon is not None
                    and math.isfinite(lat)
                    and math.isfinite(lon)
                    and -90 <= lat <= ANTARCTIC_NORTH_LIMIT
                    and -180 <= lon <= 180
                ):

                    # Rebuild as a clean 2D point.
                    gdf.at[
                        idx,
                        gdf.geometry.name,
                    ] = Point(
                        lon,
                        lat,
                    )

                    repaired_indices.append(
                        idx
                    )

                    continue

        # Could not recover this bad geometry.
        gdf.at[
            idx,
            gdf.geometry.name,
        ] = None

        unrecoverable_indices.append(
            idx
        )

    if repaired_indices:

        print(
            f"  Rebuilt bad point geometries "
            f"from lat/lon: "
            f"{len(repaired_indices):,}"
        )

        for idx in repaired_indices[:20]:
            print(
                f"    Feature {idx}"
            )

        if len(repaired_indices) > 20:

            print(
                f"    ... plus "
                f"{len(repaired_indices) - 20:,} more"
            )

    if unrecoverable_indices:

        print(
            f"  Bad point geometries with no usable "
            f"Antarctic lat/lon: "
            f"{len(unrecoverable_indices):,}"
        )

    return gdf


# =============================================================================
# POLYGON REPAIR
# =============================================================================

def extract_polygon_parts(geom):
    """
    Recursively extract polygon components from a geometry.

    Lines and points produced by make_valid() are discarded.
    """

    if geom is None or geom.is_empty:
        return []

    if geom.geom_type == "Polygon":

        if geom.area > 0:
            return [geom]

        return []

    if geom.geom_type == "MultiPolygon":

        return [
            part
            for part in geom.geoms
            if (
                not part.is_empty
                and part.area > 0
            )
        ]

    if hasattr(geom, "geoms"):

        parts = []

        for part in geom.geoms:

            parts.extend(
                extract_polygon_parts(
                    part
                )
            )

        return parts

    return []


def repair_invalid_polygon_geometry(gdf, relative_path):
    """
    Repair invalid polygon geometries using make_valid().

    Only polygonal components are retained. Line and point fragments
    produced during repair are discarded.
    """

    repaired_indices = []
    unrecoverable_indices = []

    for idx, geom in gdf.geometry.items():

        if geom is None or geom.is_empty:
            continue

        # Leave valid polygons untouched.
        if geom.is_valid:
            continue

        try:
            repaired_geom = make_valid(
                geom
            )

        except Exception:

            gdf.at[
                idx,
                gdf.geometry.name,
            ] = None

            unrecoverable_indices.append(
                idx
            )

            continue

        polygon_parts = extract_polygon_parts(
            repaired_geom
        )

        if not polygon_parts:

            gdf.at[
                idx,
                gdf.geometry.name,
            ] = None

            unrecoverable_indices.append(
                idx
            )

            continue

        if len(polygon_parts) == 1:

            new_geom = polygon_parts[0]

        else:

            new_geom = MultiPolygon(
                polygon_parts
            )

        gdf.at[
            idx,
            gdf.geometry.name,
        ] = new_geom

        repaired_indices.append(
            idx
        )

    if repaired_indices:

        print(
            f"  Repaired invalid polygon geometries: "
            f"{len(repaired_indices):,}"
        )

        for idx in repaired_indices[:20]:
            print(
                f"    Feature {idx}"
            )

        if len(repaired_indices) > 20:

            print(
                f"    ... plus "
                f"{len(repaired_indices) - 20:,} more"
            )

    if unrecoverable_indices:

        print(
            f"  Polygon geometries with no "
            f"recoverable polygon parts: "
            f"{len(unrecoverable_indices):,}"
        )

    return gdf


# =============================================================================
# GENERAL GEOMETRY CHECKS
# =============================================================================

def drop_unusable_geometries(gdf, relative_path):
    """
    Drop features with missing or unusable geometry.

    Includes:
        - NULL geometry
        - empty geometry
        - non-finite coordinates
        - extreme sentinel coordinates such as +/-1.79769313486232E+308
    """

    bad_indices = []

    for idx, geom in gdf.geometry.items():

        if geom is None or geom.is_empty:

            bad_indices.append(
                idx
            )

            continue

        try:
            bounds = geom.bounds

        except Exception:

            bad_indices.append(
                idx
            )

            continue

        if not all(
            math.isfinite(value)
            for value in bounds
        ):

            bad_indices.append(
                idx
            )

            continue

        if any(
            abs(value) > 1e100
            for value in bounds
        ):

            bad_indices.append(
                idx
            )

    if bad_indices:

        print(
            f"  Dropping features with no usable geometry: "
            f"{len(bad_indices):,}"
        )

        for idx in bad_indices[:20]:
            print(
                f"    Feature {idx}"
            )

        if len(bad_indices) > 20:

            print(
                f"    ... plus "
                f"{len(bad_indices) - 20:,} more"
            )

        gdf = gdf.drop(
            index=bad_indices
        ).copy()

    return gdf


def check_finite_geometries(gdf, relative_path):
    """
    Stop if any geometry contains non-finite coordinates after
    reprojection.
    """

    bad_features = []

    for idx, geom in gdf.geometry.items():

        if geom is None or geom.is_empty:

            bad_features.append(
                idx
            )

            continue

        try:
            bounds = geom.bounds

        except Exception:

            bad_features.append(
                idx
            )

            continue

        if not all(
            math.isfinite(value)
            for value in bounds
        ):

            bad_features.append(
                idx
            )

    if bad_features:

        preview = ", ".join(
            str(idx)
            for idx in bad_features[:10]
        )

        raise ValueError(
            f"{relative_path}: "
            f"{len(bad_features):,} geometry feature(s) "
            f"contain non-finite coordinates after "
            f"reprojection to {TARGET_CRS}. "
            f"Feature indices: {preview}"
        )


# =============================================================================
# ATTRIBUTE SCHEMA
# =============================================================================

def build_canonical_field_names(datasets):
    """
    Build a case-insensitive field-name lookup.

    The capitalisation/spelling from the first occurrence is retained.
    """

    canonical = {
        "source": "source"
    }

    for dataset in datasets:

        gdf = dataset["gdf"]

        for column in gdf.columns:

            if column == gdf.geometry.name:
                continue

            lower = column.lower()

            if lower not in canonical:
                canonical[lower] = column

    return canonical


def harmonise_fields(gdf, canonical):
    """
    Rename fields case-insensitively to their canonical names.

    If a source contains multiple fields differing only by case,
    their values are coalesced into a single field.
    """

    geometry_column = gdf.geometry.name

    result = gpd.GeoDataFrame(
        geometry=gdf.geometry.copy(),
        crs=gdf.crs,
    )

    grouped_columns = {}

    for column in gdf.columns:

        if column == geometry_column:
            continue

        canonical_name = canonical[
            column.lower()
        ]

        grouped_columns.setdefault(
            canonical_name,
            []
        ).append(column)

    for canonical_name, columns in grouped_columns.items():

        if len(columns) == 1:

            result[canonical_name] = gdf[
                columns[0]
            ]

        else:

            print(
                f"    Combining case-duplicate fields: "
                f"{', '.join(columns)} -> "
                f"{canonical_name}"
            )

            result[canonical_name] = (
                gdf[columns]
                .bfill(axis=1)
                .iloc[:, 0]
            )

    return result


# =============================================================================
# MAIN
# =============================================================================

def main():

    OUTPUT_DIR.mkdir(
        parents=True,
        exist_ok=True,
    )

    shapefiles = sorted(
        INPUT_DIR.rglob("*.shp")
    )

    if not shapefiles:

        raise FileNotFoundError(
            f"No shapefiles found beneath:\n"
            f"{INPUT_DIR}"
        )

    print(
        f"Found {len(shapefiles)} shapefile(s)"
    )

    print(
        f"Output CRS: {TARGET_CRS}"
    )

    print(
        f"Antarctic point limit: "
        f"{ANTARCTIC_NORTH_LIMIT}° latitude\n"
    )

    groups = {
        "point": [],
        "line": [],
        "polygon": [],
    }

    target_crs = CRS.from_user_input(
        TARGET_CRS
    )

    # -------------------------------------------------------------------------
    # Read, repair, clean and reproject
    # -------------------------------------------------------------------------

    for shp in shapefiles:

        relative_path = shp.relative_to(
            INPUT_DIR
        )

        print(
            f"Reading: {relative_path}"
        )

        gdf = gpd.read_file(
            shp
        )

        if gdf.empty:

            print(
                "  Skipped: no features\n"
            )

            continue

        family = geometry_family(
            gdf
        )

        if family is None:

            print(
                "  Skipped: no geometries\n"
            )

            continue

        # ---------------------------------------------------------------------
        # Point repair
        # ---------------------------------------------------------------------

        # Valid/plausible Antarctic point geometry is left unchanged.
        #
        # Bad or implausible point geometry is rebuilt from
        # Lat_DD + Long_DD / Lon_DD where possible.
        #
        # If those replacement coordinates are also implausible,
        # the feature is set to NULL and dropped below.
        if family == "point":

            gdf = repair_invalid_point_geometry(
                gdf,
                relative_path,
            )

        # ---------------------------------------------------------------------
        # Drop unusable geometry
        # ---------------------------------------------------------------------

        gdf = drop_unusable_geometries(
            gdf,
            relative_path,
        )

        if gdf.empty:

            print(
                "  Skipped: no usable geometries remain\n"
            )

            continue

        # ---------------------------------------------------------------------
        # Polygon repair
        # ---------------------------------------------------------------------

        if family == "polygon":

            gdf = repair_invalid_polygon_geometry(
                gdf,
                relative_path,
            )

            # make_valid() may leave some features with no recoverable
            # polygon component.
            gdf = drop_unusable_geometries(
                gdf,
                relative_path,
            )

            if gdf.empty:

                print(
                    "  Skipped: no usable polygon "
                    "geometries remain\n"
                )

                continue

        # ---------------------------------------------------------------------
        # CRS
        # ---------------------------------------------------------------------

        if gdf.crs is None:

            raise ValueError(
                f"{relative_path} has no defined CRS. "
                f"Cannot safely reproject to "
                f"{TARGET_CRS}."
            )

        source_crs = CRS.from_user_input(
            gdf.crs
        )

        if not source_crs.equals(
            target_crs
        ):

            print(
                f"  Reprojecting: "
                f"{source_crs.to_string()} -> "
                f"{TARGET_CRS}"
            )

            gdf = gdf.to_crs(
                TARGET_CRS
            )

        check_finite_geometries(
            gdf,
            relative_path,
        )

        # ---------------------------------------------------------------------
        # Attributes
        # ---------------------------------------------------------------------

        # Lat_DD / Long_DD / Lon_DD are deliberately retained until here,
        # so they are available for point geometry repair.
        gdf = drop_unwanted_fields(
            gdf
        )

        existing_source_fields = [
            column
            for column in gdf.columns
            if column.lower() == "source"
        ]

        if existing_source_fields:

            raise ValueError(
                f"{relative_path} already contains "
                f"a field named "
                f"{existing_source_fields[0]!r}, "
                f"which conflicts with the new "
                f"'source' field."
            )

        gdf.insert(
            0,
            "source",
            shp.stem,
        )

        groups[family].append(
            {
                "path": relative_path,
                "gdf": gdf,
            }
        )

        print(
            f"  Geometry: {family}"
        )

        print(
            f"  Features: {len(gdf):,}\n"
        )

    # -------------------------------------------------------------------------
    # Remove existing output
    # -------------------------------------------------------------------------

    if OUTPUT_FILE.exists():

        OUTPUT_FILE.unlink()

    # -------------------------------------------------------------------------
    # Merge and write each geometry group
    # -------------------------------------------------------------------------

    first_layer = True

    for family, datasets in groups.items():

        if not datasets:

            print(
                f"No {family} shapefiles found - "
                f"skipping layer."
            )

            continue

        layer_name = LAYER_NAMES[
            family
        ]

        print(
            f"\nCreating: {layer_name}"
        )

        # ---------------------------------------------------------------------
        # Harmonise field names
        # ---------------------------------------------------------------------

        canonical = build_canonical_field_names(
            datasets
        )

        prepared = []

        for dataset in datasets:

            gdf = harmonise_fields(
                dataset["gdf"],
                canonical,
            )

            prepared.append(
                gdf
            )

        # ---------------------------------------------------------------------
        # Merge schemas
        # ---------------------------------------------------------------------

        merged = gpd.GeoDataFrame(
            pd.concat(
                prepared,
                ignore_index=True,
                sort=False,
            ),
            geometry="geometry",
            crs=TARGET_CRS,
        )

        # source first, geometry last
        columns = (
            ["source"]
            + [
                column
                for column in merged.columns
                if column not in {
                    "source",
                    "geometry",
                }
            ]
            + ["geometry"]
        )

        merged = merged[
            columns
        ]

        # ---------------------------------------------------------------------
        # Write GeoPackage
        # ---------------------------------------------------------------------

        mode = (
            "w"
            if first_layer
            else "a"
        )

        merged.to_file(
            OUTPUT_FILE,
            layer=layer_name,
            driver="GPKG",
            engine="pyogrio",
            mode=mode,
        )

        first_layer = False

        print(
            f"  Features: {len(merged):,}"
        )

        print(
            f"  Fields:   "
            f"{len(merged.columns) - 1:,}"
        )

        print(
            f"  CRS:      {TARGET_CRS}"
        )

    print(
        "\nComplete"
    )

    print(
        f"Output: "
        f"{OUTPUT_FILE.relative_to(Path.cwd())}"
    )


if __name__ == "__main__":
    main()