#!/usr/bin/env python3

# Merge BYU consolidated iceberg tracking CSVs and create an EPSG:3031 shapefile.
#
# Processing:
#   - Reads all CSVs from cwd/input/consolidated
#   - Adds iceberg_id from each CSV filename
#   - Derives lat/lon using the configured sensor priority
#   - Records the selected sensor in pos_source
#   - Writes the full merged CSV to cwd/output/consolidated
#   - Writes spatial records as an EPSG:3031 shapefile
#
# The shapefile excludes raw sensor coordinate fields (*_1, *_2) and *_3
# fields, while the merged CSV retains all original attributes.
#
# Install requirements:
#   pip install pandas geopandas shapely

from pathlib import Path
import sys

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


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

BASE_DIR = Path.cwd()

INPUT_DIR = BASE_DIR / "input" / "consolidated"
OUTPUT_DIR = BASE_DIR / "output" / "consolidated"

OUTPUT_CSV = OUTPUT_DIR / "iceberg_tracking_database_v7-1_consolidated.csv"
OUTPUT_SHP = OUTPUT_DIR / "iceberg_tracking_database_v7-1_consolidated.shp"

LAT_COL = "lat"
LON_COL = "lon"
POS_SOURCE_COL = "pos_source"

SRC_EPSG = 4326
DST_EPSG = 3031

# Position source priority, highest to lowest.
# qscat and seawinds are treated as the same priority tier.
SENSOR_PRIORITY = [
    ["nic"],
    ["ascat"],
    ["oscat"],
    ["qscat", "seawinds"],
    ["nscat"],
    ["ers"],
    ["sass"],
]


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

def fail(msg: str) -> None:
    print(f"ERROR: {msg}", file=sys.stderr)
    sys.exit(1)


def get_numeric_series(df: pd.DataFrame, col_name: str) -> pd.Series:
    if col_name in df.columns:
        return pd.to_numeric(df[col_name], errors="coerce")

    return pd.Series(float("nan"), index=df.index, dtype="float64")


def derive_lat_lon_and_source(df: pd.DataFrame, csv_name: str) -> pd.DataFrame:
    """
    Derive lat/lon and pos_source using SENSOR_PRIORITY.

    A sensor position is valid only where both *_1 and *_2 fields exist
    and contain non-null, non-zero values.
    """

    df = df.copy()

    df[LAT_COL] = pd.Series(float("nan"), index=df.index, dtype="float64")
    df[LON_COL] = pd.Series(float("nan"), index=df.index, dtype="float64")
    df[POS_SOURCE_COL] = pd.Series([None] * len(df), index=df.index, dtype="object")

    found_any_source_cols = False

    for sensor_group in SENSOR_PRIORITY:
        for sensor in sensor_group:
            lat_col = f"{sensor}_1"
            lon_col = f"{sensor}_2"

            if lat_col not in df.columns or lon_col not in df.columns:
                continue

            found_any_source_cols = True

            sensor_lat = get_numeric_series(df, lat_col)
            sensor_lon = get_numeric_series(df, lon_col)

            valid_sensor = (
                sensor_lat.notna()
                & sensor_lon.notna()
                & (sensor_lat != 0)
                & (sensor_lon != 0)
            )

            use_sensor = df[LAT_COL].isna() & valid_sensor

            df.loc[use_sensor, LAT_COL] = sensor_lat.loc[use_sensor]
            df.loc[use_sensor, LON_COL] = sensor_lon.loc[use_sensor]
            df.loc[use_sensor, POS_SOURCE_COL] = sensor

    if not found_any_source_cols:
        fail(f"{csv_name}: no recognised coordinate source columns found")

    return df


def read_and_tag_csv(csv_path: Path) -> pd.DataFrame:
    df = pd.read_csv(csv_path)
    df.insert(0, "iceberg_id", csv_path.stem)
    return derive_lat_lon_and_source(df, csv_path.name)


def choose_shapefile_columns(df: pd.DataFrame) -> list[str]:
    """
    Retain derived coordinates, position source and general attributes.

    Raw sensor coordinate fields (*_1, *_2) and all *_3 fields are omitted
    from the shapefile but remain in the merged CSV.
    """

    keep_cols = []

    for col in df.columns:
        if col == "geometry":
            continue

        if col in {"iceberg_id", LAT_COL, LON_COL, POS_SOURCE_COL}:
            keep_cols.append(col)
            continue

        lower_col = col.lower()

        if any(
            lower_col == f"{sensor}_1" or lower_col == f"{sensor}_2"
            for group in SENSOR_PRIORITY
            for sensor in group
        ):
            continue

        if lower_col.endswith("_3"):
            continue

        keep_cols.append(col)

    return keep_cols


# =============================================================================
# SPATIAL OUTPUT
# =============================================================================

def build_geometry(df: pd.DataFrame) -> gpd.GeoDataFrame:
    if LAT_COL not in df.columns or LON_COL not in df.columns:
        fail(f"Required columns '{LAT_COL}' and '{LON_COL}' not found")

    work_df = df.copy()

    work_df[LON_COL] = pd.to_numeric(work_df[LON_COL], errors="coerce")
    work_df[LAT_COL] = pd.to_numeric(work_df[LAT_COL], errors="coerce")

    valid = work_df[LON_COL].notna() & work_df[LAT_COL].notna()

    if valid.sum() == 0:
        fail("No valid coordinates found")

    skipped = (~valid).sum()
    if skipped > 0:
        print(
            f"Skipping {skipped} row(s) with missing/invalid "
            f"derived lat/lon in shapefile output"
        )

    spatial_df = work_df.loc[valid].copy()
    spatial_df["geometry"] = [
        Point(xy)
        for xy in zip(spatial_df[LON_COL], spatial_df[LAT_COL])
    ]

    gdf = gpd.GeoDataFrame(
        spatial_df,
        geometry="geometry",
        crs=f"EPSG:{SRC_EPSG}",
    )

    gdf = gdf.to_crs(epsg=DST_EPSG)

    shp_cols = choose_shapefile_columns(gdf)
    return gdf[shp_cols + ["geometry"]]


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

def main() -> None:
    if not INPUT_DIR.exists():
        fail(f"Input directory not found: {INPUT_DIR}")

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

    csv_files = sorted(INPUT_DIR.glob("*.csv"))
    if not csv_files:
        fail("No CSV files found in input directory")

    print(f"Found {len(csv_files)} CSV files")

    frames = []

    for csv_path in csv_files:
        print(f"Reading {csv_path.name}")
        frames.append(read_and_tag_csv(csv_path))

    merged_df = pd.concat(frames, ignore_index=True, sort=False)

    # Keep derived fields together near the front of the merged CSV.
    preferred_front = ["iceberg_id", LAT_COL, LON_COL, POS_SOURCE_COL]
    existing_front = [c for c in preferred_front if c in merged_df.columns]
    other_cols = [c for c in merged_df.columns if c not in existing_front]
    merged_df = merged_df[existing_front + other_cols]

    print("Writing merged CSV")
    merged_df.to_csv(OUTPUT_CSV, index=False)

    print("Creating shapefile")
    gdf = build_geometry(merged_df)

    print("Writing shapefile")
    gdf.to_file(OUTPUT_SHP, driver="ESRI Shapefile")

    print("Done")
    print(f"CSV: {OUTPUT_CSV}")
    print(f"SHP: {OUTPUT_SHP}")


if __name__ == "__main__":
    main()