#!/usr/bin/env python3

# Convert ensemble CSVs of 1° grid-cell centres into multiband GeoTIFFs.
#
# Bands:
#   1 = min_change
#   2 = max_change
#   3 = mean_change
#   4 = median_change
#   5 = sd_change
#
# Install:
#   pip install numpy pandas rasterio

from pathlib import Path
import re

import numpy as np
import pandas as pd
import rasterio
from rasterio.transform import from_origin


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

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

NODATA = -9999.0

BANDS = [
    "min_change",
    "max_change",
    "mean_change",
    "median_change",
    "sd_change",
]


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

def build_output_name(csv_path):
    """
    Convert:
        ensemble_perc_change_fish_bio_ssp126_2041-2050_antarctica.csv

    Into:
        percchange_fish_bio__SSP1__2041-2050__Antarctica.tif
    """

    parts = csv_path.stem.split("_")

    if len(parts) < 7:
        raise ValueError(
            f"Unexpected filename format: {csv_path.name}"
        )

    metric = parts[1] + parts[2]
    variable = parts[3] + "_" + parts[4]
    scenario_raw = parts[5]
    decade = parts[6]
    region_raw = "_".join(parts[7:])

    match = re.match(
        r"ssp(\d)\d\d",
        scenario_raw,
        re.IGNORECASE,
    )

    scenario = (
        f"SSP{match.group(1)}"
        if match
        else scenario_raw.upper()
    )

    region = "".join(
        word.capitalize()
        for word in region_raw.split("_")
    )

    return (
        f"{metric}_{variable}"
        f"__{scenario}"
        f"__{decade}"
        f"__{region}.tif"
    )


def csv_to_tif(csv_path, output_path):

    df = pd.read_csv(csv_path)

    required = {
        "longitude",
        "latitude",
        *BANDS,
    }

    missing = required - set(df.columns)

    if missing:
        raise ValueError(
            f"{csv_path.name} missing columns: {missing}"
        )

    lons = np.sort(df["longitude"].unique())
    lats = np.sort(df["latitude"].unique())

    ncols = len(lons)
    nrows = len(lats)

    lon_to_col = {
        lon: i
        for i, lon in enumerate(lons)
    }

    lat_to_row = {
        lat: nrows - 1 - i
        for i, lat in enumerate(lats)
    }

    west = lons.min() - 0.5
    north = lats.max() + 0.5

    transform = from_origin(
        west,
        north,
        1.0,
        1.0,
    )

    band_arrays = [
        np.full(
            (nrows, ncols),
            NODATA,
            dtype=np.float32,
        )
        for _ in BANDS
    ]

    for row in df.itertuples(index=False):

        lon = getattr(row, "longitude")
        lat = getattr(row, "latitude")

        if lon not in lon_to_col or lat not in lat_to_row:
            continue

        r = lat_to_row[lat]
        c = lon_to_col[lon]

        for band_index, column in enumerate(BANDS):

            value = getattr(row, column)

            if pd.isna(value):
                continue

            band_arrays[band_index][r, c] = float(value)

    profile = {
        "driver": "GTiff",
        "height": nrows,
        "width": ncols,
        "count": len(BANDS),
        "dtype": "float32",
        "crs": "EPSG:4326",
        "transform": transform,
        "nodata": NODATA,
        "compress": "DEFLATE",
    }

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

        for band_index, (array, description) in enumerate(
            zip(band_arrays, BANDS),
            start=1,
        ):
            dst.write(array, band_index)
            dst.set_band_description(
                band_index,
                description,
            )

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


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

def main():

    if not INPUT_DIR.exists():
        raise FileNotFoundError(
            f"Could not find input folder: {INPUT_DIR}"
        )

    OUTPUT_DIR.mkdir(exist_ok=True)

    csv_files = sorted(
        INPUT_DIR.glob("*.csv")
    )

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

    for csv_path in csv_files:

        output_path = (
            OUTPUT_DIR
            / build_output_name(csv_path)
        )

        csv_to_tif(
            csv_path,
            output_path,
        )

    print("Done.")


if __name__ == "__main__":
    main()