# Create 1-day elephant seal track segments from dive-behaviour data.
#
# Usage:
#   1. Place ElephantSeal_divebehaviour.csv in ./input/
#   2. Edit CONFIG below if required.
#   3. Run: python 03_trackseg_with_attributes.py
#   Output is written to ./output/
#
# Required packages:
#   pip install pandas geopandas shapely

from pathlib import Path

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


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

INPUT_FILENAME = "ElephantSeal_divebehaviour.csv"
OUTPUT_FILENAME = "ElephantSeal_dive_segments_1d.shp"

# Only connect observations exactly one day apart
REQUIRE_CONSECUTIVE_DAYS = False

INPUT_EPSG = 4326
OUTPUT_EPSG = 3031


# =========================
# PATHS
# =========================

INPUT_PATH = Path.cwd() / "input" / INPUT_FILENAME
OUTPUT_DIR = Path.cwd() / "output"
OUTPUT_PATH = OUTPUT_DIR / OUTPUT_FILENAME

OUTPUT_DIR.mkdir(exist_ok=True)


# =========================
# COLUMN DEFINITIONS
# =========================

META_SRC = [
    "id",
    "Species",
    "Common name",
    "Sex",
    "AgeClass",
    "Location",
    "n_day",
    "day",
]

VARS_1D_SRC = [
    "desc_rate_1d",
    "asc_rate_1d",
    "hunt_1d",
    "bot_time_1d",
    "surf_res_1d",
    "dive_res_1d",
    "shelf_1d",
    "dr_1d",
    "ddr_1d",
    "lon_1d",
    "lat_1d",
]

# Short field names to avoid shapefile truncation
RENAME = {
    "id": "individ_id",
    "Species": "sci_name",
    "Common name": "comm_name",
    "Sex": "sex",
    "AgeClass": "age_class",
    "Location": "location",
    "n_day": "n_day",
    "day": "date",
    "desc_rate_1d": "desc_r_1d",
    "asc_rate_1d": "asc_r_1d",
    "hunt_1d": "hunt_t_1d",
    "bot_time_1d": "bot_t_1d",
    "surf_res_1d": "surfres_1d",
    "dive_res_1d": "diveres_1d",
    "shelf_1d": "shelf_1d",
    "dr_1d": "dr_1d",
    "ddr_1d": "ddr_1d",
    "lon_1d": "lon_1d",
    "lat_1d": "lat_1d",
}

OUT_COLS = [
    "individ_id",
    "sci_name",
    "comm_name",
    "sex",
    "age_class",
    "location",
    "n_day",
    "date",
    "desc_r_1d",
    "asc_r_1d",
    "hunt_t_1d",
    "bot_t_1d",
    "surfres_1d",
    "diveres_1d",
    "shelf_1d",
    "dr_1d",
    "ddr_1d",
    "lon_1d",
    "lat_1d",
]


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

def read_csv_any(path):
    """Read CSV using common encodings."""
    for encoding in ("utf-8", "utf-8-sig", "cp1252", "ISO-8859-1"):
        try:
            return pd.read_csv(path, encoding=encoding)
        except UnicodeDecodeError:
            pass

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


def to_date_series(series):
    return pd.to_datetime(series, errors="coerce", utc=True).dt.date


def unwrap_longitudes(lons):
    """Prevent segments taking the long route across the dateline."""
    unwrapped = []
    offset = 0.0
    previous = None

    for lon in lons:
        if pd.isna(lon):
            unwrapped.append(None)
            previous = None
            offset = 0.0
            continue

        adjusted = float(lon) + offset

        if previous is not None:
            while adjusted - previous > 180:
                offset -= 360
                adjusted = float(lon) + offset

            while adjusted - previous < -180:
                offset += 360
                adjusted = float(lon) + offset

        unwrapped.append(adjusted)
        previous = adjusted

    return unwrapped


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

def main():
    if not INPUT_PATH.exists():
        raise FileNotFoundError(f"Missing input: {INPUT_PATH}")

    df = read_csv_any(INPUT_PATH)

    required = META_SRC + VARS_1D_SRC
    missing = [col for col in required if col not in df.columns]

    if missing:
        raise ValueError(f"Input CSV is missing expected columns: {missing}")

    df = df[required].copy()

    df["day"] = to_date_series(df["day"])
    df["lon_1d"] = pd.to_numeric(df["lon_1d"], errors="coerce")
    df["lat_1d"] = pd.to_numeric(df["lat_1d"], errors="coerce")
    df["dr_1d"] = pd.to_numeric(df["dr_1d"], errors="coerce")
    df["ddr_1d"] = pd.to_numeric(df["ddr_1d"], errors="coerce")

    # Keep shelf classification as text
    df["shelf_1d"] = df["shelf_1d"].astype("string")

    df = df.dropna(
        subset=["id", "day", "lon_1d", "lat_1d"]
    )

    df.rename(columns=RENAME, inplace=True)

    # -------------------------
    # BUILD SEGMENTS
    # -------------------------

    segment_rows = []
    dropped_no_pair = 0
    dropped_gap = 0

    for individ_id, group in df.groupby("individ_id", sort=True):
        group = group.sort_values("date").reset_index(drop=True)

        if len(group) < 2:
            dropped_no_pair += 1
            continue

        unwrapped_lons = unwrap_longitudes(
            group["lon_1d"].tolist()
        )

        for i in range(len(group) - 1):
            start = group.iloc[i]
            end = group.iloc[i + 1]

            if REQUIRE_CONSECUTIVE_DAYS:
                day_gap = (
                    pd.to_datetime(end["date"])
                    - pd.to_datetime(start["date"])
                ).days

                if day_gap != 1:
                    dropped_gap += 1
                    continue

            lon_start = unwrapped_lons[i]
            lon_end = unwrapped_lons[i + 1]

            geometry = LineString([
                (lon_start, start["lat_1d"]),
                (lon_end, end["lat_1d"]),
            ])

            row = {col: start[col] for col in OUT_COLS}
            row["geometry"] = geometry

            segment_rows.append(row)

    if not segment_rows:
        raise RuntimeError(
            "No segments were created; check input data or filters."
        )

    gdf = gpd.GeoDataFrame(
        segment_rows,
        crs=f"EPSG:{INPUT_EPSG}",
    )

    gdf = gdf[
        OUT_COLS + ["geometry"]
    ].to_crs(epsg=OUTPUT_EPSG)

    gdf.to_file(
        OUTPUT_PATH,
        driver="ESRI Shapefile",
    )

    print(
        f"Wrote {len(gdf)} segments -> "
        f"{OUTPUT_PATH} (EPSG:{OUTPUT_EPSG})"
    )

    print(
        f"Individuals in segments: "
        f"{gdf['individ_id'].nunique()}"
    )

    if dropped_no_pair:
        print(
            f"Individuals dropped (only 1 point): "
            f"{dropped_no_pair}"
        )

    if REQUIRE_CONSECUTIVE_DAYS and dropped_gap:
        print(
            f"Segments dropped (non-consecutive days): "
            f"{dropped_gap}"
        )


if __name__ == "__main__":
    main()