# Convert Kerguelen Plateau CSV predictions to multiband Cloud Optimised GeoTIFFs.
#
# Required packages:
#   pip install pandas numpy
#   Also requires GDAL (from osgeo import gdal, osr)

from pathlib import Path

import numpy as np
import pandas as pd
from osgeo import gdal, osr


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

INPUT_CSV = "KP_RCP_Predictions.csv"

RCP_OUTPUT = "KerguelenPlateau_RCP_prob.tif"
HCLASS_OUTPUT = "KerguelenPlateau_HClass.tif"

CELLSIZE = 0.1
NODATA = -9999


# =========================
# SETUP
# =========================

INPUT_DIR = Path.cwd() / "input"
OUTPUT_DIR = Path.cwd() / "output"
OUTPUT_DIR.mkdir(exist_ok=True)

df = pd.read_csv(INPUT_DIR / INPUT_CSV)

outputs = {
    RCP_OUTPUT: [
        ("RCP_1", "RCP #1 prob"),
        ("RCP_2", "RCP #2 prob"),
        ("RCP_3", "RCP #3 prob"),
        ("RCP_4", "RCP #4 prob"),
        ("RCP_5", "RCP #5 prob"),
        ("RCP_6", "RCP #6 prob"),
        ("RCP_7", "RCP #7 prob"),
    ],
    HCLASS_OUTPUT: [
        ("HClass", "Hard class RCP #"),
        ("HClass_prob", "Hard class RCP prob"),
    ],
}


# =========================
# GRID
# =========================

minx = df["Longitude"].min() - CELLSIZE / 2
maxx = df["Longitude"].max() + CELLSIZE / 2
miny = df["Latitude"].min() - CELLSIZE / 2
maxy = df["Latitude"].max() + CELLSIZE / 2

width = round((maxx - minx) / CELLSIZE)
height = round((maxy - miny) / CELLSIZE)

geotransform = (minx, CELLSIZE, 0, maxy, 0, -CELLSIZE)

srs = osr.SpatialReference()
srs.ImportFromEPSG(4326)

# Calculate raster positions once
cols = np.floor((df["Longitude"].to_numpy() - minx) / CELLSIZE).astype(int)
rows = np.floor((maxy - df["Latitude"].to_numpy()) / CELLSIZE).astype(int)


# =========================
# CREATE RASTERS
# =========================

def create_cog(output_name, bands):
    ds = gdal.GetDriverByName("MEM").Create(
        "", width, height, len(bands), gdal.GDT_Float32
    )

    ds.SetGeoTransform(geotransform)
    ds.SetProjection(srs.ExportToWkt())

    for band_number, (column, description) in enumerate(bands, start=1):
        data = np.full((height, width), NODATA, dtype=np.float32)
        data[rows, cols] = df[column].to_numpy()

        band = ds.GetRasterBand(band_number)
        band.WriteArray(data)
        band.SetNoDataValue(NODATA)
        band.SetDescription(column)
        band.SetMetadata({"long_name": description})

        stats = band.GetStatistics(True, True)
        band.SetStatistics(*stats)

    output_path = OUTPUT_DIR / output_name

    gdal.Translate(
        str(output_path),
        ds,
        format="COG",
        creationOptions=["COMPRESS=DEFLATE"],
    )

    print(f"Created: {output_path}")


for output_name, bands in outputs.items():
    create_cog(output_name, bands)