#!/usr/bin/env python3

# Convert MEaSUREs Antarctic ice velocity NetCDF speed/direction/texture COGs.
#
# Processing:
#   - Read VX and VY from each NetCDF in cwd/input
#   - Calculate full-resolution ice speed and flow direction
#   - Generate a LIC flow texture at half resolution
#   - Globally normalise texture to 0–1 and resample to the full source grid
#   - Retain the native georeferencing/projection of the velocity dataset
#   - Standardise NoData to -9999
#   - Write separate Float32 COGs to cwd/output
#
# Outputs:
#   <stem>_SPEED.tif
#   <stem>_DIRECTION.tif
#   <stem>_TEXTURE.tif
#
# Install:
#   pip install numpy scipy GDAL

from __future__ import annotations

import math
from pathlib import Path

import numpy as np
from osgeo import gdal
from scipy.ndimage import map_coordinates


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

BASE_DIR = Path.cwd()
INPUT_DIR = BASE_DIR / "input"
OUTPUT_DIR = BASE_DIR / "output"

VAR_VX = "VX"
VAR_VY = "VY"

NODATA = -9999.0

COG_OPTIONS = [
    "COMPRESS=DEFLATE",
    "BIGTIFF=IF_SAFER",
]

TEMP_TIFF_OPTIONS = [
    "TILED=YES",
    "COMPRESS=DEFLATE",
    "BIGTIFF=IF_SAFER",
]

# LIC texture settings.
LIC_DOWNSAMPLE = 2
LIC_LENGTH = 12
LIC_STEP_SIZE = 1.0
LIC_TILE_SIZE = 1024
LIC_MARGIN = LIC_LENGTH * 2 + 8
LIC_SEED = 42

# Global texture normalisation percentiles.
LIC_NORM_PCT_LOW = 2.0
LIC_NORM_PCT_HIGH = 98.0


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

def nc_subdataset(path: Path, variable: str) -> str:
    return f'NETCDF:"{path}":{variable}'


def open_ds(path: str) -> gdal.Dataset:
    ds = gdal.Open(path, gdal.GA_ReadOnly)

    if ds is None:
        raise RuntimeError(f"Failed to open: {path}")

    return ds


def get_nodata(band: gdal.Band) -> float | None:
    return band.GetNoDataValue()


def create_temp_gtiff(
    path: Path,
    xsize: int,
    ysize: int,
    dtype: int,
    geotransform,
    projection: str,
) -> gdal.Dataset:

    ds = gdal.GetDriverByName("GTiff").Create(
        str(path),
        xsize,
        ysize,
        1,
        dtype,
        options=TEMP_TIFF_OPTIONS,
    )

    if ds is None:
        raise RuntimeError(f"Failed to create: {path}")

    ds.SetGeoTransform(geotransform)
    ds.SetProjection(projection)
    ds.GetRasterBand(1).SetNoDataValue(NODATA)

    return ds


def translate_to_cog(src: Path, dst: Path) -> None:
    if dst.exists():
        dst.unlink()

    output = gdal.Translate(
        str(dst),
        str(src),
        options=gdal.TranslateOptions(
            format="COG",
            creationOptions=COG_OPTIONS,
        ),
    )

    if output is None:
        raise RuntimeError(f"Failed to create COG: {dst}")

    output = None


def resample_to_match_grid(
    src: Path,
    dst: Path,
    xsize: int,
    ysize: int,
    geotransform,
    projection: str,
) -> None:

    minx = geotransform[0]
    maxy = geotransform[3]
    maxx = minx + xsize * geotransform[1] + ysize * geotransform[2]
    miny = maxy + xsize * geotransform[4] + ysize * geotransform[5]

    output = gdal.Warp(
        str(dst),
        str(src),
        options=gdal.WarpOptions(
            format="GTiff",
            outputBounds=(minx, miny, maxx, maxy),
            width=xsize,
            height=ysize,
            dstSRS=projection,
            resampleAlg="bilinear",
            srcNodata=NODATA,
            dstNodata=NODATA,
            creationOptions=TEMP_TIFF_OPTIONS,
        ),
    )

    if output is None:
        raise RuntimeError(f"Failed to resample: {dst}")

    output = None


def mask_from_vxvy(
    vx: np.ndarray,
    vy: np.ndarray,
    vx_nodata: float | None,
    vy_nodata: float | None,
) -> np.ndarray:

    mask = ~np.isfinite(vx) | ~np.isfinite(vy)

    if vx_nodata is not None:
        mask |= vx == vx_nodata

    if vy_nodata is not None:
        mask |= vy == vy_nodata

    return mask


# =============================================================================
# SPEED / DIRECTION
# =============================================================================

def compute_speed_direction(
    vx: np.ndarray,
    vy: np.ndarray,
    mask: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:

    vx = vx.astype(np.float32, copy=False)
    vy = vy.astype(np.float32, copy=False)

    speed = np.sqrt(vx * vx + vy * vy).astype(np.float32)

    direction = (
        np.degrees(np.arctan2(vx, vy)) + 360.0
    ) % 360.0
    direction = direction.astype(np.float32)

    speed[mask] = NODATA
    direction[mask] = NODATA

    return speed, direction


def write_speed_direction(
    vx_band: gdal.Band,
    vy_band: gdal.Band,
    vx_nodata: float | None,
    vy_nodata: float | None,
    speed_path: Path,
    direction_path: Path,
    xsize: int,
    ysize: int,
    geotransform,
    projection: str,
) -> None:

    speed_ds = create_temp_gtiff(
        speed_path,
        xsize,
        ysize,
        gdal.GDT_Float32,
        geotransform,
        projection,
    )

    direction_ds = create_temp_gtiff(
        direction_path,
        xsize,
        ysize,
        gdal.GDT_Float32,
        geotransform,
        projection,
    )

    speed_band = speed_ds.GetRasterBand(1)
    direction_band = direction_ds.GetRasterBand(1)

    speed_band.SetDescription("SPEED")
    speed_band.SetUnitType("m/yr")

    direction_band.SetDescription("DIRECTION")
    direction_band.SetUnitType("degrees")

    block_x, block_y = vx_band.GetBlockSize()

    if not block_x or not block_y:
        block_x, block_y = 512, 512

    for yoff in range(0, ysize, block_y):
        ywin = min(block_y, ysize - yoff)

        for xoff in range(0, xsize, block_x):
            xwin = min(block_x, xsize - xoff)

            vx = vx_band.ReadAsArray(xoff, yoff, xwin, ywin)
            vy = vy_band.ReadAsArray(xoff, yoff, xwin, ywin)

            if vx is None or vy is None:
                raise RuntimeError(
                    f"Read failed at x={xoff}, y={yoff}"
                )

            mask = mask_from_vxvy(
                vx,
                vy,
                vx_nodata,
                vy_nodata,
            )

            speed, direction = compute_speed_direction(
                vx,
                vy,
                mask,
            )

            speed_band.WriteArray(speed, xoff, yoff)
            direction_band.WriteArray(direction, xoff, yoff)

        print(
            f"speed/direction: "
            f"{min(100, int((yoff + ywin) * 100 / ysize))}%"
        )

    speed_ds.FlushCache()
    direction_ds.FlushCache()

    speed_ds = None
    direction_ds = None


# =============================================================================
# LIC TEXTURE
# =============================================================================

def deterministic_noise(
    global_x0: int,
    global_y0: int,
    width: int,
    height: int,
) -> np.ndarray:

    xs = np.arange(
        global_x0,
        global_x0 + width,
        dtype=np.float32,
    )[None, :]

    ys = np.arange(
        global_y0,
        global_y0 + height,
        dtype=np.float32,
    )[:, None]

    noise = (
        np.sin(
            xs * 12.9898
            + ys * 78.233
            + LIC_SEED * 0.12345
        )
        * 43758.5453
    )

    return (noise - np.floor(noise)).astype(np.float32)


def texture_raw(
    vx: np.ndarray,
    vy: np.ndarray,
    mask: np.ndarray,
    global_x0: int,
    global_y0: int,
) -> np.ndarray:

    height, width = vx.shape

    magnitude = np.sqrt(vx * vx + vy * vy).astype(np.float32)
    valid = (~mask) & np.isfinite(magnitude) & (magnitude > 0)

    u = np.zeros_like(vx, dtype=np.float32)
    v = np.zeros_like(vy, dtype=np.float32)

    u[valid] = vx[valid] / magnitude[valid]
    v[valid] = -vy[valid] / magnitude[valid]

    valid_float = valid.astype(np.float32)
    noise = deterministic_noise(
        global_x0,
        global_y0,
        width,
        height,
    )

    yy, xx = np.mgrid[0:height, 0:width].astype(np.float32)

    accumulator = np.zeros((height, width), dtype=np.float32)
    weights = np.zeros((height, width), dtype=np.float32)

    accumulator[valid] = noise[valid]
    weights[valid] = 1.0

    for sign in (-1.0, 1.0):
        x = xx.copy()
        y = yy.copy()

        for _ in range(LIC_LENGTH):
            uu = map_coordinates(
                u,
                [y, x],
                order=1,
                mode="nearest",
            )

            vv = map_coordinates(
                v,
                [y, x],
                order=1,
                mode="nearest",
            )

            valid_sample = (
                map_coordinates(
                    valid_float,
                    [y, x],
                    order=0,
                    mode="nearest",
                )
                > 0.5
            )

            x += sign * LIC_STEP_SIZE * uu
            y += sign * LIC_STEP_SIZE * vv

            inside = (
                (x >= 0)
                & (x <= width - 1)
                & (y >= 0)
                & (y <= height - 1)
                & valid_sample
            )

            if not np.any(inside):
                continue

            sample = map_coordinates(
                noise,
                [
                    np.clip(y, 0, height - 1),
                    np.clip(x, 0, width - 1),
                ],
                order=1,
                mode="nearest",
            )

            accumulator[inside] += sample[inside]
            weights[inside] += 1.0

    output = np.full(
        (height, width),
        NODATA,
        dtype=np.float32,
    )

    valid_output = valid & (weights > 0)
    output[valid_output] = (
        accumulator[valid_output]
        / weights[valid_output]
    )

    return output


def write_halfres_texture(
    vx_band: gdal.Band,
    vy_band: gdal.Band,
    vx_nodata: float | None,
    vy_nodata: float | None,
    output_path: Path,
    full_xsize: int,
    full_ysize: int,
    full_geotransform,
    projection: str,
) -> None:

    half_xsize = math.ceil(full_xsize / LIC_DOWNSAMPLE)
    half_ysize = math.ceil(full_ysize / LIC_DOWNSAMPLE)

    half_geotransform = (
        full_geotransform[0],
        full_geotransform[1] * LIC_DOWNSAMPLE,
        full_geotransform[2],
        full_geotransform[3],
        full_geotransform[4],
        full_geotransform[5] * LIC_DOWNSAMPLE,
    )

    output_ds = create_temp_gtiff(
        output_path,
        half_xsize,
        half_ysize,
        gdal.GDT_Float32,
        half_geotransform,
        projection,
    )

    output_band = output_ds.GetRasterBand(1)
    output_band.SetDescription("TEXTURE")
    output_band.SetUnitType("unitless")

    for hy in range(0, half_ysize, LIC_TILE_SIZE):
        core_height = min(
            LIC_TILE_SIZE,
            half_ysize - hy,
        )

        for hx in range(0, half_xsize, LIC_TILE_SIZE):
            core_width = min(
                LIC_TILE_SIZE,
                half_xsize - hx,
            )

            ex0 = max(0, hx - LIC_MARGIN)
            ey0 = max(0, hy - LIC_MARGIN)

            ex1 = min(
                half_xsize,
                hx + core_width + LIC_MARGIN,
            )

            ey1 = min(
                half_ysize,
                hy + core_height + LIC_MARGIN,
            )

            extended_width = ex1 - ex0
            extended_height = ey1 - ey0

            fx0 = ex0 * LIC_DOWNSAMPLE
            fy0 = ey0 * LIC_DOWNSAMPLE

            fx1 = min(
                full_xsize,
                ex1 * LIC_DOWNSAMPLE,
            )

            fy1 = min(
                full_ysize,
                ey1 * LIC_DOWNSAMPLE,
            )

            full_width = fx1 - fx0
            full_height = fy1 - fy0

            vx = vx_band.ReadAsArray(
                fx0,
                fy0,
                full_width,
                full_height,
                buf_xsize=extended_width,
                buf_ysize=extended_height,
            )

            vy = vy_band.ReadAsArray(
                fx0,
                fy0,
                full_width,
                full_height,
                buf_xsize=extended_width,
                buf_ysize=extended_height,
            )

            if vx is None or vy is None:
                raise RuntimeError(
                    f"Texture read failed at hx={hx}, hy={hy}"
                )

            vx = vx.astype(np.float32)
            vy = vy.astype(np.float32)

            mask = mask_from_vxvy(
                vx,
                vy,
                vx_nodata,
                vy_nodata,
            )

            texture = texture_raw(
                vx,
                vy,
                mask,
                ex0,
                ey0,
            )

            cy0 = hy - ey0
            cx0 = hx - ex0

            core = texture[
                cy0:cy0 + core_height,
                cx0:cx0 + core_width,
            ].astype(np.float32)

            core_mask = mask[
                cy0:cy0 + core_height,
                cx0:cx0 + core_width,
            ]

            core[core_mask] = NODATA

            output_band.WriteArray(
                core,
                hx,
                hy,
            )

        print(
            f"texture raw half-res: "
            f"{min(100, int((hy + core_height) * 100 / half_ysize))}%"
        )

    output_ds.FlushCache()
    output_ds = None


def normalize_texture(
    src_path: Path,
    dst_path: Path,
) -> None:

    src_ds = gdal.Open(
        str(src_path),
        gdal.GA_ReadOnly,
    )

    if src_ds is None:
        raise RuntimeError(
            f"Failed to open texture: {src_path}"
        )

    band = src_ds.GetRasterBand(1)

    xsize = src_ds.RasterXSize
    ysize = src_ds.RasterYSize
    geotransform = src_ds.GetGeoTransform()
    projection = src_ds.GetProjection()

    block_x, block_y = band.GetBlockSize()

    if not block_x or not block_y:
        block_x, block_y = 512, 512

    samples = []
    sample_stride = 8

    for yoff in range(0, ysize, block_y):
        ywin = min(block_y, ysize - yoff)

        for xoff in range(0, xsize, block_x):
            xwin = min(block_x, xsize - xoff)

            array = band.ReadAsArray(
                xoff,
                yoff,
                xwin,
                ywin,
            )

            if array is None:
                raise RuntimeError(
                    f"Failed reading texture at x={xoff}, y={yoff}"
                )

            valid = array[
                (array != NODATA)
                & np.isfinite(array)
            ]

            if valid.size:
                samples.append(valid[::sample_stride])

    if not samples:
        raise RuntimeError(
            "No valid texture pixels found for normalisation"
        )

    sample = np.concatenate(samples)

    low = float(
        np.percentile(
            sample,
            LIC_NORM_PCT_LOW,
        )
    )

    high = float(
        np.percentile(
            sample,
            LIC_NORM_PCT_HIGH,
        )
    )

    if (
        not np.isfinite(low)
        or not np.isfinite(high)
        or high <= low
    ):
        low = float(sample.min())
        high = float(sample.max())

        if high <= low:
            high = low + 1e-6

    print(
        f"Texture normalisation: "
        f"low={low:.6f}, high={high:.6f}"
    )

    output_ds = create_temp_gtiff(
        dst_path,
        xsize,
        ysize,
        gdal.GDT_Float32,
        geotransform,
        projection,
    )

    output_band = output_ds.GetRasterBand(1)
    output_band.SetDescription("TEXTURE")
    output_band.SetUnitType("unitless")

    for yoff in range(0, ysize, block_y):
        ywin = min(block_y, ysize - yoff)

        for xoff in range(0, xsize, block_x):
            xwin = min(block_x, xsize - xoff)

            array = band.ReadAsArray(
                xoff,
                yoff,
                xwin,
                ywin,
            ).astype(np.float32)

            valid = (
                (array != NODATA)
                & np.isfinite(array)
            )

            output = np.full(
                array.shape,
                NODATA,
                dtype=np.float32,
            )

            if np.any(valid):
                scaled = (
                    array[valid] - low
                ) / (
                    high - low
                )

                output[valid] = np.clip(
                    scaled,
                    0.0,
                    1.0,
                )

            output_band.WriteArray(
                output,
                xoff,
                yoff,
            )

        print(
            f"texture normalize: "
            f"{min(100, int((yoff + ywin) * 100 / ysize))}%"
        )

    output_ds.FlushCache()

    output_ds = None
    src_ds = None


# =============================================================================
# PROCESS ONE NETCDF
# =============================================================================

def process_nc(nc_path: Path) -> None:
    stem = nc_path.stem

    speed_cog = OUTPUT_DIR / f"{stem}_SPEED.tif"
    direction_cog = OUTPUT_DIR / f"{stem}_DIRECTION.tif"
    texture_cog = OUTPUT_DIR / f"{stem}_TEXTURE.tif"

    tmp_speed = OUTPUT_DIR / f"_{stem}_tmp_speed.tif"
    tmp_direction = OUTPUT_DIR / f"_{stem}_tmp_direction.tif"
    tmp_texture_raw = OUTPUT_DIR / f"_{stem}_tmp_texture_half_raw.tif"
    tmp_texture_norm = OUTPUT_DIR / f"_{stem}_tmp_texture_half_norm.tif"
    tmp_texture_full = OUTPUT_DIR / f"_{stem}_tmp_texture_full.tif"

    vx_ds = open_ds(
        nc_subdataset(nc_path, VAR_VX)
    )

    vy_ds = open_ds(
        nc_subdataset(nc_path, VAR_VY)
    )

    vx_band = vx_ds.GetRasterBand(1)
    vy_band = vy_ds.GetRasterBand(1)

    xsize = vx_ds.RasterXSize
    ysize = vx_ds.RasterYSize

    geotransform = vx_ds.GetGeoTransform()
    projection = vx_ds.GetProjection()

    vx_nodata = get_nodata(vx_band)
    vy_nodata = get_nodata(vy_band)

    print(f"\n[INPUT] {nc_path.relative_to(BASE_DIR)}")
    print(f"[GRID] {xsize} x {ysize}")

    # Speed and direction.
    write_speed_direction(
        vx_band,
        vy_band,
        vx_nodata,
        vy_nodata,
        tmp_speed,
        tmp_direction,
        xsize,
        ysize,
        geotransform,
        projection,
    )

    translate_to_cog(
        tmp_speed,
        speed_cog,
    )

    translate_to_cog(
        tmp_direction,
        direction_cog,
    )

    # Half-resolution LIC texture.
    write_halfres_texture(
        vx_band,
        vy_band,
        vx_nodata,
        vy_nodata,
        tmp_texture_raw,
        xsize,
        ysize,
        geotransform,
        projection,
    )

    # Global 0–1 normalisation.
    normalize_texture(
        tmp_texture_raw,
        tmp_texture_norm,
    )

    # Resample texture back to the source grid.
    resample_to_match_grid(
        tmp_texture_norm,
        tmp_texture_full,
        xsize,
        ysize,
        geotransform,
        projection,
    )

    translate_to_cog(
        tmp_texture_full,
        texture_cog,
    )

    vx_ds = None
    vy_ds = None

    for temp_path in (
        tmp_speed,
        tmp_direction,
        tmp_texture_raw,
        tmp_texture_norm,
        tmp_texture_full,
    ):
        temp_path.unlink(missing_ok=True)

    print(f"[OK] {speed_cog.relative_to(BASE_DIR)}")
    print(f"[OK] {direction_cog.relative_to(BASE_DIR)}")
    print(f"[OK] {texture_cog.relative_to(BASE_DIR)}")


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

def main() -> None:
    gdal.UseExceptions()
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    nc_files = sorted(
        INPUT_DIR.glob("*.nc")
    )

    if not nc_files:
        raise FileNotFoundError(
            f"No NetCDF files found in: {INPUT_DIR}"
        )

    for nc_path in nc_files:
        try:
            process_nc(nc_path)

        except Exception as error:
            print(
                f"[FAIL] {nc_path.name}: {error}"
            )

    print()
    print("FINISHED")


if __name__ == "__main__":
    main()