# ruff: noqa: PD003
"""Download and preprocess ERA5 data.

Variables of Interest:
- 2 metre temperature (t2m) [instant]
- Total precipitation (tp) [accum]
- Snow Fall (sf) [accum]
- Snow cover (snowc) [instant]
- Snow depth (sde) [instant]
- Surface sensible heat flux (sshf) [accum]
- Lake ice bottom temperature (lblt) [instant]

Snow Fall and Total precipitation are both further accumulated over a day -
thus instead of taking the daily sum, only the last value of the day is taken.
This is 00:00 of the next day, e.g.:
tp_2020_06_23 <- tp_2020_06_24_00:00
See: https://confluence.ecmwf.int/display/CKB/ERA5-Land%3A+data+documentation#ERA5Land:datadocumentation-accumulationsAccumulations

Naming patterns:
- Instant Variables are downloaded already as statistically aggregated (lossy),
  therefore their names get the aggregation as suffix
- Accumulation Variables are downloaded as totals (sum), their names stay the same

Daily Variables (downloaded from hourly data):
- t2m_max
- t2m_min
- t2m_mean
- snowc_mean
- sde_mean
- lblt_max
- tp
- sf
- sshf

Derived Daily Variables:
- t2m_range [instant]: t2m_max - t2m_min
- t2m_avg [instant]: (t2m_max - t2m_min) / 2
- t2m_skew [instant]: (t2m_mean - t2m_min) / t2m_range
- thawing_degree_days [accum]: (t2m_avg - 273.15).clip(min=0)
- freezing_degree_days [accum]: (273.15 - t2m_avg).clip(min=0)
- thawing_days [accum]: (t2m_avg > 273.15).astype(int)
- freezing_days [accum]: (t2m_avg < 273.15).astype(int)
- precipitation_occurrences [accum]: (tp > 0.001).astype(int)
- snowfall_occurrences [accum]: (sf > 0.001).astype(int)
- naive_snow_isolation [instant]: snowc_mean * sde_mean


Monthly, Winter, Summer & Yearly Aggregations (Names don't change):
- instant variables:
    - *_min -> min
    - *_max -> max
    - *_rest -> median
- accum variables: sum

Additionally:
- snow_cover_min [instant]: min of snowc_mean over month/winter/summer/year
- snow_cover_max [instant]: max of snowc_mean over month/winter/summer/year

Derived & (from monthly) Aggregated Winter Variables:
- effective_snow_depth [instant]: (sde_mean * M + 1 - m).sum(M) / (m).sum(M),see also https://tc.copernicus.org/articles/11/989/2017/tc-11-989-2017.pdf

Derived & (from daily) Aggregated Yearly Variables:
- day_of_first_thaw [yearly]: First day in year where t2m_daily_avg > 273.15
- day_of_last_thaw [yearly]: Last day in year where t2m_daily_avg > 273.15
- thawing_period_length [yearly]: day_of_last_thaw - day_of_first_thaw
- day_of_first_freeze [yearly]: First day in year where t2m_daily_avg < 273.15
- day_of_last_freeze [yearly]: Last day in year where t2m_daily_avg < 273.15

About yearly aggregates:
- A year always starts on 1st October and ends on 30th September of the next year
  to better capture the Arctic seasonal cycle.
- Thus year == 2020 means 1st Oct 2019 - 30th Sep 2020
- Thus winter == 2020 means 1st Oct 2019 - 31th March 2020
- Thus summer == 2020 means 1st April 2020 - 30th Sep 2020

Author: Tobias Hölzer
Date: June to October 2025
"""

import cProfile
import time
from typing import Literal

import cyclopts
import dask.distributed as dd
import numpy as np
import odc.geo
import odc.geo.xr
import pandas as pd
import shapely.geometry
import ultraplot as uplt
import xarray as xr
import xdggs
import xvec
from rasterio.features import shapes
from rich import pretty, print, traceback
from stopuhr import stopwatch

from entropice.spatial import grids, watermask
from entropice.spatial.aggregators import _Aggregations, aggregate_raster_into_grid
from entropice.spatial.xvec import to_xvec
from entropice.utils import codecs
from entropice.utils.paths import FIGURES_DIR, get_era5_stores
from entropice.utils.types import Grid

traceback.install(show_locals=True, suppress=[cyclopts, xr, pd, cProfile])
pretty.install()

cli = cyclopts.App()


min_lat = 50
max_lat = 83.7  # Ensures the right Chunks Size (90 - 64 / 10 + 0.1)
min_time = "1990-01-01"
max_time = "2024-12-31"


today = time.strftime("%Y-%m-%d")

instants = {
    "t2m_max",
    "t2m_min",
    "t2m_mean",
    "snowc_mean",
    "sde_mean",
    "lblt_max",
    "t2m_range",
    "t2m_avg",
    "t2m_skew",
    "naive_snow_isolation",
}

accums = {
    "tp",
    "sf",
    "sshf",
    "thawing_degree_days",
    "freezing_degree_days",
    "thawing_days",
    "freezing_days",
    "precipitation_occurrences",
    "snowfall_occurrences",
}


# ================
# === Download ===
# ================


def download_daily_aggregated():
    """Download and aggregate ERA5 data to daily resolution.

    Downloads ERA5 reanalysis data from the DESTINE Earth Data Hub and aggregates
    it to daily resolution. Includes temperature extremes, precipitation, snow,
    and surface heat flux variables.

    The function downloads hourly data and creates daily aggregates:
    - Temperature: daily min/max
    - Precipitation and snowfall: daily totals
    - Snow cover and depth: daily means
    - Surface heat flux: daily totals
    - Lake ice temperature: daily max

    The aggregated data is saved to a zarr file with compression.
    """
    era5 = xr.open_dataset(
        "https://data.earthdatahub.destine.eu/era5/reanalysis-era5-land-no-antartica-v0.zarr",
        storage_options={"client_kwargs": {"trust_env": True}},
        chunks={},
        # chunks={},
        engine="zarr",
    ).rename({"valid_time": "time"})
    subset = {
        "latitude": slice(max_lat, min_lat),
    }

    # Compute the clostest chunk-start to min_time, to avoid problems with cropped chunks at the start
    tchunksize = era5.chunksizes["time"][0]
    era5_chunk_starts = pd.date_range(era5.time.min().item(), era5.time.max().item(), freq=f"{tchunksize}h")
    closest_chunk_start = era5_chunk_starts[
        era5_chunk_starts.get_indexer([pd.to_datetime(min_time)], method="ffill")[0]  # ty:ignore[invalid-argument-type]
    ]
    subset["time"] = slice(str(closest_chunk_start), max_time)

    era5 = era5.sel(subset)

    daily_raw = xr.merge(
        [
            # Instant
            era5.t2m.resample(time="1D").max().rename("t2m_max").astype(np.float32),
            era5.t2m.resample(time="1D").min().rename("t2m_min").astype(np.float32),
            era5.t2m.resample(time="1D").mean().rename("t2m_mean").astype(np.float32),
            era5.snowc.resample(time="1D").mean().rename("snowc_mean").astype(np.float32),
            era5.sde.resample(time="1D").mean().rename("sde_mean").astype(np.float32),
            era5.lblt.resample(time="1D").max().rename("lblt_max").astype(np.float32),
            # Accum
            era5.sshf.resample(time="1D").sum(skipna=False).rename("sshf").astype(np.float32),
            # Precipitation and snow fall are special
            # Take only the last value of the day (00:00 of next day)
            era5.tp.resample(time="1D").first().shift(time=-1).rename("tp").astype(np.float32),
            era5.sf.resample(time="1D").first().shift(time=-1).rename("sf").astype(np.float32),
        ]
    )

    # Assign attributes
    daily_raw["t2m_max"].attrs = {
        "long_name": "Daily maximum 2 metre temperature",
        "units": "K",
    }
    daily_raw["t2m_min"].attrs = {
        "long_name": "Daily minimum 2 metre temperature",
        "units": "K",
    }
    daily_raw["t2m_mean"].attrs = {
        "long_name": "Daily mean 2 metre temperature",
        "units": "K",
    }
    daily_raw["snowc_mean"].attrs = {"long_name": "Daily mean snow cover", "units": "m"}
    daily_raw["sde_mean"].attrs = {"long_name": "Daily mean snow depth", "units": "m"}
    daily_raw["lblt_max"].attrs = {
        "long_name": "Daily maximum lake ice bottom temperature",
        "units": "K",
    }
    daily_raw["tp"].attrs = {
        "long_name": "Daily total precipitation",
        "units": "m",
    }  # Units are rather m^3 / m^2
    daily_raw["sf"].attrs = {
        "long_name": "Daily total snow fall",
        "units": "m",
    }  # Units are rather m^3 / m^2
    daily_raw["sshf"].attrs = {
        "long_name": "Daily total surface sensible heat flux",
        "units": "J/m²",
    }

    daily_raw = daily_raw.odc.assign_crs("epsg:4326")
    daily_raw = daily_raw.drop_vars(["surface", "number", "depthBelowLandLayer"])
    daily_raw.attrs = era5.attrs
    daily_store = get_era5_stores("daily")
    print(f"Saving downloaded and daily aggregated ERA5 data to {daily_store}.")
    daily_raw.to_zarr(daily_store, mode="w", encoding=codecs.from_ds(daily_raw), consolidated=False)


@cli.command
def download():
    """Download ERA5 data using Dask cluster for parallel processing.

    Creates a local Dask cluster and downloads daily aggregated ERA5 data.
    The cluster is configured with a single worker with 10 threads and 100GB
    memory limit for optimal performance.
    """
    with (
        dd.LocalCluster(n_workers=1, threads_per_worker=10, memory_limit="100GB") as cluster,
        dd.Client(cluster) as client,
    ):
        print(client)
        print(client.dashboard_link)
        download_daily_aggregated()
        print("Downloaded and aggregated ERA5 data.")


# ============================
# === Temporal Aggregation ===
# ============================


def daily_enrich():
    """Enrich daily ERA5 data with derived climate variables.

    Loads downloaded daily ERA5 data and computes additional climate variables.
    Creates derived variables including temperature statistics, degree days, and occurrence indicators.

    Derived variables include:
    - Daily average and range temperature
    - Temperature skewness
    - Thawing and freezing degree days
    - Thawing and freezing day counts
    - Precipitation and snowfall occurrences
    - Snow isolation index

    """
    daily_store = get_era5_stores("daily")
    daily = xr.open_zarr(daily_store, consolidated=False).set_coords("spatial_ref")
    assert "time" in daily.dims, f"Expected dim 'time' to be in {daily.dims=}"
    # Keep track of stored variables for re-calculation and encodings
    stored_vars = list(daily.data_vars)

    # For better dask performance, all variables are written immediately after calculation
    # The smart scheduling which could performantly handle more is according to docs not yet implemented.
    # See https://docs.xarray.dev/en/stable/user-guide/dask.html#best-practices point 3 or https://github.com/dask/dask/issues/874
    def _store(v: str):
        nonlocal daily

        encoding = codecs.from_ds(daily[[v]], include_coords=False) if v not in stored_vars else None
        print(f"Storing enriched daily variable {v} to {daily_store}...")
        with stopwatch("Storing enriched daily variable"):
            daily[[v]].to_zarr(daily_store, mode="a", encoding=encoding, consolidated=False)
        stored_vars.append(v)
        daily = xr.open_zarr(daily_store, consolidated=False).set_coords("spatial_ref")

    # Formulas based on Groeke et. al. (2025) Stochastic Weather generation...
    daily["t2m_avg"] = (daily.t2m_max + daily.t2m_min) / 2
    daily.t2m_avg.attrs = {
        "long_name": "Daily average 2 metre temperature",
        "units": "K",
    }
    _store("t2m_avg")

    daily["t2m_range"] = daily.t2m_max - daily.t2m_min
    daily.t2m_range.attrs = {
        "long_name": "Daily range of 2 metre temperature",
        "units": "K",
    }
    _store("t2m_range")

    with np.errstate(invalid="ignore"):
        daily["t2m_skew"] = (daily.t2m_mean - daily.t2m_min) / daily.t2m_range
        daily.t2m_skew.attrs = {"long_name": "Daily skewness of 2 metre temperature"}
        _store("t2m_skew")

    daily["thawing_degree_days"] = (daily.t2m_avg - 273.15).clip(min=0)
    daily.thawing_degree_days.attrs = {"long_name": "Thawing degree days", "units": "K"}
    _store("thawing_degree_days")

    daily["freezing_degree_days"] = (273.15 - daily.t2m_avg).clip(min=0)
    daily.freezing_degree_days.attrs = {
        "long_name": "Freezing degree days",
        "units": "K",
    }
    _store("freezing_degree_days")

    daily["thawing_days"] = (daily.t2m_avg > 273.15).where(~daily.t2m_avg.isnull())
    daily.thawing_days.attrs = {"long_name": "Thawing days"}
    _store("thawing_days")

    daily["freezing_days"] = (daily.t2m_avg < 273.15).where(~daily.t2m_avg.isnull())
    daily.freezing_days.attrs = {"long_name": "Freezing days"}
    _store("freezing_days")

    daily["precipitation_occurrences"] = (daily.tp > 0.001).where(~daily.tp.isnull())
    daily.precipitation_occurrences.attrs = {"long_name": "Precipitation occurrences"}
    _store("precipitation_occurrences")

    daily["snowfall_occurrences"] = (daily.sf > 0.001).where(~daily.sf.isnull())
    daily.snowfall_occurrences.attrs = {"long_name": "Snowfall occurrences"}
    _store("snowfall_occurrences")

    daily["naive_snow_isolation"] = daily.snowc_mean * daily.sde_mean
    daily.naive_snow_isolation.attrs = {"long_name": "Naive snow isolation"}
    _store("naive_snow_isolation")


def _load_daily():
    daily_store = get_era5_stores("daily")
    daily = xr.open_zarr(daily_store, consolidated=False, chunks={"latitude": -1, "longitude": 640})
    daily = daily.set_coords("spatial_ref")
    assert "time" in daily.dims, f"Expected dim 'time' to be in {daily.dims=}"
    return daily


def monthly_aggregate():
    """Aggregate enriched daily ERA5 data to monthly resolution.

    Takes the enriched daily ERA5 data and creates monthly aggregates using
    appropriate statistical functions for each variable type.
    Instant variables use min, max, or median aggregations, while accumulative
    variables are summed over the month.

    The aggregated monthly data is saved to a zarr file for further processing.

    """
    daily = _load_daily()
    daily = daily.sel(time=slice(min_time, max_time))

    monthly = (
        xr.Dataset(
            {},
            coords={
                "time": daily.time.resample(time="1MS").last(),
                "latitude": daily.latitude,
                "longitude": daily.longitude,
            },
            attrs=daily.attrs,
        )
        .odc.assign_crs("epsg:4326")
        .chunk({"time": 12, "latitude": -1, "longitude": 640})
    )

    monthly_store = get_era5_stores("monthly")
    print(f"Saving empty monthly ERA5 data to {monthly_store}.")
    monthly.to_zarr(monthly_store, mode="w", encoding=codecs.from_ds(monthly), consolidated=False)

    # Monthly instant aggregates
    for var in instants | accums:
        if var in accums:
            monthly[var] = daily[var].resample(time="1MS").sum(skipna=False).rename(var)
        elif var.endswith("_min"):
            monthly[var] = daily[var].resample(time="1MS").min().rename(var)
        elif var.endswith("_max"):
            monthly[var] = daily[var].resample(time="1MS").max().rename(var)
        else:
            monthly[var] = daily[var].resample(time="1MS").median().rename(var)

        monthly[var].attrs = daily[var].attrs
        if var in daily or var.endswith(("_min", "_max")):
            monthly[var].attrs["long_name"] = monthly[var].attrs.get("long_name", "").replace("Daily", "Monthly")
        else:
            monthly[var].attrs["long_name"] = f"Monthly median of {monthly[var].attrs.get('long_name', var)}"

        # Similar storing logic as in daily
        encoding = codecs.from_ds(monthly[[var]], include_coords=False)
        print(f"Storing enriched monthly variable {var} to {monthly_store}...")
        with stopwatch("Storing enriched monthly variable"):
            monthly[[var]].chunk({"time": 12, "latitude": -1, "longitude": 640}).to_zarr(
                monthly_store, mode="a", encoding=encoding, consolidated=False
            )
        monthly = xr.open_zarr(monthly_store, consolidated=False).set_coords("spatial_ref")


def multi_monthly_aggregate(agg: Literal["yearly", "seasonal", "shoulder"] = "yearly"):
    """Aggregate monthly ERA5 data to a multi-month resolution.

    Takes monthly aggregated data and creates multi-month aggregates using a shifted
    calendar (October to September) to better capture Arctic seasonal patterns.

    Args:
        agg (Literal["yearly", "seasonal", "shoulder"], optional): The type of multi-month aggregation.


    """
    n = {"yearly": 12, "seasonal": 6, "shoulder": 3}[agg]
    daily = _load_daily()

    # "Shift" the calendar by slicing the first Jan-Sep and the last Oct-Dec months
    first_year = daily.time.dt.year.min().item()
    last_year = daily.time.dt.year.max().item()
    daily = daily.sel(time=slice(f"{first_year}-10-01", f"{last_year}-09-30"))

    chunks = {"time": 3, "latitude": -1, "longitude": 640}
    multimonthly = (
        xr.Dataset(
            {},
            coords={
                "time": daily.resample(time=f"{n}MS").last().time,
                "latitude": daily.latitude,
                "longitude": daily.longitude,
            },
            attrs=daily.attrs,
        )
        .odc.assign_crs("epsg:4326")
        .chunk(chunks)
    )

    multimonthly_store = get_era5_stores(agg)
    print(f"Saving empty multi-monthly ERA5 data to {multimonthly_store}.")
    multimonthly.to_zarr(
        multimonthly_store,
        mode="w",
        encoding=codecs.from_ds(multimonthly),
        consolidated=False,
    )

    def _store(var):
        nonlocal multimonthly

        # Similar storing logic as in daily
        encoding = codecs.from_ds(multimonthly[[var]], include_coords=False)
        print(f"Storing enriched multi-monthly variable {var} to {multimonthly_store}...")
        with stopwatch("Storing enriched multi-monthly variable"):
            multimonthly[[var]].chunk(chunks).to_zarr(
                multimonthly_store, mode="a", encoding=encoding, consolidated=False
            )
        multimonthly = xr.open_zarr(multimonthly_store, consolidated=False).set_coords("spatial_ref")

    # Multi-Monthly instant aggregates
    for var in instants | accums:
        if var in accums:
            multimonthly[var] = daily[var].resample(time=f"{n}MS").sum(skipna=False).rename(var)
        elif var.endswith("_min"):
            multimonthly[var] = daily[var].resample(time=f"{n}MS").min().rename(var)
        elif var.endswith("_max"):
            multimonthly[var] = daily[var].resample(time=f"{n}MS").max().rename(var)
        else:
            multimonthly[var] = daily[var].resample(time=f"{n}MS").median().rename(var)

        multimonthly[var].attrs = daily[var].attrs
        if var in accums or var.endswith(("_min", "_max")):
            multimonthly[var].attrs["long_name"] = (
                multimonthly[var].attrs.get("long_name", "").replace("Daily", agg.capitalize())
            )
        else:
            multimonthly[var].attrs["long_name"] = (
                f"{agg.capitalize()} median of {multimonthly[var].attrs.get('long_name', var)}"
            )
        _store(var)

    # Effective snow depth
    m = (
        daily.time.resample(time=f"{n}MS")
        .map(lambda x: np.arange(1, x.count() + 1))
        .assign_coords({"time": daily.time})
    )
    nd = (
        daily.time.resample(time=f"{n}MS")
        .map(lambda x: np.ones(x.count().item()) * x.count().item())
        .assign_coords({"time": daily.time})
    )
    n_sum = (nd * (nd + 1) // 2).resample(time=f"{n}MS").sum()
    efd = (daily.sde_mean * (nd + 1 - m)).resample(time=f"{n}MS").sum(skipna=False) / n_sum
    multimonthly["effective_snow_depth"] = efd.rename("effective_snow_depth")
    multimonthly["effective_snow_depth"].attrs = {
        "long_name": "Effective Snow Density",
        "reference": "Slater et. al. (2017)",
        "link": "https://tc.copernicus.org/articles/11/989/2017/tc-11-989-2017.pdf",
    }
    _store("effective_snow_depth")


def _get_first_day(x):
    return x.argmax(dim="time") + 1


def _store_in_yearly(da, var):
    yearly_store = get_era5_stores("yearly")
    yearly = xr.open_zarr(yearly_store, consolidated=False).set_coords("spatial_ref")
    assert "time" in yearly.dims, f"Expected dim 'time' to be in {yearly.dims=}"

    da["time"] = yearly.time
    yearly[var] = da.rename(var)

    # Similar storing logic as in daily
    encoding = codecs.from_ds(yearly[[var]], include_coords=False)
    chunks = {"time": 20, "latitude": -1, "longitude": 640}
    print(f"Storing enriched yearly variable {var} to {yearly_store}...")
    with stopwatch("Storing enriched yearly variable"):
        yearly[[var]].chunk(chunks).to_zarr(yearly_store, mode="a", encoding=encoding, consolidated=False)


def yearly_thaw_periods():
    daily = _load_daily()
    first_year = daily.time.dt.year.min().item() + 1
    last_year = daily.time.dt.year.max().item()
    daily = daily.sel(time=slice(f"{first_year}-01-01", f"{last_year}-12-31"))

    never_thaws = (daily.thawing_days.resample(time="12MS").sum(dim="time") == 0).persist()

    first_thaw_day = daily.thawing_days.resample(time="12MS").map(_get_first_day).where(~never_thaws)
    first_thaw_day.attrs = {
        "long_name": "Day of first thaw in year",
        "units": "day of year",
    }
    _store_in_yearly(first_thaw_day, "day_of_first_thaw")

    n_days_in_year = xr.where(366, 365, daily.time[::-1].dt.is_leap_year).resample(time="12MS").first()
    last_thaw_day = n_days_in_year - daily.thawing_days[::-1].resample(time="12MS").map(_get_first_day).where(
        ~never_thaws
    )
    last_thaw_day.attrs = {
        "long_name": "Day of last thaw in year",
        "units": "day of year",
    }
    _store_in_yearly(last_thaw_day, "day_of_last_thaw")


def yearly_freeze_periods():
    daily = _load_daily()
    first_year = daily.time.dt.year.min().item()
    last_year = daily.time.dt.year.max().item() - 1
    daily = daily.sel(time=slice(f"{first_year}-07-01", f"{last_year}-06-30"))

    never_freezes = (daily.freezing_days.resample(time="12MS").sum(dim="time") == 0).persist()

    first_freeze_day = daily.freezing_days.resample(time="12MS").map(_get_first_day).where(~never_freezes)
    first_freeze_day.attrs = {
        "long_name": "Day of first freeze in year",
        "units": "day of year",
    }
    _store_in_yearly(first_freeze_day, "day_of_first_freeze")

    n_days_in_year = xr.where(366, 365, daily.time[::-1].dt.is_leap_year).resample(time="12MS").last()
    last_freeze_day = n_days_in_year - daily.freezing_days[::-1].resample(time="12MS").map(_get_first_day).where(
        ~never_freezes
    )
    last_freeze_day.attrs = {
        "long_name": "Day of last freeze in year",
        "units": "day of year",
    }
    _store_in_yearly(last_freeze_day, "day_of_last_freeze")


@cli.command
def enrich(n_workers: int = 10, monthly: bool = True, yearly: bool = True, daily: bool = True):
    """Enrich data and pPerform temporal aggregation of ERA5 data using Dask cluster.

    Creates a Dask cluster and runs both monthly and yearly aggregation
    functions to generate temporally aggregated climate datasets. The
    processing uses parallel workers for efficient computation.

    Args:
        n_workers (int, optional): Number of Dask workers to use. Defaults to 10.
        monthly (bool, optional): Whether to perform monthly aggregation. Defaults to True.
        yearly (bool, optional): Whether to perform yearly aggregation. Defaults to True.
        daily (bool, optional): Whether to perform daily enrichment. Defaults to True.

    """
    with (
        dd.LocalCluster(n_workers=n_workers, threads_per_worker=20, memory_limit="10GB") as cluster,
        dd.Client(cluster) as client,
    ):
        print(client)
        print(client.dashboard_link)
        if daily:
            daily_enrich()
        if monthly:
            monthly_aggregate()
        if yearly:
            multi_monthly_aggregate(agg="shoulder")
            multi_monthly_aggregate(agg="seasonal")
            multi_monthly_aggregate(agg="yearly")
            # yearly_thaw_periods()
            # yearly_freeze_periods()
        print("Enriched ERA5 data with additional features and aggregated it temporally.")
        stopwatch.summary()


@cli.command
def viz(
    grid: Grid | None = None,
    level: int | None = None,
    agg: Literal["daily", "monthly", "yearly", "summer", "winter", "seasonal", "shoulder"] = "monthly",
    high_qual: bool = False,
):
    """Visualize a small overview of ERA5 variables for a given aggregation.

    Args:
        grid (Grid, optional): Grid type for spatial representation.
            If provided along with level, the ERA5 data will be decoded onto the specified grid.
        level (int, optional): Level of the grid for spatial representation.
            If provided along with grid, the ERA5 data will be decoded onto the specified grid.
        agg (Literal["daily", "monthly", "yearly", "summer", "winter", "seasonal", "shoulder"], optional):
            Aggregation identifier used to locate the appropriate ERA5 Zarr store via
            get_era5_stores. Determines which dataset is opened and visualized.
        high_qual (bool, optional): Whether to use high quality plotting settings.
            If True, the plot will be generated with higher resolution and quality settings.
            Defaults to False.

    Example:
        >>> # produce and save an overview for monthly ERA5 data
        >>> viz("monthly")

    """
    is_grid = grid is not None and level is not None
    store = get_era5_stores(agg, grid, level)
    ds = xr.open_zarr(store, consolidated=False).set_coords("spatial_ref")
    if is_grid:
        ds = xdggs.decode(ds)
        # Get cell boundaries for plotting
        ds = to_xvec(ds)

    tis = [0, 1, -2, -1]
    ts = [str(ds.time.isel(time=t).values)[:10] for t in tis]
    vnames = [ds[var].attrs.get("long_name", var) for var in ds.data_vars]
    vunits = [ds[var].attrs.get("units", "") for var in ds.data_vars]
    vlabels = [f"{name} [{unit}]" if unit else name for name, unit in zip(vnames, vunits)]

    proj = uplt.Proj("npaeqd")
    fig, axs = uplt.subplots(ncols=len(tis), nrows=len(vlabels), proj=proj, share=0)
    axs.format(boundinglat=50, coast=True, toplabels=ts, leftlabels=vlabels)
    for t in tis:
        for i, var in enumerate(ds.data_vars):
            subset = ds[var].isel(time=t)
            if is_grid:
                subset.xvec.to_geodataframe().to_crs(proj.proj4_init).plot(ax=axs[i, t], column=var)
            else:
                if not high_qual:
                    subset = subset.isel(latitude=slice(None, None, 4), longitude=slice(None, None, 4))
                m = axs[i, t].pcolormesh(subset, cmap="viridis")
                axs[i, t].colorbar(m, loc="ll", label=var)
    fig.format(suptitle="ERA5 Data")
    figdir = (FIGURES_DIR / "era5").resolve()
    figdir.mkdir(parents=True, exist_ok=True)
    if high_qual:
        print(f"Saving ERA5 {agg} overview figure to {figdir / f'{agg}_overview_unaligned.pdf'}.")
        fig.savefig(figdir / f"{agg}_overview_unaligned.pdf")
    else:
        print(f"Saving ERA5 {agg} overview figure to {figdir / f'{agg}_overview_unaligned.png'}.")
        fig.savefig(figdir / f"{agg}_overview_unaligned.png", dpi=100)


# ===========================
# === Spatial Aggregation ===
# ===========================


@stopwatch("Correcting longitudes to -180 to 180")
def _correct_longs(ds: xr.Dataset) -> xr.Dataset:
    return ds.assign_coords(longitude=(((ds.longitude + 180) % 360) - 180)).sortby("longitude")


@cli.command
def spatial_agg(
    grid: Grid,
    level: int,
    concurrent_partitions: int = 20,
):
    """Perform spatial aggregation of ERA5 data to grid cells.

    Loads a grid and spatially aggregates ERA5 data to each grid cell using
    parallel processing. Creates an empty zarr file first, then fills it
    with extracted data for each cell.

    Args:
        grid ("hex" | "healpix"): Grid type.
        level (int): Grid resolution level.
        concurrent_partitions (int, optional): Number of concurrent partitions to process.
            Defaults to 20.

    """
    with stopwatch(f"Loading {grid} grid at level {level}"):
        grid_gdf = grids.open(grid, level)
        # ? Mask out water, since we don't want to aggregate over oceans
        grid_gdf = watermask.clip_grid(grid_gdf)
        grid_gdf = grid_gdf.to_crs("epsg:4326")
        # There is a small group of invalid geometry in healpix 10, filter it out manually
        # They are invalid because the watermask clipping does not handle the antimeridian well
        if grid == "healpix" and level == 10:
            invalid_cell_id = [3059646, 3063547]
            grid_gdf = grid_gdf[~grid_gdf.cell_id.isin(invalid_cell_id)]

    aggregations_by_gridlevel: dict[str, dict[int, _Aggregations | Literal["interpolate"]]] = {
        "hex": {
            3: _Aggregations.common(),
            4: _Aggregations.common(),
            5: _Aggregations(mean=True),
            6: "interpolate",  # nearest neighbor interpolation
        },
        "healpix": {
            6: _Aggregations.common(),
            7: _Aggregations.common(),
            8: _Aggregations(mean=True),
            9: _Aggregations(mean=True),
            10: "interpolate",  # nearest neighbor interpolation
        },
    }
    aggregations = aggregations_by_gridlevel[grid][level]

    for agg in ("yearly", "seasonal", "shoulder"):
        unaligned_store = get_era5_stores(agg)
        with stopwatch(f"Loading {agg} ERA5 data"):
            unaligned = xr.open_zarr(unaligned_store, consolidated=False).set_coords("spatial_ref").load()
            assert {"latitude", "longitude", "time"} == set(unaligned.dims)
            assert unaligned.odc.crs == "epsg:4326", f"Expected CRS 'epsg:4326', got {unaligned.odc.crs}"
            unaligned = _correct_longs(unaligned)

        # Filter out Grid Cells that are completely outside the ERA5 valid area
        valid_geoms = []
        for g, v in shapes(
            unaligned.t2m_mean.isel(time=0).isnull().astype("uint8").values,
            transform=unaligned.odc.transform,
        ):
            if v == 0:
                valid_geoms.append(shapely.geometry.shape(g))
        grid_gdf_filtered = grid_gdf[grid_gdf.geometry.intersects(shapely.geometry.MultiPolygon(valid_geoms))]

        aggregated = aggregate_raster_into_grid(
            unaligned,
            grid_gdf_filtered,
            aggregations,
            grid,
            level,
            n_partitions=40,
            concurrent_partitions=concurrent_partitions,
            pxbuffer=10,
        )

        # These are artifacts from previous processing steps, drop them
        if "latitude" in aggregated.coords and "longitude" in aggregated.coords:
            aggregated = aggregated.drop_vars(["latitude", "longitude"])

        aggregated = aggregated.chunk(
            {
                "cell_ids": min(len(aggregated.cell_ids), 10000),
                "time": len(aggregated.time),
            }
        )
        store = get_era5_stores(agg, grid, level)
        with stopwatch(f"Saving spatially aggregated {agg} ERA5 data to {store}"):
            aggregated.to_zarr(store, mode="w", consolidated=False, encoding=codecs.from_ds(aggregated))
        print(f"Finished spatial matching for {agg} data.")

    print("### Stopwatch Summary ###")
    stopwatch.summary()


# ===========================
# === Spatial Aggregation ===
# ===========================


def unstack_era5_time(era5: xr.Dataset, aggregation: Literal["yearly", "seasonal", "shoulder"]) -> xr.Dataset:
    """Unstacks the time dimension of an ERA5 dataset into year and month/season dimensions.

    Args:
        era5 (xr.Dataset): The ERA5 dataset with a time dimension to be unstack
        aggregation (Literal["yearly", "seasonal", "shoulder"]): The type of aggregation to perform.
            - "yearly": No unstacking, just renames time to year.
            - "seasonal": Unstacks into year and season (winter/summer).
            - "shoulder": Unstacks into year and shoulder season (OND, JFM, AMJ, JAS).

    """
    # In the yearly case, no unstacking is necessary, we can just rename the time dimension to year and change the coord
    if aggregation == "yearly":
        era5 = era5.rename({"time": "year"})
        era5.coords["year"] = era5["year"].dt.year
        return era5

    # Make the time index a MultiIndex of year and month
    era5.coords["year"] = era5.time.dt.year
    era5.coords["month"] = era5.time.dt.month
    era5["time"] = pd.MultiIndex.from_arrays(
        [
            era5.time.dt.year.values,  # noqa: PD011
            era5.time.dt.month.values,  # noqa: PD011
        ],
        names=("year", "month"),
    )
    era5 = era5.unstack("time")  # noqa: PD010
    seasons = {10: "winter", 4: "summer"}
    shoulder_seasons = {10: "OND", 1: "JFM", 4: "AMJ", 7: "JAS"}
    month_map = seasons if aggregation == "seasonal" else shoulder_seasons
    era5.coords["month"] = era5["month"].to_series().map(month_map)
    return era5


@cli.command
def compute_synopsis(
    grid: Grid,
    level: int,
):
    """Create synopsis datasets for spatially aggregated ERA5 data.

    Loads spatially aggregated ERA5 data and computes mean and trend for each variable.
    The resulting synopsis datasets are saved to new zarr stores.

    Args:
        grid (Grid): Grid type.
        level (int): Grid resolution level.

    """
    for agg in ("yearly", "seasonal", "shoulder"):
        store = get_era5_stores(agg, grid=grid, level=level)
        era5 = xr.open_zarr(store, consolidated=False).compute()

        # For the trend calculation it is necessary to unstack the time dimension
        era5 = unstack_era5_time(era5, agg)

        era5_mean = era5.mean(dim="year")
        era5_trend = era5.polyfit(dim="year", deg=1).sel(degree=1, drop=True)
        # Rename all cols from "{var}_polyfit_coefficients" to "{var}_trend"
        era5_trend = era5_trend.rename(
            {var: str(var).replace("_polyfit_coefficients", "_trend") for var in era5_trend.data_vars}
        )
        era5_synopsis = xr.merge([era5_mean, era5_trend])
        synopsis_store = get_era5_stores(agg, grid=grid, level=level, temporal="synopsis")
        encoding = codecs.from_ds(era5_synopsis)
        for var in era5_synopsis.data_vars:
            encoding[var]["chunks"] = tuple(era5_synopsis[var].shape)

        print(f"Saving ERA5 {agg} synopsis data to {synopsis_store}.")
        era5_synopsis.to_zarr(synopsis_store, mode="w", encoding=encoding, consolidated=False)


if __name__ == "__main__":
    cli()
