#!/usr/bin/env python3

# Apply adaptive contrast stretching to processed NSIDC ice shelf imagery.
#
# Processing:
#   - Find TIFFs recursively in cwd/imagery/<ice_shelf>/int16_3031/
#   - Apply a standard 2–98% stretch
#   - Use a stronger stretch for low-dynamic-range images
#   - Optionally brighten flat/dark images while avoiding already-bright scenes
#   - Convert to Byte with NoData 0
#   - Write COGs to cwd/imagery/<ice_shelf>/colcorrect/
#
# Existing non-empty outputs are skipped.
#
# Install:
#   pip install GDAL

from __future__ import annotations

from pathlib import Path

from osgeo import gdal


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

BASE_DIR = Path.cwd()
IMAGERY_DIR = BASE_DIR / "imagery"

INPUT_FOLDER_NAME = "int16_3031"
OUTPUT_FOLDER_NAME = "colcorrect"

OUT_NODATA = 0

# True = allow extra highlight compression for flat/dark images.
# False = never apply the extra highlight boost.
BOOST_BRIGHT = True

# Standard stretch.
P_LOW = 2.0
P_HIGH = 98.0

# Stronger bounds for flat images.
FLAT_P_LOW = 0.5
FLAT_P_HIGH = 99.5

# Treat images below this percentile-derived dynamic range as flat.
RANGE_THRESHOLD = 800.0

# For flat, non-bright images, use only this fraction of the flat range.
FLAT_HIGHLIGHT_FRACTION = 0.80

# Already-bright detector.
BRIGHT_MED_NORM = 0.70
BRIGHT_SAT_FRAC = 0.20

COG_OPTIONS = [
    "COMPRESS=DEFLATE",
    "BLOCKSIZE=512",
    "BIGTIFF=IF_SAFER",
    "NUM_THREADS=ALL_CPUS",
]


# =============================================================================
# HISTOGRAM / STATISTICS
# =============================================================================

def percentile_from_hist(
    hist: list[int],
    minv: float,
    maxv: float,
    percentile: float,
) -> float:

    total = sum(hist)

    if total <= 0:
        return float("nan")

    target = total * (percentile / 100.0)
    cumulative = 0
    bins = len(hist)

    for i, count in enumerate(hist):
        cumulative += count

        if cumulative >= target:
            fraction = (i + 0.5) / bins
            return minv + fraction * (maxv - minv)

    return maxv


def get_histogram(
    src: Path,
) -> tuple[float | None, float | None, list[int]]:

    ds = gdal.Open(str(src), gdal.GA_ReadOnly)

    if ds is None:
        raise RuntimeError(f"Could not open: {src}")

    band = ds.GetRasterBand(1)
    result = band.GetDefaultHistogram(force=True)

    ds = None

    if result is None:
        return None, None, []

    minv, maxv, _, hist = result

    return float(minv), float(maxv), list(hist)


def get_minmax_stats(
    src: Path,
) -> tuple[float, float, float | None]:

    ds = gdal.Open(str(src), gdal.GA_ReadOnly)

    if ds is None:
        raise RuntimeError(f"Could not open: {src}")

    band = ds.GetRasterBand(1)

    stats = band.GetStatistics(
        False,
        True,
    )

    ds = None

    if stats is None:
        raise RuntimeError(
            f"Could not derive statistics for {src.name}"
        )

    vmin, vmax, mean, _ = stats

    return float(vmin), float(vmax), float(mean)


def get_bounds(
    src: Path,
    low: float,
    high: float,
) -> tuple[float, float, str]:

    minv, maxv, hist = get_histogram(src)

    if minv is not None and maxv is not None and hist:
        lo = percentile_from_hist(
            hist, minv, maxv, low
        )
        hi = percentile_from_hist(
            hist, minv, maxv, high
        )

        return lo, hi, "hist"

    vmin, vmax, _ = get_minmax_stats(src)

    return vmin, vmax, "stats"


def is_already_bright(
    src: Path,
) -> tuple[bool, str]:

    minv, maxv, hist = get_histogram(src)

    if minv is not None and maxv is not None and hist:
        total = sum(hist) or 1

        median = percentile_from_hist(
            hist,
            minv,
            maxv,
            50.0,
        )

        denom = maxv - minv if maxv > minv else 1.0
        med_norm = (median - minv) / denom

        # Fraction of pixels within approximately the top 1% of histogram bins.
        top_bins = max(
            1,
            int(round(len(hist) * 0.01)),
        )

        sat = sum(hist[-top_bins:]) / total

        bright = (
            med_norm >= BRIGHT_MED_NORM
            or sat >= BRIGHT_SAT_FRAC
        )

        return (
            bright,
            f"hist med_norm={med_norm:.3f}, "
            f"sat(top~1%)={sat:.3f}",
        )

    vmin, vmax, mean = get_minmax_stats(src)

    if mean is None:
        return False, "stats mean unavailable"

    denom = vmax - vmin if vmax > vmin else 1.0
    mean_norm = (mean - vmin) / denom

    bright = mean_norm >= BRIGHT_MED_NORM

    return bright, f"stats mean_norm={mean_norm:.3f}"


# =============================================================================
# STRETCH
# =============================================================================

def safe_scale_bounds(
    in_low: float,
    in_high: float,
) -> tuple[float, float]:

    if in_low != in_low or in_high != in_high:
        return 0.0, 1.0

    if in_high <= in_low:
        return in_low, in_low + 1.0

    return in_low, in_high


def translate_stretch(
    src: Path,
    dst: Path,
    in_low: float,
    in_high: float,
) -> None:

    dst.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    in_low, in_high = safe_scale_bounds(
        in_low,
        in_high,
    )

    if dst.exists():
        dst.unlink()

    result = gdal.Translate(
        str(dst),
        str(src),
        options=gdal.TranslateOptions(
            format="COG",
            outputType=gdal.GDT_Byte,
            noData=OUT_NODATA,
            scaleParams=[
                [in_low, in_high, 1, 255]
            ],
            creationOptions=COG_OPTIONS,
        ),
    )

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

    result = None


# =============================================================================
# PROCESSING
# =============================================================================

def process_tif(src: Path) -> None:
    shelf_dir = src.parent.parent

    output_dir = (
        shelf_dir
        / OUTPUT_FOLDER_NAME
    )

    dst = (
        output_dir
        / src.name
    )

    if dst.exists() and dst.stat().st_size > 0:
        print(
            f"[SKIP] "
            f"{dst.relative_to(BASE_DIR)}"
        )
        return

    lo, hi, method = get_bounds(
        src,
        P_LOW,
        P_HIGH,
    )

    dynamic_range = hi - lo

    bright, bright_reason = is_already_bright(
        src
    )

    print()
    print(
        f"[INPUT] "
        f"{src.relative_to(BASE_DIR)}"
    )

    print(
        f"  standard ({method}): "
        f"low={lo:.2f}, high={hi:.2f}, "
        f"range={dynamic_range:.2f}"
    )

    print(
        f"  bright? {bright} "
        f"({bright_reason}) "
        f"BOOST_BRIGHT={BOOST_BRIGHT}"
    )

    if dynamic_range < RANGE_THRESHOLD:
        flat_lo, flat_hi, flat_method = get_bounds(
            src,
            FLAT_P_LOW,
            FLAT_P_HIGH,
        )

        flat_range = flat_hi - flat_lo

        print(
            f"  -> FLAT: bounds ({flat_method}) "
            f"{FLAT_P_LOW}-{FLAT_P_HIGH} "
            f"low={flat_lo:.2f}, "
            f"high={flat_hi:.2f}, "
            f"range={flat_range:.2f}"
        )

        if not BOOST_BRIGHT or bright:

            if bright:
                print(
                    "     flat but already bright "
                    "-> STANDARD bounds"
                )

                translate_stretch(
                    src,
                    dst,
                    lo,
                    hi,
                )

            else:
                print(
                    "     flat and BOOST_BRIGHT=False "
                    "-> FLAT bounds"
                )

                translate_stretch(
                    src,
                    dst,
                    flat_lo,
                    flat_hi,
                )

        else:
            lift_hi = (
                flat_lo
                + flat_range
                * FLAT_HIGHLIGHT_FRACTION
            )

            print(
                f"     highlight fraction "
                f"{FLAT_HIGHLIGHT_FRACTION:.2f} "
                f"-> scale high={lift_hi:.2f}"
            )

            translate_stretch(
                src,
                dst,
                flat_lo,
                lift_hi,
            )

    else:
        print(
            f"  -> OK: using "
            f"{P_LOW}-{P_HIGH}"
        )

        translate_stretch(
            src,
            dst,
            lo,
            hi,
        )

    print(
        f"[OK] "
        f"{dst.relative_to(BASE_DIR)}"
    )


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

def main() -> None:
    gdal.UseExceptions()

    tif_files = sorted(
        path
        for path in IMAGERY_DIR.rglob("*.tif")
        if path.parent.name == INPUT_FOLDER_NAME
    )

    if not tif_files:
        raise FileNotFoundError(
            f"No TIFFs found in "
            f"imagery/**/{INPUT_FOLDER_NAME}/"
        )

    print(
        f"Found {len(tif_files)} TIFF(s)"
    )

    for src in tif_files:
        try:
            process_tif(src)

        except Exception as error:
            print(
                f"[FAIL] "
                f"{src.relative_to(BASE_DIR)}: "
                f"{error}"
            )

    print()
    print("FINISHED")


if __name__ == "__main__":
    main()