# Grid seal dives to the deepest observed dive within each cell.
#
# Required packages:
#   pip install numpy pandas pyproj rasterio

import math
from pathlib import Path

import numpy as np
import pandas as pd
import rasterio as rio
from pyproj import Transformer
from rasterio.shutil import copy as rio_copy
from rasterio.transform import from_origin


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

INPUT_FILENAME = "dives.csv"
OUTPUT_FILENAME = "seal_dives_maxdepth_25.tif"

GRID_SIZE = 2500.0
TARGET_EPSG = 3031

LON_COL = "lon"
LAT_COL = "lat"
VALUE_COL = "max_depth"      # Negative depth: more negative = deeper

DEEPER_DIVES_ONLY = True
FILTER_COL = "dfb_adj"       # Retain only values > 0 when filtering

NODATA = -32768.0
DTYPE = "float32"
BLOCKSIZE = 512


# =========================
# 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)


# =========================
# LOAD AND FILTER
# =========================

df = pd.read_csv(INPUT_PATH)

required = [LON_COL, LAT_COL, VALUE_COL]
if DEEPER_DIVES_ONLY:
    required.append(FILTER_COL)

missing = [col for col in required if col not in df.columns]
if missing:
    raise ValueError(f"Missing required columns: {', '.join(missing)}")

df = df[required].copy()

for col in required:
    df[col] = pd.to_numeric(df[col], errors="coerce")

df = df.dropna()

if DEEPER_DIVES_ONLY:
    df = df[df[FILTER_COL] > 0]

if df.empty:
    raise ValueError("No valid dives remain after filtering.")


# =========================
# PROJECT POINTS
# =========================

transformer = Transformer.from_crs(
    4326,
    TARGET_EPSG,
    always_xy=True,
)

x, y = transformer.transform(
    df[LON_COL].to_numpy(),
    df[LAT_COL].to_numpy(),
)

values = df[VALUE_COL].to_numpy(dtype=DTYPE)


# =========================
# CREATE GRID
# =========================

# Snap raster extent to exact GRID_SIZE multiples
xmin = math.floor(np.min(x) / GRID_SIZE) * GRID_SIZE
xmax = math.ceil(np.max(x) / GRID_SIZE) * GRID_SIZE
ymin = math.floor(np.min(y) / GRID_SIZE) * GRID_SIZE
ymax = math.ceil(np.max(y) / GRID_SIZE) * GRID_SIZE

width = int(round((xmax - xmin) / GRID_SIZE))
height = int(round((ymax - ymin) / GRID_SIZE))

if width <= 0 or height <= 0:
    raise ValueError("Computed raster has invalid dimensions.")

# Convert projected coordinates to raster cells
cols = ((x - xmin) // GRID_SIZE).astype(np.int64)
rows = ((ymax - y) // GRID_SIZE).astype(np.int64)

# Remove any points falling outside due to boundary/rounding effects
valid = (
    (cols >= 0)
    & (cols < width)
    & (rows >= 0)
    & (rows < height)
)

rows = rows[valid]
cols = cols[valid]
values = values[valid]

if values.size == 0:
    raise ValueError("No points fall within the calculated raster grid.")

transform = from_origin(
    xmin,
    ymax,
    GRID_SIZE,
    GRID_SIZE,
)


# =========================
# DEEPEST DIVE PER CELL
# =========================

flat_index = rows * width + cols

# max_depth is negative, so minimum = deepest
raster = np.full(width * height, np.inf, dtype=DTYPE)
np.minimum.at(raster, flat_index, values)

raster[np.isposinf(raster)] = NODATA
raster = raster.reshape(height, width)


# =========================
# WRITE COG
# =========================

TEMP_PATH = OUTPUT_DIR / f"{OUTPUT_PATH.stem}_TEMP.tif"

profile = {
    "driver": "GTiff",
    "height": height,
    "width": width,
    "count": 1,
    "dtype": DTYPE,
    "crs": f"EPSG:{TARGET_EPSG}",
    "transform": transform,
    "nodata": NODATA,
    "tiled": True,
    "blockxsize": BLOCKSIZE,
    "blockysize": BLOCKSIZE,
    "compress": "DEFLATE",
    "predictor": 2,
    "BIGTIFF": "IF_SAFER",
    "NUM_THREADS": "ALL_CPUS",
}

with rio.open(TEMP_PATH, "w", **profile) as dst:
    dst.write(raster, 1)

    # Build nearest-neighbour overviews
    factors = []
    level = 2

    while min(dst.width // level, dst.height // level) >= 1:
        factors.append(level)
        level *= 2

    if factors:
        dst.build_overviews(factors, rio.enums.Resampling.nearest)
        dst.update_tags(ns="rio_overview", resampling="nearest")

rio_copy(
    TEMP_PATH,
    OUTPUT_PATH,
    driver="COG",
    BLOCKSIZE=BLOCKSIZE,
    COMPRESS="DEFLATE",
    PREDICTOR=2,
    BIGTIFF="IF_SAFER",
    NUM_THREADS="ALL_CPUS",
    OVERVIEW_RESAMPLING="NEAREST",
)

TEMP_PATH.unlink()

print(f"Created: {OUTPUT_PATH}")