When to Use Batch vs Streaming for Real-Time AIS Tracking

"""
ais_stream_handler.py — Lambda handler for real-time AIS processing.

Environment variables required:
  GDAL_DATA       = /opt/share/gdal
  PROJ_LIB        = /opt/share/proj
  H3_GEOFENCE_BUCKET = S3 bucket holding geofence_cells.json
  AIS_STREAM_NAME = Kinesis stream name (informational, used in logs)

Trigger: Kinesis Data Stream with batch size 100, bisect-on-error enabled.
"""

import json
import logging
import os
import base64
from datetime import datetime, timezone, timedelta
from typing import Optional

import boto3
import h3
from pyais import decode as ais_decode
from pyais.exceptions import InvalidNMEAMessageException

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# ---------------------------------------------------------------------------
# Geofence index — loaded once per container lifetime (warm-start caching).
# The index is a dict: {"restricted_cells": [...], "monitoring_cells": [...]}
# stored as JSON in S3 and precomputed at H3 resolution 7 (~5 km² per cell).
# ---------------------------------------------------------------------------
_GEOFENCE_INDEX: Optional[dict] = None

def _load_geofence_index() -> dict:
    global _GEOFENCE_INDEX
    if _GEOFENCE_INDEX is not None:
        return _GEOFENCE_INDEX
    s3 = boto3.client("s3")
    bucket = os.environ["H3_GEOFENCE_BUCKET"]
    obj = s3.get_object(Bucket=bucket, Key="geofence_cells.json")
    raw = json.loads(obj["Body"].read())
    # Convert lists to sets for O(1) containment checks.
    _GEOFENCE_INDEX = {
        "restricted_cells": set(raw.get("restricted_cells", [])),
        "monitoring_cells": set(raw.get("monitoring_cells", [])),
    }
    logger.info(
        "Loaded geofence index: %d restricted cells, %d monitoring cells",
        len(_GEOFENCE_INDEX["restricted_cells"]),
        len(_GEOFENCE_INDEX["monitoring_cells"]),
    )
    return _GEOFENCE_INDEX


# ---------------------------------------------------------------------------
# Watermark: reject messages older than MAX_EVENT_AGE_MINUTES to prevent
# stale satellite retransmissions from triggering live geofence alerts.
# 20 minutes absorbs the upper bound of satellite downlink latency (~15 min)
# while still catching most late-arriving terrestrial messages (~2–3 min).
# ---------------------------------------------------------------------------
MAX_EVENT_AGE_MINUTES = 20

# MMSI ranges to drop outright (base stations, aids to navigation, test MMSIs)
_SKIP_MMSI_PREFIXES = (
    "0",    # base stations (00xxxxxxx)
    "99",   # aids to navigation (99xxxxxxx)
    "111",  # SAR aircraft
)


def _decode_record(record: dict) -> Optional[dict]:
    """
    Decode one Kinesis record into a structured AIS position dict.
    Returns None when the record should be silently dropped.
    Raises on hard decode errors so Kinesis bisect-on-error can isolate them.
    """
    # Kinesis wraps payload in base64.
    raw_bytes = base64.b64decode(record["kinesis"]["data"])

    try:
        msg = ais_decode(raw_bytes)
    except InvalidNMEAMessageException as exc:
        # Malformed sentences — log and drop; do not raise (would block the shard).
        logger.warning("Invalid NMEA sentence (dropped): %s — %s", raw_bytes[:80], exc)
        return None

    mmsi = str(getattr(msg, "mmsi", "") or "")
    if any(mmsi.startswith(p) for p in _SKIP_MMSI_PREFIXES):
        return None  # Infrastructure MMSI — not a vessel

    lat = getattr(msg, "lat", None)
    lon = getattr(msg, "lon", None)
    if lat is None or lon is None:
        return None  # Message type does not carry position (e.g., type 5 static data)
    if not (-90.0 <= lat <= 90.0 and -180.0 <= lon <= 180.0):
        logger.debug("MMSI %s: out-of-range coordinates (%.4f, %.4f) — dropped", mmsi, lat, lon)
        return None

    # Event time comes from Kinesis approximate arrival, not the NMEA sentence
    # (AIS NMEA carries only a UTC second-of-minute, not a full timestamp).
    event_ts_str = record["kinesis"]["approximateArrivalTimestamp"]
    # approximateArrivalTimestamp is a Unix epoch float from the Kinesis SDK.
    event_ts = datetime.fromtimestamp(float(event_ts_str), tz=timezone.utc)

    age = datetime.now(tz=timezone.utc) - event_ts
    if age > timedelta(minutes=MAX_EVENT_AGE_MINUTES):
        # Late satellite message — route to batch reconciliation, not live alerts.
        logger.info(
            "MMSI %s: message age %s exceeds watermark (%d min) — skipping live path",
            mmsi, age, MAX_EVENT_AGE_MINUTES,
        )
        return None

    sog = getattr(msg, "speed", None)   # Speed over ground in knots
    cog = getattr(msg, "course", None)  # Course over ground in degrees

    # H3 resolution 7: ~5 km² cells, fast enough for port-density volumes.
    h3_cell = h3.latlng_to_cell(lat, lon, 7)

    return {
        "mmsi": mmsi,
        "event_ts": event_ts.isoformat(),
        "lat": lat,
        "lon": lon,
        "sog_knots": sog,
        "cog_deg": cog,
        "h3_r7": h3_cell,
        "msg_type": msg.msg_type,
    }


def _check_geofences(position: dict, index: dict) -> list[str]:
    """Return a list of alert types triggered for this position."""
    alerts = []
    cell = position["h3_r7"]
    if cell in index["restricted_cells"]:
        alerts.append("RESTRICTED_ZONE_ENTRY")
    if cell in index["monitoring_cells"]:
        alerts.append("MONITORING_ZONE_ENTRY")
    # Abnormal speed: SOG > 30 knots is unusual in restricted/coastal zones.
    sog = position.get("sog_knots")
    if sog is not None and sog > 30.0 and cell in index["monitoring_cells"]:
        alerts.append("ABNORMAL_SPEED_IN_ZONE")
    return alerts


def lambda_handler(event: dict, context) -> dict:
    """
    Kinesis trigger entry point.
    Configure the trigger with:
      - BisectBatchOnFunctionError: true   (isolate bad records)
      - StartingPosition: LATEST           (live tracking)
      - BatchSize: 100
      - ParallelizationFactor: 4 per shard (for port-density bursts)
    """
    index = _load_geofence_index()
    records = event.get("Records", [])
    processed = dropped = alerted = 0

    for record in records:
        position = _decode_record(record)
        if position is None:
            dropped += 1
            continue

        alerts = _check_geofences(position, index)
        if alerts:
            # In production: publish to SNS topic or EventBridge for downstream handlers.
            logger.warning(
                "ALERT mmsi=%s ts=%s h3=%s alerts=%s",
                position["mmsi"], position["event_ts"], position["h3_r7"], alerts,
            )
            alerted += 1

        processed += 1

    logger.info(
        "Batch complete: processed=%d dropped=%d alerted=%d total=%d",
        processed, dropped, alerted, len(records),
    )
    return {"processed": processed, "dropped": dropped, "alerted": alerted}

"""
ais_batch_reconcile.py — PySpark job for nightly AIS trajectory reconciliation.
Run on EMR Serverless (AWS), Dataproc Serverless (GCP), or Azure Synapse Spark pool.

Environment variables:
  RAW_AIS_PATH    = s3://ais-raw/year={}/month={}/day={}
  ANALYTICS_PATH  = s3://ais-analytics/trajectories/
  PROCESSING_DATE = YYYY-MM-DD (injected by orchestrator)
"""

import os
from datetime import datetime, timedelta
from pyspark.sql import SparkSession, functions as F, Window

spark = SparkSession.builder.appName("ais-reconcile").getOrCreate()

processing_date = os.environ["PROCESSING_DATE"]
dt = datetime.strptime(processing_date, "%Y-%m-%d")
# Include prior day to catch late satellite downlinks (up to 15-min delay
# means end-of-day messages may not appear until early next day UTC).
lookback_start = (dt - timedelta(days=1)).strftime("%Y-%m-%d")

raw_path = os.environ["RAW_AIS_PATH"]
df = spark.read.parquet(
    raw_path.format(dt.year, f"{dt.month:02d}", f"{dt.day:02d}"),
    raw_path.format(
        (dt - timedelta(days=1)).year,
        f"{(dt - timedelta(days=1)).month:02d}",
        f"{(dt - timedelta(days=1)).day:02d}",
    ),
)

# Normalise timestamps and filter to processing window.
df = (
    df.withColumn("event_ts", F.to_utc_timestamp(F.col("event_ts"), "UTC"))
      .filter(F.col("event_ts") >= F.lit(lookback_start))
      .filter(F.col("event_ts") < F.lit((dt + timedelta(days=1)).strftime("%Y-%m-%d")))
)

# Deduplicate: for each (MMSI, minute-bucket), keep the record with the
# highest-quality source (terrestrial preferred over satellite) then earliest ts.
dedup_window = Window.partitionBy(
    "mmsi", F.date_trunc("minute", "event_ts")
).orderBy(
    F.col("source_type").asc(),  # "terrestrial" < "satellite" alphabetically
    F.col("event_ts").asc(),
)
df_deduped = (
    df.withColumn("row_num", F.row_number().over(dedup_window))
      .filter(F.col("row_num") == 1)
      .drop("row_num")
)

# Aggregate daily trajectory per vessel.
trajectory_df = (
    df_deduped
    .groupBy("mmsi", F.date_trunc("day", "event_ts").alias("day"))
    .agg(
        F.count("*").alias("message_count"),
        F.avg("sog_knots").alias("avg_sog_knots"),
        F.max("sog_knots").alias("max_sog_knots"),
        F.collect_list(
            F.struct("event_ts", "lat", "lon", "sog_knots", "cog_deg", "h3_r7")
        ).alias("trajectory"),
    )
)

# Write with Z-ordering on MMSI for efficient per-vessel queries.
(
    trajectory_df.write
    .mode("overwrite")
    .partitionBy("day")
    .parquet(os.environ["ANALYTICS_PATH"])
)

print(f"Reconciliation complete for {processing_date}: {df_deduped.count():,} deduped records")

"""
verify_ais_handler.py — Smoke test for the streaming Lambda handler.
Run locally with AWS credentials configured.
"""

import base64
import json
from datetime import datetime, timezone
from unittest.mock import patch, MagicMock

# Minimal valid AIS Class A position report (MMSI 123456789, 51.5°N 0.1°W)
SAMPLE_NMEA = b"!AIVDM,1,1,,A,13HOI:0P0000W9N0000000000000,0*1B"

def make_kinesis_record(nmea_bytes: bytes, age_seconds: float = 2.0) -> dict:
    ts = (datetime.now(tz=timezone.utc).timestamp() - age_seconds)
    return {
        "kinesis": {
            "data": base64.b64encode(nmea_bytes).decode(),
            "approximateArrivalTimestamp": str(ts),
        }
    }

# Patch S3 geofence load so the test runs without live AWS access.
mock_s3 = MagicMock()
mock_s3.get_object.return_value = {
    "Body": MagicMock(read=lambda: json.dumps({
        "restricted_cells": [],
        "monitoring_cells": [],
    }).encode())
}

with patch("boto3.client", return_value=mock_s3):
    import ais_stream_handler as handler
    handler._GEOFENCE_INDEX = None  # reset warm-state cache

    event = {"Records": [make_kinesis_record(SAMPLE_NMEA, age_seconds=2.0)]}
    result = handler.lambda_handler(event, None)

assert result["processed"] == 1, f"Expected 1 processed, got {result}"
assert result["dropped"] == 0, f"Expected 0 dropped, got {result}"
print("PASS:", result)
# Expected output: PASS: {'processed': 1, 'dropped': 0, 'alerted': 0}

# Verify watermark rejects old satellite messages (>20 min).
event_stale = {"Records": [make_kinesis_record(SAMPLE_NMEA, age_seconds=1300.0)]}
with patch("boto3.client", return_value=mock_s3):
    handler._GEOFENCE_INDEX = None
    result_stale = handler.lambda_handler(event_stale, None)

assert result_stale["processed"] == 0, "Stale satellite message should be dropped"
print("PASS watermark:", result_stale)
# Expected output: PASS watermark: {'processed': 0, 'dropped': 1, 'alerted': 0}

# On EMR or local PySpark shell:
python -c "
from pyspark.sql import SparkSession
spark = SparkSession.builder.getOrCreate()
df = spark.read.parquet('s3://ais-analytics/trajectories/day=2024-10-01/')
print('Vessel trajectories:', df.count())
df.select('mmsi','message_count','avg_sog_knots').orderBy('message_count', ascending=False).show(10)
"