#!/usr/bin/env python3

# Create monthly and annual whale habitat-suitability raster products.
#
# Usage:
#   Run from the folder containing the four species folders:
#
#   python 02_process_habitat.py
#
# Expected structure:
#
#   cwd/
#   ├── Antarctic_Blue_Whale/
#   │   └── input/
#   │       ├── ABW_Mean_01_Jan.tif
#   │       ├── ABW_SD_01_Jan.tif
#   │       ├── ABW_Qnt_1_01_Jan.tif
#   │       ├── ABW_Qnt_9_01_Jan.tif
#   │       └── ...
#   ├── Antarctic_Minke_Whale/
#   │   └── input/
#   ├── Fin_Whale/
#   │   └── input/
#   └── Humpback_Whale/
#       └── input/
#
# Output for each species:
#
#   <species>/output/<PREFIX>_Habitat_Monthly.tif
#   <species>/output/<PREFIX>_Habitat_Annual.tif
#
# Monthly output contains 48 bands:
#   January_mean ... December_mean
#   January_SD   ... December_SD
#   January_Q10  ... December_Q10
#   January_Q90  ... December_Q90
#
# Annual output contains 2 bands:
#   mean_all_months
#   SD_all_months
#
# Install:
#   pip install numpy rasterio

from pathlib import Path
import tempfile

import numpy as np
import rasterio
from rasterio.warp import calculate_default_transform, reproject, Resampling
from rasterio.shutil import copy as rio_copy


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

BASE_DIR = Path.cwd()

SPECIES_FOLDERS = [
    "Antarctic_Blue_Whale",
    "Antarctic_Minke_Whale",
    "Fin_Whale",
    "Humpback_Whale",
]

TARGET_CRS = "EPSG:3031"
NODATA = -3.4e38

# Number of years represented by each monthly summary
N_YEARS = 18

MONTHS = [
    ("01_Jan", "January"),
    ("02_Feb", "February"),
    ("03_Mar", "March"),
    ("04_Apr", "April"),
    ("05_May", "May"),
    ("06_Jun", "June"),
    ("07_Jul", "July"),
    ("08_Aug", "August"),
    ("09_Sep", "September"),
    ("10_Oct", "October"),
    ("11_Nov", "November"),
    ("12_Dec", "December"),
]

STATISTICS = [
    ("Mean", "mean"),
    ("SD", "SD"),
    ("Qnt_1", "Q10"),
    ("Qnt_9", "Q90"),
]


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

def find_species_prefix(input_dir: Path) -> str:
    """Detect species prefix from files such as ABW_Mean_01_Jan.tif."""

    mean_files = sorted(input_dir.glob("*_Mean_01_Jan.tif"))

    if not mean_files:
        raise RuntimeError(
            f"No file matching '*_Mean_01_Jan.tif' found in {input_dir}"
        )

    return mean_files[0].name.split("_")[0]


def get_target_grid(reference_file: Path):
    """Create the EPSG:3031 output grid from the first input raster."""

    with rasterio.open(reference_file) as src:

        transform, width, height = calculate_default_transform(
            src.crs,
            TARGET_CRS,
            src.width,
            src.height,
            *src.bounds,
        )

        profile = src.profile.copy()

    profile.update(
        driver="GTiff",
        crs=TARGET_CRS,
        transform=transform,
        width=width,
        height=height,
        count=1,
        dtype="float32",
        nodata=NODATA,
        compress="DEFLATE",
        tiled=True,
        blockxsize=512,
        blockysize=512,
        BIGTIFF="YES",
    )

    return profile


def read_reprojected(path: Path, target_profile: dict) -> np.ndarray:
    """Read one raster and reproject it to the common EPSG:3031 grid."""

    destination = np.full(
        (target_profile["height"], target_profile["width"]),
        NODATA,
        dtype="float32",
    )

    with rasterio.open(path) as src:

        source = src.read(1)

        reproject(
            source=source,
            destination=destination,
            src_transform=src.transform,
            src_crs=src.crs,
            src_nodata=src.nodata,
            dst_transform=target_profile["transform"],
            dst_crs=target_profile["crs"],
            dst_nodata=NODATA,
            resampling=Resampling.bilinear,
        )

    return destination


def valid_mask(array: np.ndarray) -> np.ndarray:
    """Return valid, finite cells excluding output NoData."""

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


def write_cog(
    output_path: Path,
    arrays: list,
    descriptions: list,
    target_profile: dict,
):
    """Write arrays to a temporary GeoTIFF and convert to a COG."""

    profile = target_profile.copy()

    profile.update(
        count=len(arrays),
        driver="GTiff",
        compress="DEFLATE",
        tiled=True,
        BIGTIFF="YES",
    )

    with tempfile.TemporaryDirectory() as temp_dir:

        temp_path = Path(temp_dir) / "temp.tif"

        with rasterio.open(temp_path, "w", **profile) as dst:

            for band_number, (array, description) in enumerate(
                zip(arrays, descriptions),
                start=1,
            ):
                dst.write(array.astype("float32"), band_number)
                dst.set_band_description(band_number, description)

        rio_copy(
            temp_path,
            output_path,
            driver="COG",
            compress="DEFLATE",
            BIGTIFF="YES",
        )


def calculate_annual(mean_arrays, sd_arrays):
    """
    Calculate annual mean and pooled annual sample SD.

    Each calendar month represents N_YEARS observations.

    Annual variance combines:
      1. within-month variance
      2. between-month variance
    """

    means = np.stack(mean_arrays).astype("float64")
    sds = np.stack(sd_arrays).astype("float64")

    mean_valid = np.isfinite(means) & (means != NODATA)
    sd_valid = np.isfinite(sds) & (sds != NODATA)

    # ---------------------------------------------------------
    # ANNUAL MEAN
    # ---------------------------------------------------------

    means_nan = np.where(mean_valid, means, np.nan)

    with np.errstate(invalid="ignore"):
        annual_mean = np.nanmean(means_nan, axis=0)

    annual_mean_out = np.full(
        annual_mean.shape,
        NODATA,
        dtype="float32",
    )

    good_mean = np.isfinite(annual_mean)

    annual_mean_out[good_mean] = annual_mean[good_mean].astype(
        "float32"
    )

    # ---------------------------------------------------------
    # ANNUAL SD
    # ---------------------------------------------------------

    # A month can contribute to annual SD only where both its
    # monthly mean and monthly SD are valid.
    pair_valid = mean_valid & sd_valid

    means_pair = np.where(pair_valid, means, np.nan)
    sds_pair = np.where(pair_valid, sds, np.nan)

    number_months = pair_valid.sum(axis=0)

    with np.errstate(invalid="ignore"):
        combined_mean = np.nanmean(means_pair, axis=0)

    # Within-month sum of squares:
    #   (n - 1) * SD_month^2
    within_ss = np.nansum(
        (N_YEARS - 1) * (sds_pair ** 2),
        axis=0,
    )

    # Between-month sum of squares:
    #   n * (mean_month - overall_mean)^2
    between_ss = np.nansum(
        N_YEARS
        * ((means_pair - combined_mean[np.newaxis, :, :]) ** 2),
        axis=0,
    )

    total_n = number_months * N_YEARS

    annual_variance = np.full(
        combined_mean.shape,
        np.nan,
        dtype="float64",
    )

    usable = total_n > 1

    annual_variance[usable] = (
        within_ss[usable] + between_ss[usable]
    ) / (total_n[usable] - 1)

    annual_sd = np.sqrt(annual_variance)

    annual_sd_out = np.full(
        annual_sd.shape,
        NODATA,
        dtype="float32",
    )

    good_sd = np.isfinite(annual_sd)

    annual_sd_out[good_sd] = annual_sd[good_sd].astype(
        "float32"
    )

    return annual_mean_out, annual_sd_out


# =========================
# PROCESS ONE SPECIES
# =========================

def process_species(species_folder: Path):

    input_dir = species_folder / "input"
    output_dir = species_folder / "output"

    if not input_dir.is_dir():
        print(f"Skipped: input folder not found: {input_dir}")
        return False

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

    prefix = find_species_prefix(input_dir)

    monthly_output = (
        output_dir / f"{prefix}_Habitat_Monthly.tif"
    )

    annual_output = (
        output_dir / f"{prefix}_Habitat_Annual.tif"
    )

    print(f"\nProcessing {species_folder.name}")
    print(f"  Prefix: {prefix}")

    # Reference raster determines the common EPSG:3031 grid.
    reference_file = (
        input_dir / f"{prefix}_Mean_01_Jan.tif"
    )

    target_profile = get_target_grid(reference_file)

    # Store arrays by statistic so the final monthly product
    # is ordered Mean -> SD -> Q10 -> Q90.
    monthly_arrays = {
        "mean": [],
        "SD": [],
        "Q10": [],
        "Q90": [],
    }

    monthly_descriptions = {
        "mean": [],
        "SD": [],
        "Q10": [],
        "Q90": [],
    }

    # ---------------------------------------------------------
    # READ + REPROJECT ALL MONTHLY INPUTS
    # ---------------------------------------------------------

    for file_stat, description_stat in STATISTICS:

        print(f"  Processing {description_stat} rasters...")

        for month_code, month_name in MONTHS:

            input_file = (
                input_dir
                / f"{prefix}_{file_stat}_{month_code}.tif"
            )

            if not input_file.is_file():
                raise RuntimeError(
                    f"Missing required input: {input_file}"
                )

            array = read_reprojected(
                input_file,
                target_profile,
            )

            monthly_arrays[description_stat].append(array)

            monthly_descriptions[description_stat].append(
                f"{month_name}_{description_stat}"
            )

    # ---------------------------------------------------------
    # MONTHLY 48-BAND PRODUCT
    # ---------------------------------------------------------

    output_arrays = []
    output_descriptions = []

    for statistic in ["mean", "SD", "Q10", "Q90"]:

        output_arrays.extend(
            monthly_arrays[statistic]
        )

        output_descriptions.extend(
            monthly_descriptions[statistic]
        )

    print("  Creating monthly 48-band raster...")

    write_cog(
        monthly_output,
        output_arrays,
        output_descriptions,
        target_profile,
    )

    print(f"  Created: {monthly_output.name}")

    # ---------------------------------------------------------
    # ANNUAL PRODUCT
    # ---------------------------------------------------------

    print("  Calculating annual mean and SD...")

    annual_mean, annual_sd = calculate_annual(
        monthly_arrays["mean"],
        monthly_arrays["SD"],
    )

    write_cog(
        annual_output,
        [annual_mean, annual_sd],
        ["mean_all_months", "SD_all_months"],
        target_profile,
    )

    print(f"  Created: {annual_output.name}")

    return True


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

def main():

    processed = 0
    skipped = 0

    print("Processing whale habitat suitability rasters...")

    for folder_name in SPECIES_FOLDERS:

        species_folder = BASE_DIR / folder_name

        if not species_folder.is_dir():
            print(f"\nSkipped: species folder not found: {species_folder}")
            skipped += 1
            continue

        try:
            success = process_species(species_folder)

            if success:
                processed += 1
            else:
                skipped += 1

        except Exception as e:
            print(f"  Failed: {e}")
            skipped += 1

    print("\nComplete.")
    print(f"Species processed: {processed}")
    print(f"Species skipped/failed: {skipped}")


if __name__ == "__main__":
    main()