# Convert elephant seal dive-behaviour CSV data to track and point shapefiles.
#
# Usage:
#   1. Place ElephantSeal_divebehaviour.csv in ./input/
#   2. Edit CONFIG below if required.
#   3. Run: python 02_process_csv.py
#   Outputs are 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, MultiLineString


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

INPUT_FILENAME = "ElephantSeal_divebehaviour.csv"

LINES_OUTPUT = "ElephantSeal_dive_tracks.shp"
POINTS_PREFIX = "ElephantSeal_dive_behaviour" # point outputs: f"{POINTS_PREFIX}_{Nd}_pts.shp"

# Smoothing windows to export as point files
POINT_WINDOWS = [1, 7]

# True = skip requested windows if their columns are missing
SKIP_MISSING_WINDOWS = False

INPUT_EPSG = 4326
OUTPUT_EPSG = 3031


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

INPUT_PATH = Path.cwd() / "input" / INPUT_FILENAME
OUTPUT_DIR = Path.cwd() / "output"
LINES_PATH = OUTPUT_DIR / LINES_OUTPUT

OUTPUT_DIR.mkdir(exist_ok=True)


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

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

META_RENAME = {
    "id": "individ_id",
    "Species": "sci_name",
    "Common name": "comm_name",
    "Sex": "sex",
    "AgeClass": "age_class",
    "Location": "location",
    "n_day": "n_day",
    "day": "date",
}

METRIC_SRC_PATTERNS = [
    "desc_rate_{d}d",
    "asc_rate_{d}d",
    "hunt_{d}d",
    "bot_time_{d}d",
    "surf_res_{d}d",
    "dive_res_{d}d",
    "lon_{d}d",
    "lat_{d}d",
    "shelf_{d}d",
    "dr_{d}d",
    "ddr_{d}d",
]

# Short field names to avoid shapefile truncation
METRIC_RENAME_PATTERNS = {
    "desc_rate_{d}d": "desc_r_{d}d",
    "asc_rate_{d}d": "asc_r_{d}d",
    "hunt_{d}d": "hunt_t_{d}d",
    "bot_time_{d}d": "bot_t_{d}d",
    "surf_res_{d}d": "surfres_{d}d",
    "dive_res_{d}d": "diveres_{d}d",
    "lon_{d}d": "lon_{d}d",
    "lat_{d}d": "lat_{d}d",
    "shelf_{d}d": "shelf_{d}d",
    "dr_{d}d": "dr_{d}d",
    "ddr_{d}d": "ddr_{d}d",
}

POINT_META_ORDER = [
    "individ_id",
    "sci_name",
    "comm_name",
    "sex",
    "age_class",
    "location",
    "n_day",
    "date",
]

LINE_COL_ORDER = [
    "individ_id",
    "sci_name",
    "comm_name",
    "sex",
    "age_class",
    "location",
    "date_start",
    "date_end",
]


# =========================
# 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 split_on_dateline(lons, lats):
    """Split tracks where longitude crosses the dateline."""
    coords = list(zip(lons, lats))

    if len(coords) < 2:
        return []

    segments = []
    current = [coords[0]]

    for (lon_prev, lat_prev), (lon, lat) in zip(coords[:-1], coords[1:]):
        if pd.isna(lon_prev) or pd.isna(lat_prev) or pd.isna(lon) or pd.isna(lat):
            if len(current) >= 2:
                segments.append(LineString(current))
            current = []
            continue

        if abs(lon - lon_prev) > 180:
            if len(current) >= 2:
                segments.append(LineString(current))
            current = [(lon, lat)]
        else:
            current.append((lon, lat))

    if len(current) >= 2:
        segments.append(LineString(current))

    return segments


def build_line_geom(group, lon_col, lat_col):
    group = group.sort_values("date")

    segments = split_on_dateline(
        group[lon_col].tolist(),
        group[lat_col].tolist(),
    )

    if not segments:
        return None

    return segments[0] if len(segments) == 1 else MultiLineString(segments)


def first_nonnull(series):
    values = series.dropna()
    return values.iloc[0] if len(values) else None


def window_required_cols(days):
    return META_SRC + [
        pattern.format(d=days)
        for pattern in METRIC_SRC_PATTERNS
    ]


def window_rename_map(days):
    rename = dict(META_RENAME)

    for source, output in METRIC_RENAME_PATTERNS.items():
        rename[source.format(d=days)] = output.format(d=days)

    return rename


def point_output_order(days):
    return POINT_META_ORDER + [
        f"desc_r_{days}d",
        f"asc_r_{days}d",
        f"hunt_t_{days}d",
        f"bot_t_{days}d",
        f"surfres_{days}d",
        f"diveres_{days}d",
        f"shelf_{days}d",
        f"dr_{days}d",
        f"ddr_{days}d",
        f"lon_{days}d",
        f"lat_{days}d",
    ]


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

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

    df_all = read_csv_any(INPUT_PATH)
    print(f"Read {len(df_all)} rows from {INPUT_PATH.name}")

    # -------------------------
    # TRACK LINES FROM 1-DAY DATA
    # -------------------------

    required = META_SRC + ["lon_1d", "lat_1d"]
    missing = [col for col in required if col not in df_all.columns]

    if missing:
        raise ValueError(
            f"Input CSV is missing required columns for lines: {missing}"
        )

    lines_df = df_all[required].copy()

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

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

    lines_df.rename(columns=META_RENAME, inplace=True)

    line_rows = []

    for individ_id, group in lines_df.groupby("individ_id", sort=True):
        geom = build_line_geom(
            group,
            lon_col="lon_1d",
            lat_col="lat_1d",
        )

        if geom is None:
            continue

        line_rows.append({
            "individ_id": individ_id,
            "sci_name": first_nonnull(group["sci_name"]),
            "comm_name": first_nonnull(group["comm_name"]),
            "sex": first_nonnull(group["sex"]),
            "age_class": first_nonnull(group["age_class"]),
            "location": first_nonnull(group["location"]),
            "date_start": group["date"].min(),
            "date_end": group["date"].max(),
            "geometry": geom,
        })

    gdf_lines = gpd.GeoDataFrame(
        line_rows,
        crs=f"EPSG:{INPUT_EPSG}",
    )

    gdf_lines = gdf_lines[
        LINE_COL_ORDER + ["geometry"]
    ].to_crs(epsg=OUTPUT_EPSG)

    gdf_lines.to_file(
        LINES_PATH,
        driver="ESRI Shapefile",
    )

    print(
        f"Wrote {len(gdf_lines)} tracks -> "
        f"{LINES_PATH} (EPSG:{OUTPUT_EPSG})"
    )

    # -------------------------
    # POINT FILES
    # -------------------------

    for days in POINT_WINDOWS:
        required = window_required_cols(days)
        missing = [
            col for col in required
            if col not in df_all.columns
        ]

        if missing:
            message = f"Window {days}d missing columns: {missing}"

            if SKIP_MISSING_WINDOWS:
                print(f"Skipping {days}d points: {message}")
                continue

            raise ValueError(message)

        wdf = df_all[required].copy()

        wdf["day"] = to_date_series(wdf["day"])

        lon_col = f"lon_{days}d"
        lat_col = f"lat_{days}d"

        wdf[lon_col] = pd.to_numeric(wdf[lon_col], errors="coerce")
        wdf[lat_col] = pd.to_numeric(wdf[lat_col], errors="coerce")
        wdf[f"dr_{days}d"] = pd.to_numeric(wdf[f"dr_{days}d"], errors="coerce")
        wdf[f"ddr_{days}d"] = pd.to_numeric(wdf[f"ddr_{days}d"], errors="coerce")

        # Keep shelf values as text
        wdf[f"shelf_{days}d"] = wdf[f"shelf_{days}d"].astype("string")

        wdf = wdf.dropna(
            subset=["id", "day", lon_col, lat_col]
        )

        wdf.rename(
            columns=window_rename_map(days),
            inplace=True,
        )

        wdf = wdf.sort_values(
            ["individ_id", "date"]
        ).reset_index(drop=True)

        output_cols = point_output_order(days)
        points_df = wdf[output_cols].copy()

        lon_col = f"lon_{days}d"
        lat_col = f"lat_{days}d"

        gdf_points = gpd.GeoDataFrame(
            points_df,
            geometry=gpd.points_from_xy(
                points_df[lon_col],
                points_df[lat_col],
            ),
            crs=f"EPSG:{INPUT_EPSG}",
        ).to_crs(epsg=OUTPUT_EPSG)

        output_path = (
            OUTPUT_DIR
            / f"{POINTS_PREFIX}_{days}d_pts.shp"
        )

        gdf_points.to_file(
            output_path,
            driver="ESRI Shapefile",
        )

        print(
            f"Wrote {len(gdf_points)} points ({days}d) -> "
            f"{output_path} (EPSG:{OUTPUT_EPSG})"
        )

    print("Done.")


if __name__ == "__main__":
    main()