import datetime
import multiprocessing as mp
from dataclasses import dataclass
from math import ceil

import cupy as cp
import cupy_xarray
import cupyx.scipy.signal
import cyclopts
import dask.array
import dask.distributed as dd
import geopandas as gpd
import icechunk
import icechunk.xarray
import numpy as np
import smart_geocubes
import xarray as xr
import xdggs
import xrspatial
import xrspatial.convolution
import zarr
from cupyx.scipy.ndimage import binary_dilation, binary_erosion, distance_transform_edt
from rich import pretty, print, traceback
from stopuhr import stopwatch
from xrspatial.aspect import _run_cupy as aspect_cupy
from xrspatial.curvature import _run_cupy as curvature_cupy
from xrspatial.slope import _run_cupy as slope_cupy
from zarr.codecs import BloscCodec

from entropice.spatial import grids, watermask
from entropice.spatial.aggregators import _Aggregations, aggregate_raster_into_grid
from entropice.utils import codecs
from entropice.utils.paths import get_arcticdem_stores
from entropice.utils.types import Grid

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

zarr.config.set({"async.concurrency": 128})

cli = cyclopts.App(name="arcticdem")


@cli.command()
def download():
    arcticdem_store = get_arcticdem_stores()
    adem = smart_geocubes.ArcticDEM32m(arcticdem_store)
    adem.create(exists_ok=True)
    with stopwatch("Download ArcticDEM data"):
        adem.procedural_download(adem.extent.extent, None)


@dataclass
class _KernelFactory:
    res: int
    size_px: int

    @property
    def size(self):
        return self.size_px * self.res

    @property
    def inner(self):
        return ceil(self.size / 3)

    @staticmethod
    def _to_cupy_f32(kernel):
        return cp.asarray(kernel.astype("float32"))

    def ring(self):
        return self._to_cupy_f32(xrspatial.convolution.annulus_kernel(self.res, self.res, self.size, self.inner))

    def tri(self):
        kernel = np.ones((self.size_px, self.size_px), dtype=float)
        kernel[self.size_px // 2, self.size_px // 2] = 0  # Set the center cell to 0
        return self._to_cupy_f32(xrspatial.convolution.custom_kernel(kernel))

    def vrm(self):
        kernel = np.ones((self.size_px, self.size_px), dtype=float) / self.size_px**2
        return self._to_cupy_f32(xrspatial.convolution.custom_kernel(kernel))


def tpi_cupy(chunk, kernels: _KernelFactory):
    kernel = kernels.ring()
    kernel = kernel / cp.nansum(kernel)
    tpi = chunk - cupyx.scipy.signal.convolve2d(chunk, kernel, mode="same")
    return tpi


def tri_cupy(chunk, kernels: _KernelFactory):
    kernel = kernels.tri()
    c2 = chunk**2
    focal_sum = cupyx.scipy.signal.convolve2d(chunk, kernel, mode="same")
    focal_sum_squared = cupyx.scipy.signal.convolve2d(c2, kernel, mode="same")
    tri = np.sqrt((kernel.size - 1) * c2 - 2 * chunk * focal_sum + focal_sum_squared)
    return tri


def ruggedness_cupy(chunk, slope, aspect, kernels: _KernelFactory):
    slope_rad = slope * (cp.pi / 180)
    aspect_rad = aspect * (cp.pi / 180)
    aspect_rad = cp.where(aspect_rad == -1, 0, aspect_rad)

    xy = cp.sin(slope_rad)
    z = cp.cos(slope_rad)
    x = xy * cp.sin(aspect_rad)
    y = xy * cp.cos(aspect_rad)

    kernel = kernels.vrm()

    # Calculate sums of x, y, and z components in the neighborhood
    x_sum = cupyx.scipy.signal.convolve2d(x, kernel, mode="same")
    y_sum = cupyx.scipy.signal.convolve2d(y, kernel, mode="same")
    z_sum = cupyx.scipy.signal.convolve2d(z, kernel, mode="same")

    # Calculate the resultant vector magnitude
    vrm = 1 - np.sqrt(x_sum**2 + y_sum**2 + z_sum**2)
    return vrm


def _get_xy_chunk(chunk: np.ndarray, x: np.ndarray, y: np.ndarray, block_info=None) -> tuple[cp.ndarray, cp.ndarray]:
    assert isinstance(block_info, list) and len(block_info) >= 1
    chunk_loc = block_info[0]["chunk-location"]
    d = 15
    cs = 3600

    # Calculate safe slice bounds for edge chunks
    y_start, y_end = (
        max(0, cs * chunk_loc[0] - d),
        min(len(y), cs * chunk_loc[0] + cs + d),
    )
    x_start, x_end = (
        max(0, cs * chunk_loc[1] - d),
        min(len(x), cs * chunk_loc[1] + cs + d),
    )

    # Extract coordinate arrays with safe bounds
    y_chunk = cp.asarray(y[y_start:y_end])
    x_chunk = cp.asarray(x[x_start:x_end])

    # Handle cases where the extracted chunk doesn't match expected dimensions
    if len(y_chunk) != chunk.shape[0] or len(x_chunk) != chunk.shape[1]:
        print(
            f"Adjusting coordinate chunk sizes: y_chunk {len(y_chunk)} vs chunk {chunk.shape[0]}, "
            f"x_chunk {len(x_chunk)} vs chunk {chunk.shape[1]}"
        )
        # Pad coordinates to match chunk dimensions if needed
        if len(y_chunk) < chunk.shape[0]:
            # Pad with the edge values
            pad_start = chunk.shape[0] - len(y_chunk)
            y_chunk = cp.pad(y_chunk, (pad_start, 0), mode="edge")
        if len(x_chunk) < chunk.shape[1]:
            pad_start = chunk.shape[1] - len(x_chunk)
            x_chunk = cp.pad(x_chunk, (pad_start, 0), mode="edge")

    yy = cp.repeat(y_chunk.reshape(-1, 1), x_chunk.shape[0], axis=1)
    xx = cp.repeat(x_chunk.reshape(1, -1), y_chunk.shape[0], axis=0)
    return xx, yy


def _enrich_chunk(chunk: np.ndarray, x: np.ndarray, y: np.ndarray, block_info=None) -> np.ndarray:
    res = 32  # 32m resolution
    small_kernels = _KernelFactory(res=res, size_px=3)  # ~3x3 kernels (96m)
    medium_kernels = _KernelFactory(res=res, size_px=7)  # ~7x7 kernels (224m)
    large_kernels = _KernelFactory(res=res, size_px=15)  # ~15x15 kernels (480m)

    # Check if there is data in the chunk
    if np.all(np.isnan(chunk)):
        # Return an array of NaNs with the expected shape
        return np.full(
            (
                12,
                chunk.shape[0] - 2 * large_kernels.size_px,
                chunk.shape[1] - 2 * large_kernels.size_px,
            ),
            np.nan,
            dtype=np.float32,
        )

    chunk = cp.asarray(chunk)

    # Interpolate missing values in chunk (for patches smaller than 7x7 pixels)
    mask = cp.isnan(chunk)
    mask &= ~binary_dilation(
        binary_erosion(mask, iterations=3, brute_force=True),
        iterations=3,
        brute_force=True,
    )
    if cp.any(mask):
        # Find indices of valid values
        indices = distance_transform_edt(mask, return_distances=False, return_indices=True)
        chunk = chunk[tuple(indices)]

    # TPI
    tpi_small = tpi_cupy(chunk, small_kernels)
    tpi_medium = tpi_cupy(chunk, medium_kernels)
    tpi_large = tpi_cupy(chunk, large_kernels)
    # Slope
    slope = slope_cupy(chunk, res, res)
    # Aspect & correction
    aspect = aspect_cupy(chunk)
    xx, yy = _get_xy_chunk(chunk, x, y, block_info)
    aspect_correction = cp.arctan2(yy, xx) * (180 / cp.pi) + 90
    aspect = (aspect + aspect_correction) % 360
    # Curvature
    curvature = curvature_cupy(chunk, res)
    # TRI
    tri_small = tri_cupy(chunk, small_kernels)
    tri_medium = tri_cupy(chunk, medium_kernels)
    tri_large = tri_cupy(chunk, large_kernels)
    # Ruggedness
    vrm_small = ruggedness_cupy(chunk, slope, aspect, small_kernels)
    vrm_medium = ruggedness_cupy(chunk, slope, aspect, medium_kernels)
    vrm_large = ruggedness_cupy(chunk, slope, aspect, large_kernels)

    # Stack on GPU, then move to CPU once
    res_gpu = cp.stack(
        [
            tpi_small,
            tpi_medium,
            tpi_large,
            slope,
            aspect,
            curvature,
            tri_small,
            tri_medium,
            tri_large,
            vrm_small,
            vrm_medium,
            vrm_large,
        ],
        axis=0,
    )

    # Trim edges on GPU
    e = large_kernels.size_px
    res_gpu = res_gpu[:, e:-e, e:-e]

    # Single CPU transfer at the end
    res = cp.asnumpy(res_gpu)
    return res


def _enrich(terrain: xr.DataArray):
    features = [
        "tpi_small",
        "tpi_medium",
        "tpi_large",
        "slope",
        "aspect",
        "curvature",
        "tri_small",
        "tri_medium",
        "tri_large",
        "vrm_small",
        "vrm_medium",
        "vrm_large",
    ]
    output_chunks = (len(features), *tuple(terrain.data.chunks))
    enriched_da = dask.array.map_overlap(
        _enrich_chunk,
        terrain.data,
        x=terrain.x.to_numpy(),
        y=terrain.y.to_numpy(),
        depth=15,  # large_kernels.size_px
        chunks=output_chunks,
        new_axis=0,
        dtype=np.float32,
        meta=np.array((), dtype=np.float32),
        trim=False,
        boundary=np.nan,
        token="enrich_arcticdem",
    )
    enriched_da = xr.DataArray(
        enriched_da,
        dims=("feature", "y", "x"),
        coords={
            "feature": features,
            "y": terrain.y,
            "x": terrain.x,
        },
    )
    return enriched_da, features


@cli.command()
def enrich():
    with (
        dd.LocalCluster(n_workers=7, threads_per_worker=2, memory_limit="30GB") as cluster,
        dd.Client(cluster) as client,
    ):
        print(client)
        print(client.dashboard_link)

        arcticdem_store = get_arcticdem_stores()
        accessor = smart_geocubes.ArcticDEM32m(arcticdem_store)

        # Garbage collect from previous runs
        accessor.repo.garbage_collect(datetime.datetime.now(datetime.UTC))

        adem = accessor.open_xarray()
        del adem.y.attrs["_FillValue"]
        del adem.x.attrs["_FillValue"]
        enriched, new_features = _enrich(adem.dem)

        for feature in new_features:
            adem[feature] = enriched.sel(feature=feature)
        print(adem[new_features])
        # subset = adem[new_features].isel(x=slice(190800, 220000), y=slice(61200, 100000))

        session = accessor.repo.writable_session("main")
        encodings = {feature: {"compressors": [BloscCodec(clevel=5)]} for feature in new_features}
        icechunk.xarray.to_icechunk(
            adem[new_features],
            session,
            mode="a",
            encoding=encodings,
        )
        session.commit("Add terrain features: TPI, slope, aspect, curvature, TRI, ruggedness")

    print("Enrichment complete.")


def _open_adem():
    arcticdem_store = get_arcticdem_stores()
    accessor = smart_geocubes.ArcticDEM32m(arcticdem_store)
    # config = icechunk.RepositoryConfig.default()
    # config.caching = icechunk.CachingConfig(
    #     num_snapshot_nodes=None,
    #     num_chunk_refs=None,
    #     num_transaction_changes=None,
    #     num_bytes_attributes=None,
    #     num_bytes_chunks=None,
    # )
    # repo = icechunk.Repository.open(accessor.repo.storage, config=config)
    # session = repo.readonly_session("main")
    # adem = xr.open_zarr(session.store, mask_and_scale=False, consolidated=False).set_coords("spatial_ref")
    adem = accessor.open_xarray()
    # Don't use the datamask
    adem = adem.drop_vars("datamask")
    assert {"x", "y"} == set(adem.dims)
    assert adem.odc.crs == "EPSG:3413"
    return adem


@cli.command()
def aggregate(grid: Grid, level: int, concurrent_partitions: int = 20):
    mp.set_start_method("forkserver", force=True)
    with (
        dd.LocalCluster(n_workers=1, threads_per_worker=32, memory_limit="20GB") as cluster,
    ):
        print(cluster)
        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)
            # ? Mask out cells that do not intersect with ArcticDEM extent
            arcticdem_store = get_arcticdem_stores()
            accessor = smart_geocubes.ArcticDEM32m(arcticdem_store)
            extent_info = gpd.read_parquet(accessor._aux_dir / "ArcticDEM_Mosaic_Index_v4_1_32m.parquet")
            grid_gdf = grid_gdf[grid_gdf.intersects(extent_info.geometry.union_all())]

        aggregations = _Aggregations.common()
        aggregated = aggregate_raster_into_grid(
            _open_adem,
            grid_gdf,
            aggregations,
            grid,
            level,
            concurrent_partitions=concurrent_partitions,
            n_partitions=400,  # 400  # Results in ~10GB large patches
            pxbuffer=30,
        )
        store = get_arcticdem_stores(grid, level)
        with stopwatch(f"Saving spatially aggregated ArcticDEM data to {store}"):
            aggregated.to_zarr(store, mode="w", consolidated=False, encoding=codecs.from_ds(aggregated))

        print("Aggregation complete.")

        print("### Finished ArcticDEM processing ###")
        stopwatch.summary()


if __name__ == "__main__":
    cli()
