#!/usr/bin/env python3

# Calculate mean and interannual standard deviation of annual, summer and
# winter particulate organic carbon flux across all available NetCDF files,
# then reproject to EPSG:3031 COG.
#
# Input:
#   input/POC.export_YYYY.nc
#   input/POC.export_YYYY_summer.nc
#   input/POC.export_YYYY_winter.nc
#
# Output:
#   output/particulate_organic_C_flux.tif
#
# Bands:
#   1 annual_mean
#   2 summer_mean
#   3 winter_mean
#   4 annual_SD
#   5 summer_SD
#   6 winter_SD
#
# Install requirements:
#   pip install numpy rasterio

from pathlib import Path
import re
import numpy as np
import rasterio
from rasterio.warp import calculate_default_transform, reproject, Resampling

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

INPUT_DIR = Path("input")
OUTPUT_DIR = Path("output")
OUTPUT_FILE = OUTPUT_DIR / "particulate_organic_C_flux.tif"

TARGET_CRS = "EPSG:3031"
NODATA = -999.0

GROUPS = [
    ("annual", re.compile(r"^POC\.export_(\d{4})\.nc$")),
    ("summer", re.compile(r"^POC\.export_(\d{4})_summer\.nc$")),
    ("winter", re.compile(r"^POC\.export_(\d{4})_winter\.nc$")),
]


# =========================
# FUNCTIONS
# =========================

def find_files(pattern):
    matches = []

    for path in INPUT_DIR.glob("*.nc"):
        match = pattern.match(path.name)
        if match:
            matches.append((int(match.group(1)), path))

    return sorted(matches)


def calculate_stats(files, reference):
    total = np.zeros((reference.height, reference.width), dtype=np.float64)
    total_sq = np.zeros((reference.height, reference.width), dtype=np.float64)
    count = np.zeros((reference.height, reference.width), dtype=np.uint16)

    for year, path in files:
        print(f"    {year}: {path.name}")

        with rasterio.open(path) as src:
            if src.width != reference.width or src.height != reference.height:
                raise RuntimeError(f"Grid size differs: {path}")

            if src.transform != reference.transform:
                raise RuntimeError(f"Transform differs: {path}")

            data = src.read(1).astype(np.float64)
            valid = np.isfinite(data)

            if src.nodata is not None:
                valid &= data != src.nodata

            total[valid] += data[valid]
            total_sq[valid] += data[valid] ** 2
            count[valid] += 1

    mean = np.full(total.shape, NODATA, dtype=np.float32)
    sd = np.full(total.shape, NODATA, dtype=np.float32)

    valid_mean = count > 0
    mean[valid_mean] = (total[valid_mean] / count[valid_mean]).astype(np.float32)

    # Sample standard deviation across years (n)
    valid_sd = count > 1
    variance = (
        total_sq[valid_sd]
        - (total[valid_sd] ** 2 / count[valid_sd])
    ) / (count[valid_sd] - 1)

    # Protect against tiny negative values caused by floating-point rounding
    variance = np.maximum(variance, 0)
    sd[valid_sd] = np.sqrt(variance).astype(np.float32)

    return mean, sd


def warp_to_3031(data, reference, transform, width, height):
    output = np.full((height, width), NODATA, dtype=np.float32)

    reproject(
        source=data,
        destination=output,
        src_transform=reference.transform,
        src_crs=reference.crs or "EPSG:4326",
        src_nodata=NODATA,
        dst_transform=transform,
        dst_crs=TARGET_CRS,
        dst_nodata=NODATA,
        resampling=Resampling.bilinear,
    )

    return output


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

def main():
    groups = [(name, find_files(pattern)) for name, pattern in GROUPS]

    for name, files in groups:
        if not files:
            raise RuntimeError(f"No files found for {name}.")

        print(f"{name}: {len(files)} files ({files[0][0]}-{files[-1][0]})")

    reference_path = groups[0][1][0][1]

    with rasterio.open(reference_path) as reference:
        src_crs = reference.crs or "EPSG:4326"

        # Let GDAL determine an appropriate EPSG:3031 output resolution
        transform, width, height = calculate_default_transform(
            src_crs,
            TARGET_CRS,
            reference.width,
            reference.height,
            *reference.bounds,
        )

        print(f"\nOutput grid: {width} x {height}")
        print(f"Resolution: {abs(transform.a):,.1f} m")
        print(f"CRS: {TARGET_CRS}")

        # Calculate all statistics before writing, so means can occupy bands 1-3
        # and standard deviations bands 4-6.
        results = {}

        for name, files in groups:
            print(f"\nCalculating {name} statistics")
            mean, sd = calculate_stats(files, reference)

            print("  Reprojecting mean to EPSG:3031")
            mean = warp_to_3031(mean, reference, transform, width, height)

            print("  Reprojecting SD to EPSG:3031")
            sd = warp_to_3031(sd, reference, transform, width, height)

            results[name] = (mean, sd)

        profile = {
            "driver": "COG",
            "width": width,
            "height": height,
            "count": 6,
            "dtype": "float32",
            "crs": TARGET_CRS,
            "transform": transform,
            "nodata": NODATA,
            "compress": "DEFLATE",
            "blocksize": 512,
            "overview_resampling": "AVERAGE",
            "interleave": "band", # Optimise for web mapping display (read each band singly)
            "BIGTIFF": "IF_SAFER",
        }

        OUTPUT_DIR.mkdir(exist_ok=True)

        if OUTPUT_FILE.exists():
            OUTPUT_FILE.unlink()

        band_data = [
            ("annual_mean", results["annual"][0]),
            ("summer_mean", results["summer"][0]),
            ("winter_mean", results["winter"][0]),
            ("annual_SD", results["annual"][1]),
            ("summer_SD", results["summer"][1]),
            ("winter_SD", results["winter"][1]),
        ]

        with rasterio.open(OUTPUT_FILE, "w", **profile) as dst:
            for band_no, (name, data) in enumerate(band_data, 1):
                dst.write(data, band_no)
                dst.set_band_description(band_no, name)
                dst.update_tags(band_no, units="g C m-2 day-1")

                print(
                    f"Band {band_no}: {name} "
                    f"({np.count_nonzero(data != NODATA):,} valid cells)"
                )

    print(f"\nCreated: {OUTPUT_FILE}")


if __name__ == "__main__":
    main()