"""Extract satellite embeddings from Google Earth Engine and map them to a grid.

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

import warnings
from collections.abc import Generator

import cudf
import cuml.cluster
import cyclopts
import ee
import geemap
import geopandas as gpd
import numpy as np
import pandas as pd
import sklearn.cluster
import xarray as xr
import xdggs
from rich import pretty, print, traceback
from rich.progress import track

from entropice.spatial import grids
from entropice.utils import codecs
from entropice.utils.paths import get_annual_embeddings_file, get_embeddings_store
from entropice.utils.types import Grid

# Filter out the GeoDataFrame.swapaxes deprecation warning
warnings.filterwarnings("ignore", message=".*GeoDataFrame.swapaxes.*", category=FutureWarning)

pretty.install()
traceback.install()
ee.Initialize(project="ee-tobias-hoelzer")

cli = cyclopts.App(name="alpha-earth")

# 7454521782,230147807,10000000.


def _batch_grid(grid_gdf: gpd.GeoDataFrame, n_partitions: int) -> Generator[gpd.GeoDataFrame]:
    # Simple partitioning by splitting the GeoDataFrame into n_partitions parts
    centroids = pd.DataFrame({"x": grid_gdf.geometry.centroid.x, "y": grid_gdf.geometry.centroid.y})

    # use cuml and cudf if len of centroids is larger than 100000
    if len(centroids) > 100000:
        print(f"Using cuML KMeans for partitioning {len(centroids)} centroids")
        centroids_cudf = cudf.DataFrame.from_pandas(centroids)
        kmeans = cuml.cluster.KMeans(n_clusters=n_partitions, random_state=42)
        labels = kmeans.fit_predict(centroids_cudf).to_pandas().to_numpy()
    else:
        labels = sklearn.cluster.KMeans(n_clusters=n_partitions, random_state=42).fit_predict(centroids)
    for i in range(n_partitions):
        yield grid_gdf.iloc[np.where(labels == i)[0]]


@cli.command()
def download(grid: Grid, level: int):
    """Extract satellite embeddings from Google Earth Engine and map them to a grid.

    Args:
        grid (Grid): The grid type to use.
        level (int): The grid level to use.

    """
    grid_gdf = grids.open(grid, level)

    # Hardcoded scale factors, depending on grid and level, so that approx. 10000px are sampled per grid cell
    scale_factors = {
        "hex": {3: 1600, 4: 600, 5: 240, 6: 90},
        "healpix": {6: 1600, 7: 800, 8: 400, 9: 200, 10: 100},
    }
    scale_factor = scale_factors[grid][level]
    print(f"Using scale factor of {scale_factor} for grid {grid} at level {level}.")

    # 2024-2025 for hex-6
    for year in track(range(2024, 2025), total=1, description="Processing years..."):
        embedding_collection = ee.ImageCollection("GOOGLE/SATELLITE_EMBEDDING/V1/ANNUAL")
        embedding_collection = embedding_collection.filterDate(f"{year}-01-01", f"{year}-12-31")
        aggs = [
            "mean",
            "stdDev",
            "min",
            "max",
            "count",
            "median",
            "p1",
            "p5",
            "p25",
            "p75",
            "p95",
            "p99",
        ]
        bands = [f"A{str(i).zfill(2)}_{agg}" for i in range(64) for agg in aggs]

        def extract_embedding(feature):
            # Filter collection by geometry
            geom = feature.geometry()
            embedding = embedding_collection.filterBounds(geom).mosaic()
            # Get mean embedding value for the geometry
            mean_dict = embedding.reduceRegion(
                reducer=ee.Reducer.mean()
                .combine(ee.Reducer.stdDev(), sharedInputs=True)
                .combine(ee.Reducer.minMax(), sharedInputs=True)
                .combine(ee.Reducer.count(), sharedInputs=True)
                .combine(ee.Reducer.median(), sharedInputs=True)
                .combine(ee.Reducer.percentile([1, 5, 25, 75, 95, 99]), sharedInputs=True),
                geometry=geom,
                scale=scale_factor,
            )
            # Add mean embedding values as properties to the feature
            return feature.set(mean_dict)

        # Process grid in batches of 100
        batch_size = 25
        all_results = []
        n_batches = len(grid_gdf) // batch_size
        for batch_grid in track(
            _batch_grid(grid_gdf, n_batches),
            description="Processing batches...",
            total=n_batches,
        ):
            # Convert batch to EE FeatureCollection
            eegrid_batch = ee.FeatureCollection(batch_grid.to_crs("epsg:4326").__geo_interface__)

            # Apply embedding extraction to batch
            eeegrid_batch = eegrid_batch.map(extract_embedding)
            df_batch = geemap.ee_to_df(eeegrid_batch)

            # Store batch results
            all_results.append(df_batch)

        # Combine all batch results
        df = pd.concat(all_results, ignore_index=True)
        embeddings_on_grid = grid_gdf.merge(df[[*bands, "cell_id"]], on="cell_id", how="left")
        embeddings_file = get_annual_embeddings_file(grid, level, year)
        embeddings_on_grid.to_parquet(embeddings_file)
        print(f"Saved embeddings for year {year} to {embeddings_file}.")


@cli.command()
def combine_to_zarr(grid: Grid, level: int):
    """Combine yearly embeddings parquet files into a single zarr store.

    Args:
        grid (Grid): The grid type to use.
        level (int): The grid level to use.

    """
    cell_ids = grids.get_cell_ids(grid, level)
    years = list(range(2018, 2025))
    aggs = [
        "mean",
        "stdDev",
        "min",
        "max",
        "count",
        "median",
        "p1",
        "p5",
        "p25",
        "p75",
        "p95",
        "p99",
    ]
    bands = [f"A{str(i).zfill(2)}" for i in range(64)]

    a = xr.DataArray(
        np.nan,
        dims=("year", "cell_ids", "band", "agg"),
        coords={"year": years, "cell_ids": cell_ids, "band": bands, "agg": aggs},
        name="embeddings",
    ).astype(np.float32)

    # ? These attributes are needed for xdggs
    a.cell_ids.attrs = {
        "grid_name": "h3" if grid == "hex" else "healpix",
        "level": level,
    }
    if grid == "healpix":
        a.cell_ids.attrs["indexing_scheme"] = "nested"

    for year in track(years, total=len(years), description="Processing years..."):
        embeddings_file = get_annual_embeddings_file(grid, level, year)
        embs = gpd.read_parquet(embeddings_file)
        for band in bands:
            for agg in aggs:
                col = f"{band}_{agg}"
                a.loc[{"band": band, "agg": agg, "year": year}] = embs[col].to_list()

    a = xdggs.decode(a)

    a = a.chunk({"year": 1, "cell_ids": 100000, "band": 64, "agg": 12})
    zarr_path = get_embeddings_store(grid, level)
    a.to_zarr(zarr_path, consolidated=False, mode="w", encoding=codecs.from_da(a))
    print(f"Saved combined embeddings to {zarr_path}.")


def _safe_polyfit(embeddings: xr.Dataset) -> xr.Dataset:
    batch_size = 1000
    embedding_trends = []
    for i in track(list(range(0, embeddings.sizes["cell_ids"], batch_size))):
        embeddings_islice = embeddings.isel(cell_ids=slice(i, i + batch_size))
        try:
            embeddings_trend = embeddings_islice.polyfit(dim="year", deg=1, skipna=True).sel(degree=1, drop=True)
            embedding_trends.append(embeddings_trend)
        except Exception as e:
            print(f"Error processing embeddings {i} to {i + batch_size}: {e}")
    return xr.concat(embedding_trends, dim="cell_ids")


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

    Loads spatially aggregated AlphaEarth embeddings 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.

    """
    store = get_embeddings_store(grid=grid, level=level)
    embeddings = xr.open_zarr(store, consolidated=False).compute()

    embeddings_mean = embeddings.mean(dim="year")
    print(f"{embeddings.sizes=} -> {embeddings.nbytes * 1e-9:.2f} GB")

    if (grid == "hex" and level >= 6) or (grid == "healpix" and level >= 10):
        # We need to process the trend in chunks, because very large arrays can lead to
        # numerical issues in the least squares solver used by polyfit.
        embeddings_trend = _safe_polyfit(embeddings)
    else:
        embeddings_trend = embeddings.polyfit(dim="year", deg=1, skipna=True).sel(degree=1, drop=True)

    # Rename all cols from "{var}_polyfit_coefficients" to "{var}_trend"
    embeddings_trend = embeddings_trend.rename(
        {var: str(var).replace("_polyfit_coefficients", "_trend") for var in embeddings_trend.data_vars}
    )
    embeddings_synopsis = xr.merge([embeddings_mean, embeddings_trend])
    synopsis_store = get_embeddings_store(grid=grid, level=level, temporal="synopsis")
    encoding = codecs.from_ds(embeddings_synopsis)
    for var in embeddings_synopsis.data_vars:
        encoding[var]["chunks"] = (min(v, 100000) for v in embeddings_synopsis[var].shape)
    print(f"Saving AlphaEarth embeddings synopsis data to {synopsis_store}.")
    embeddings_synopsis.to_zarr(synopsis_store, mode="w", encoding=encoding, consolidated=False)


def main():  # noqa: D103
    cli()


if __name__ == "__main__":
    main()
