#!/usr/bin/env python3

# Create one Antarctic tracking line per RAATD individual and join metadata.
#
# Input:
#   ./input/RAATD_metadata.csv
#   ./input/*.csv              all other CSVs are treated as tracking data
#
# Output:
#   ./output/SCAR_RAATD_filtered_tracks.shp
#
# Install:
#   pip install pandas geopandas shapely

from pathlib import Path

import pandas as pd
import geopandas as gpd
from shapely.geometry import LineString


# =========================
# CONFIG
# =========================

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

META_FILE = INPUT_DIR / "RAATD_metadata.csv"
OUTPUT_SHP = OUTPUT_DIR / "SCAR_RAATD_filtered_tracks.shp"


# =========================
# HELPERS
# =========================

def read_tracking_csv(path: Path) -> pd.DataFrame:
    """Read and standardise one tracking CSV."""

    df = pd.read_csv(path)

    df = df[
        [
            "individual_id",
            "abbreviated_name",
            "date",
            "decimal_latitude",
            "decimal_longitude",
        ]
    ].copy()

    df.rename(
        columns={
            "individual_id": "individ_id",
            "abbreviated_name": "abbrv_name",
            "decimal_latitude": "latitude",
            "decimal_longitude": "longitude",
        },
        inplace=True,
    )

    df["date"] = pd.to_datetime(df["date"], errors="coerce")
    df["latitude"] = pd.to_numeric(df["latitude"], errors="coerce")
    df["longitude"] = pd.to_numeric(df["longitude"], errors="coerce")

    df = df.dropna(
        subset=["individ_id", "date", "latitude", "longitude"]
    )

    return df


def read_metadata_csv(path: Path) -> pd.DataFrame:
    """Read metadata using fallback encodings if required."""

    for encoding in ("utf-8", "utf-8-sig", "cp1252", "ISO-8859-1"):
        try:
            return pd.read_csv(path, encoding=encoding)
        except UnicodeDecodeError:
            continue

    return pd.read_csv(
        path,
        encoding="latin-1",
        engine="python",
    )


def to_line(group: pd.DataFrame) -> pd.Series:
    """Convert one individual's ordered positions to a LineString."""

    ordered = group.sort_values("date")

    coords = list(
        zip(
            ordered["longitude"],
            ordered["latitude"],
        )
    )

    geometry = LineString(coords) if len(coords) >= 2 else None

    individ_id, abbrv_name = group.name

    return pd.Series(
        {
            "individ_id": individ_id,
            "abbrv_name": abbrv_name,
            "date_start": ordered["date"].min().date(),
            "date_end": ordered["date"].max().date(),
            "geometry": geometry,
        }
    )


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

def main():

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

    if not META_FILE.is_file():
        raise SystemExit(f"Metadata file not found: {META_FILE}")

    # Find all tracking CSVs except the metadata file
    tracking_files = sorted(
        path
        for path in INPUT_DIR.glob("*.csv")
        if path.name != META_FILE.name
    )

    if not tracking_files:
        raise SystemExit(f"No tracking CSV files found in: {INPUT_DIR}")

    print(f"Found {len(tracking_files)} tracking CSV file(s).")

    # Read and combine tracking data
    frames = [
        read_tracking_csv(path)
        for path in tracking_files
    ]

    tracks = pd.concat(
        frames,
        ignore_index=True,
    )

    total_ids = tracks["individ_id"].nunique()
    pts_per_id = tracks.groupby("individ_id").size()

    n_ge2 = (pts_per_id >= 2).sum()
    n_eq1 = (pts_per_id == 1).sum()

    print(f"Tracked individuals: {total_ids}")
    print(f"  With >=2 points (lines): {n_ge2}")
    print(f"  With <2 points (dropped): {n_eq1}")

    # Create one LineString per unique individual/species abbreviation
    agg = (
        tracks
        .groupby(["individ_id", "abbrv_name"])
        .apply(to_line, include_groups=False)
        .reset_index(drop=True)
    )

    agg = (
        agg
        .dropna(subset=["geometry"])
        .reset_index(drop=True)
    )

    # Read and prepare metadata
    meta = read_metadata_csv(META_FILE)

    if "individual_id" in meta.columns:
        meta = meta.rename(
            columns={"individual_id": "individ_id"}
        )

    meta = meta[
        [
            "individ_id",
            "dataset_identifier",
            "scientific_name",
            "common_name",
            "sex",
            "age_class",
        ]
    ].rename(
        columns={
            "dataset_identifier": "dataset_id",
            "scientific_name": "sci_name",
            "common_name": "comm_name",
        }
    )

    # Join metadata
    merged = agg.merge(
        meta,
        on="individ_id",
        how="left",
    )

    # Report individuals without matching metadata
    meta_cols = [
        "dataset_id",
        "sci_name",
        "comm_name",
        "sex",
        "age_class",
    ]

    unmatched = (
        merged[
            merged[meta_cols]
            .isna()
            .all(axis=1)
        ]["individ_id"]
        .dropna()
        .unique()
        .tolist()
    )

    if unmatched:
        print(
            "No metadata match for:",
            ", ".join(map(str, sorted(unmatched))),
        )

    # Create spatial data in WGS84
    gdf_4326 = gpd.GeoDataFrame(
        merged,
        geometry="geometry",
        crs="EPSG:4326",
    )[
        [
            "individ_id",
            "date_start",
            "date_end",
            "dataset_id",
            "abbrv_name",
            "sci_name",
            "comm_name",
            "sex",
            "age_class",
            "geometry",
        ]
    ]

    # Reproject to Antarctic Polar Stereographic
    gdf_3031 = gdf_4326.to_crs("EPSG:3031")

    # Write shapefile
    gdf_3031.to_file(
        OUTPUT_SHP,
        driver="ESRI Shapefile",
    )

    print(
        f"Wrote {len(gdf_3031)} tracks to "
        f"{OUTPUT_SHP} (EPSG:3031)"
    )


if __name__ == "__main__":
    main()