"""Implements build methods for the land surface submodel, responsible for land surface characteristics and processes."""

import copy
from pathlib import Path

import geopandas as gpd
import numpy as np
import xarray as xr
from pyflwdir.dem import fill_depressions

from geb.build.methods import build_method
from geb.workflows.io import get_window, parse_and_set_zarr_CRS
from geb.workflows.raster import (
    calculate_cell_area_m2,
    convert_nodata,
    interpolate_na_2d,
    interpolate_na_along_dim,
    pad_xy,
    rasterize_like,
    reclassify,
    repeat_grid,
    resample_chunked,
    resample_like,
)

from ..workflows.soilgrids import load_soilgrids
from .base import BuildModelBase


class LandSurface(BuildModelBase):
    """Implements land surface submodel, responsible for land surface characteristics and processes."""

    def __init__(self) -> None:
        """Initialize the LandSurface class."""
        pass

    @build_method(depends_on=["setup_regions_and_land_use"], required=True)
    def setup_cell_area(self) -> None:
        """Sets up the cell area map for the model.

        Notes:
            This method prepares the cell area map for the model by calculating the area of each cell in the grid. It first
            retrieves the grid mask from the `mask` attribute of the grid, and then calculates the cell area
            using the `calculate_cell_area()` function. The resulting cell area map is then set as the `cell_area`
            attribute of the grid.

            Additionally, this method sets up a subgrid for the cell area map by creating a new grid with the same extent as
            the subgrid, and then repeating the cell area values from the main grid to the subgrid using the `repeat_grid()`
            function, and correcting for the subgrid factor. Thus, every subgrid cell within a grid cell has the same value.
            The resulting subgrid cell area map is then set as the `cell_area` attribute of the subgrid.
        """
        mask = self.grid["mask"]

        cell_area = self.full_like(
            mask, fill_value=np.nan, nodata=np.nan, dtype=np.float32
        )

        height, width = cell_area.shape
        cell_area.data = calculate_cell_area_m2(
            mask.rio.transform(recalc=True), height, width
        )
        cell_area = cell_area.where(~mask, cell_area.attrs["_FillValue"])
        self.set_grid(cell_area, name="cell_area")

        sub_cell_area = self.full_like(
            self.subgrid["mask"],
            fill_value=np.nan,
            nodata=np.nan,
            dtype=np.float32,
        )

        sub_cell_area.data = (
            repeat_grid(cell_area.data, self.subgrid_factor) / self.subgrid_factor**2
        )
        self.set_subgrid(sub_cell_area, name="cell_area")

        region_subgrid_cell_area = self.full_like(
            self.region_subgrid["mask"],
            fill_value=np.nan,
            nodata=np.nan,
            dtype=np.float32,
        )

        height, width = region_subgrid_cell_area.shape
        region_subgrid_cell_area.data = calculate_cell_area_m2(
            region_subgrid_cell_area.rio.transform(recalc=True),
            height,
            width,
        )

        # set the cell area for the region subgrid
        self.set_region_subgrid(
            region_subgrid_cell_area,
            name="cell_area",
        )

    @build_method(depends_on=["setup_hydrography"], required=True)
    def setup_elevation(
        self,
        DEMs: list[dict[str, str | float]] = [
            {
                "name": "fabdem",
                "zmin": 0.001,
                "coastal_zmin": 30.0,
                "fill_depressions": True,
            },
            {
                "name": "delta_dtm",
                "zmax": 30,
                "zmin": 0.001,
                "fill_depressions": True,
                "coastal_only": True,
            },
            {
                "name": "gebco",
                "zmax": 0.0,
                "fill_depressions": False,
                "coastal_only": True,
            },
        ],
    ) -> None:
        """Sets up the elevation data for the model.

        Configuration parameters:
            name: The name of the DEM to use. Supported names are 'fabdem', 'delta_dtm', 'gebco'. If it is not supported,
                the path must be set.
            path: The path to the DEM file. Only required if the name is not supported.
            zmin: The minimum elevation where to use the DEM. Elevations below this value will be set to NaN.
            zmin_coastal: The minimum elevation where to use the DEM for coastal subbasins. Elevations below this value will be set to NaN.
            zmax: The maximum elevation where to use the DEM. Elevations above this value will be set to NaN.
            zmax_coastal: The maximum elevation where to use the DEM for coastal subbasins. Elevations above this value will be set to NaN.
            fill_depressions: Whether to fill depressions in the DEM.
            nodata: The nodata value in the DEM. Optional, only required if the DEM does not have a nodata value defined.
            crs: The CRS to set for custom DEMs when the file does not define one (EPSG code or CRS string).
            coastal_only: DEMs with this value set to True will be skipped if there are no coastal subbasins in the model.
                Default is False.


        Args:
            DEMs: A list of dictionaries containing the names and parameters of the DEMs to use. Each dictionary should
                be configured as described above.

        Raises:
            ValueError: If no DEMs are provided.
            ValueError: If the DEMs are not provided as a list of dictionaries.
            ValueError: If CRS is missing or invalid in a custom DEM.
            ValueError: If nodata value is missing in a custom DEM.
            ValueError: If DeltaDTM DEM is not provided when coastal subbasins are present.
            ValueError: If a custom DEM CRS is not a valid EPSG code or CRS string.
        """
        DEMs = copy.deepcopy(DEMs)

        if not DEMs:
            raise ValueError("At least one DEM must be provided.")

        if not isinstance(DEMs, list) or not all(isinstance(DEM, dict) for DEM in DEMs):
            raise ValueError("DEMs must be provided as a list of dictionaries.")

        if self.geom["routing/subbasins"]["is_coastal"].any():
            # deltaDTM must be present if coastal DEMs are used
            if not any(DEM.get("name", "") == "delta_dtm" for DEM in DEMs):
                raise ValueError(
                    "DeltaDTM DEM must be provided when coastal DEMs are used."
                )

            for DEM in DEMs:
                if "coastal_zmin" in DEM:
                    DEM["zmin"] = DEM["coastal_zmin"]
                if "coastal_zmax" in DEM:
                    DEM["zmax"] = DEM["coastal_zmax"]
        else:
            # Remove coastal DEMs if no coastal subbasins are present
            DEMs = [DEM for DEM in DEMs if not DEM.get("coastal_only", False)]

        # here we use the bounds of all subbasins, which may include downstream
        # subbasins that are not part of the study area
        bounds: tuple[float, float, float, float] = tuple(
            self.geom["routing/subbasins"].total_bounds
        )

        buffer: float = 0.5
        xmin: float = bounds[0] - buffer
        ymin: float = bounds[1] - buffer
        xmax: float = bounds[2] + buffer
        ymax: float = bounds[3] + buffer

        fabdem: xr.DataArray = self.data_catalog.fetch(
            "fabdem",
            xmin=xmin - 0.5,  # extra buffer to also support forcing
            xmax=xmax + 0.5,  # extra buffer to also support forcing
            ymin=ymin - 0.5,  # extra buffer to also support forcing
            ymax=ymax + 0.5,  # extra buffer to also support forcing
        ).read()

        target: xr.DataArray = self.subgrid["mask"]
        assert target.rio.crs is not None, "target grid must have a crs"

        self.set_subgrid(
            resample_like(fabdem, target, method="bilinear"),
            name="landsurface/elevation",
        )

        DEM_raster: xr.DataArray
        for DEM in DEMs:
            # FABDEM is already handled above, so we just use it from there
            if DEM["name"] == "fabdem":
                DEM_raster: xr.DataArray = fabdem

            elif DEM["name"] == "delta_dtm":
                DEM_raster: xr.DataArray = self.data_catalog.fetch(
                    "delta_dtm",
                    xmin=xmin,
                    xmax=xmax,
                    ymin=ymin,
                    ymax=ymax,
                ).read()

                # Create low elevation coastal zone mask based on DeltaDTM
                low_elevation_coastal_zone = DEM_raster < 10
                low_elevation_coastal_zone.values = (
                    low_elevation_coastal_zone.values.astype(np.float32)
                )
                self.set_other(
                    low_elevation_coastal_zone,
                    name="landsurface/low_elevation_coastal_zone",
                )  # Maybe remove this

            elif DEM["name"] == "gebco":
                DEM_raster: xr.DataArray = self.data_catalog.fetch("gebco").read()
                # set maximum values for DEM_raster if zmax is set for the DEM
                if "zmax" in DEM:
                    DEM_raster = DEM_raster.where(
                        DEM_raster <= DEM["zmax"], DEM["zmax"]
                    )

            else:
                # custom DEMs must have a path
                if "path" not in DEM:
                    raise ValueError(
                        f"DEM name '{DEM['name']}' is not supported by default. Please provide a valid path."
                    )
                if not isinstance(DEM["path"], str):
                    raise ValueError("DEM path must be a string.")

                DEM_raster: xr.DataArray = xr.open_dataarray(Path(DEM["path"]))

                # Handle CRS for custom DEMs
                # Zarrs need special handling to set the CRS
                if DEM["path"].endswith(".zarr") or DEM["path"].endswith(".zarr.zip"):
                    DEM_raster = parse_and_set_zarr_CRS(DEM_raster)

                if "crs" in DEM:
                    crs_value = DEM["crs"]
                    if not isinstance(crs_value, (int, str)):
                        raise ValueError(
                            "Custom DEM CRS must be an EPSG code (int) or CRS string."
                        )
                    DEM_raster = DEM_raster.rio.write_crs(crs_value)

                if DEM_raster.rio.crs is None:
                    raise ValueError(
                        f"DEM at path '{DEM['path']}' does not have a valid CRS."
                    )

                # Handle nodata for custom DEMs
                if "nodata" in DEM:
                    DEM_raster.attrs["_FillValue"] = DEM["nodata"]
                else:
                    if "_FillValue" not in DEM_raster.attrs:
                        raise ValueError(
                            f"DEM at path '{DEM['path']}' does not have a nodata value defined."
                        )

            if "bands" in DEM_raster.dims:
                DEM_raster: xr.DataArray = DEM_raster.isel(band=0)

            DEM_raster: xr.DataArray = DEM_raster.isel(
                get_window(
                    DEM_raster.x,
                    DEM_raster.y,
                    tuple(
                        self.geom["routing/subbasins"]
                        .to_crs(DEM_raster.rio.crs)
                        .total_bounds
                    ),
                    buffer=100,
                    raise_on_out_of_bounds=False,
                    raise_on_buffer_out_of_bounds=False,
                ),
            )

            DEM_raster = convert_nodata(
                DEM_raster.astype(np.float32, keep_attrs=True), np.nan
            )

            if DEM.get("fill_depressions", False):
                DEM_raster.values, d8 = fill_depressions(
                    DEM_raster.values, nodata=DEM_raster.attrs["_FillValue"]
                )

            self.set_other(
                DEM_raster,
                name=f"DEM/{DEM['name']}",
            )
            DEM["path"] = f"DEM/{DEM['name']}"

        self.set_params(DEMs, name="hydrodynamics/DEM_config")

    @build_method(depends_on=["setup_coastal_sfincs_model_regions"], required=True)
    def setup_regions_and_land_use(
        self,
        region_database: str = "GADM_level1",
        unique_region_id: str = "GID_1",
        ISO3_column: str = "GID_0",
        land_cover: str = "esa_worldcover_2021",
    ) -> None:
        """Sets up the (administrative) regions and land use data for GEB.

        The regions can be used for multiple purposes, for example for creating the
        agents in the model, assigning unique crop prices and other economic variables
        per region and for aggregating the results.

        Args:
            region_database: The name of the region database to use. Default is 'GADM_level1'.
            unique_region_id: The name of a column in the region database that contains a unique region ID. Default is 'UID',
                which is the unique identifier for the GADM database.
            ISO3_column: The name of a column in the region database that contains the ISO3 code for the region. Default is 'ISO3'.
            land_cover: The name of the land cover dataset to use. Default is 'esa_worldcover_2021'.

        Notes:
            This method sets up the regions and land use data for GEB. It first retrieves the region data from
            the specified region database and sets it as a geometry in the model. It then pads the subgrid to cover the entire
            region and retrieves the land use data from the ESA WorldCover dataset. The land use data is reprojected to the
            padded subgrid and the region ID is rasterized onto the subgrid. The cell area for each region is calculated and
            set as a grid in the model. The MERIT dataset is used to identify rivers, which are set as a grid in the model. The
            land use data is reclassified into five classes and set as a grid in the model. Finally, the cultivated land is
            identified and set as a grid in the model.
        """
        regions: gpd.GeoDataFrame = (
            self.data_catalog.fetch(region_database)
            .read(geom=self.region.union_all())
            .rename(columns={unique_region_id: "region_id", ISO3_column: "ISO3"})
        )

        global_countries: gpd.GeoDataFrame = (
            self.data_catalog.fetch("GADM_level0")
            .read()
            .rename(columns={"GID_0": "ISO3"})
        )

        global_countries["geometry"] = global_countries.to_crs(
            "ESRI:54009"
        ).centroid.to_crs(global_countries.crs)
        global_countries = global_countries.set_index("ISO3")

        self.set_geom(global_countries, name="global_countries")

        assert np.unique(regions["region_id"]).shape[0] == regions.shape[0], (
            f"Region database must contain unique region IDs"
        )

        region_id_mapping = {
            i: region_id for region_id, i in enumerate(regions["region_id"])
        }
        regions["region_id"] = regions["region_id"].map(region_id_mapping)

        self.set_params(region_id_mapping, name="region_id_mapping")

        assert "ISO3" in regions.columns, f"Region database must contain ISO3 column)"

        self.set_geom(regions, name="regions")

        resolution_x, resolution_y = self.subgrid["mask"].rio.resolution()

        regions_bounds: tuple[float, float, float, float] = self.geom[
            "regions"
        ].total_bounds
        mask_bounds: tuple[float, float, float, float] = self.grid["mask"].rio.bounds(
            recalc=True
        )

        # The bounds should be set to a bit larger than the regions to avoid edge effects
        # and also larger than the mask, to ensure that the entire grid is covered.
        pad_minx = min(regions_bounds[0], mask_bounds[0]) - abs(resolution_x) / 2.0
        pad_miny = min(regions_bounds[1], mask_bounds[1]) - abs(resolution_y) / 2.0
        pad_maxx = max(regions_bounds[2], mask_bounds[2]) + abs(resolution_x) / 2.0
        pad_maxy = max(regions_bounds[3], mask_bounds[3]) + abs(resolution_y) / 2.0

        region_mask, region_subgrid_slice = pad_xy(
            self.subgrid["mask"],
            pad_minx,
            pad_miny,
            pad_maxx,
            pad_maxy,
            return_slice=True,
            constant_values=1,
        )
        region_mask.attrs["_FillValue"] = None
        region_mask = self.set_region_subgrid(region_mask, name="mask")

        # Get the combined bounds of regions and subbasins to ensure coverage of both
        buffer = 0.1

        regions_bounds = self.geom["regions"].total_bounds
        subbasins_bounds = self.geom["routing/subbasins"].total_bounds

        union_geometry = (
            self.geom["regions"]
            .union_all()
            .union(self.geom["routing/subbasins"].union_all())
        )
        if len(self.geom["coastal/low_elevation_coastal_zone_mask"]) > 0:
            union_geometry = union_geometry.union(
                self.geom["coastal/low_elevation_coastal_zone_mask"].union_all()
            )
        xmin = union_geometry.bounds[0] - buffer
        ymin = union_geometry.bounds[1] - buffer
        xmax = union_geometry.bounds[2] + buffer
        ymax = union_geometry.bounds[3] + buffer

        land_use_classification_source: xr.DataArray = self.data_catalog.fetch(
            land_cover
        ).read(union_geometry)

        # we want to save the original land use data for the flood model
        # but also not save more data than necessary, as the land use classification source can be quite large.
        # therefore we clip the land use classification source to the union of the subbasins and
        # coastal low elevation zone
        flood_mask = [self.geom["routing/subbasins"].union_all()]
        if len(self.geom["coastal/low_elevation_coastal_zone_mask"]) > 0:
            flood_mask.append(
                self.geom["coastal/low_elevation_coastal_zone_mask"].union_all()
            )
        self.set_other(
            land_use_classification_source.rio.clip(
                flood_mask,
                drop=True,
                all_touched=True,
            ),
            name="landcover/classification",
        )

        land_use_classification_source_region_subgrid: xr.DataArray = resample_chunked(
            land_use_classification_source,
            region_mask,
            method="nearest",
        )

        subgrid_region_ids: xr.DataArray = rasterize_like(
            gdf=self.geom["regions"],
            column="region_id",
            raster=region_mask,
            dtype=np.int32,
            nodata=-1,
            all_touched=True,
        )
        self.set_region_subgrid(subgrid_region_ids, name="region_ids")

        land_use_classification_source_subgrid = (
            land_use_classification_source_region_subgrid.isel(region_subgrid_slice)
        ).astype(np.int16)
        land_use_classification_source_subgrid = xr.where(
            ~self.subgrid["mask"], land_use_classification_source_subgrid, -1
        )
        land_use_classes_subgrid = reclassify(
            land_use_classification_source_subgrid,
            remap_dict={
                -1: np.int8(-1),  # outside of the model domain
                0: np.int8(
                    5
                ),  # map nodata in source to permanent water bodies, as these are mostly ocean in the land cover dataset
                10: np.int8(0),  # tree cover
                20: np.int8(1),  # shrubland
                30: np.int8(1),  # grassland
                40: np.int8(
                    1
                ),  # cropland, setting to non-irrigated. Initiated as irrigated based on agents
                50: np.int8(4),  # built-up
                60: np.int8(1),  # bare / sparse vegetation
                70: np.int8(1),  # snow and ice
                80: np.int8(5),  # permanent water bodies
                90: np.int8(1),  # herbaceous wetland
                95: np.int8(5),  # mangroves
                100: np.int8(1),  # moss and lichen
            },
        )
        land_use_classes_subgrid.attrs["_FillValue"] = -1
        self.set_subgrid(land_use_classes_subgrid, name="landsurface/land_use_classes")

        cultivated_land_full_region = xr.where(
            land_use_classification_source_region_subgrid == 40,
            True,
            False,
        )

        cultivated_land_full_region.attrs["_FillValue"] = None
        self.set_region_subgrid(
            cultivated_land_full_region, name="landsurface/full_region_cultivated_land"
        )

    @build_method(depends_on=[], required=False)
    def setup_land_use_parameters(
        self,
        land_cover: str = "esa_worldcover_2021",
    ) -> None:
        """This method is removed."""
        self.logger.warning(
            "setup_land_use_parameters is removed, please remove it from your build configuration"
        )

    @build_method(depends_on=[], required=False)
    def setup_soil_parameters(self) -> None:
        """Deprecated method for setting up soil parameters.

        Raises:
            NotImplementedError: This method is removed; use setup_soil instead.
        """
        raise NotImplementedError(
            "setup_soil_parameters is removed; use setup_soil instead."
        )

    @build_method(depends_on=[], required=True)
    def setup_soil(self) -> None:
        """Sets up the soil parameters for the model.

        Notes:
            This method sets up the soil parameters for the model by retrieving soil data from the CWATM dataset and interpolating
            the data to the model grid. It first retrieves the soil dataset from the `data_catalog`, and
            then retrieves the soil parameters and storage depth data for each soil layer. It then interpolates the data to the
            model grid using the specified interpolation method and sets the resulting grids as attributes of the model.

            Additionally, this method sets up the percolation impeded and crop group data by retrieving the corresponding data
            from the soil dataset and interpolating it to the model grid.

            The resulting soil parameters are set as attributes of the model with names of the form 'soil/{parameter}{soil_layer}',
            where {parameter} is the name of the soil parameter (e.g. 'alpha', 'ksat', etc.) and {soil_layer} is the index of the
            soil layer (1-3; 1 is the top layer). The storage depth data is set as attributes of the model with names of the
            form 'soil/storage_depth{soil_layer}'. The percolation impeded and crop group data are set as attributes of the model
            with names 'soil/percolation_impeded' and 'soil/cropgrp', respectively.
        """
        # Keep this commented code because we want to include differentiation between valley and hillslope soils later
        # soil_depth = self.data_catalog.fetch("GlobalSoilRegolithSediment").read()
        # assert isinstance(soil_depth, xr.Dataset)
        # soil_depth = soil_depth.isel(
        #     get_window(
        #         soil_depth.x,
        #         soil_depth.y,
        #         self.bounds,
        #         buffer=10,
        #     ),
        # )

        soilgrids: xr.Dataset = load_soilgrids(
            self.data_catalog, self.subgrid["mask"], self.region
        )
        self.set_subgrid(soilgrids["silt"], name="soil/silt_percentage")
        self.set_subgrid(soilgrids["clay"], name="soil/clay_percentage")
        self.set_subgrid(soilgrids["bdod"], name="soil/bulk_density_kg_per_dm3")
        self.set_subgrid(soilgrids["soc"], name="soil/soil_organic_carbon_percentage")
        self.set_subgrid(soilgrids["height"], name="soil/soil_layer_height_m")

        depth_to_bedrock_cm = self.data_catalog.fetch("soilgridsv1").read(
            variable="BDTICM_M_250m_ll"
        )
        assert isinstance(depth_to_bedrock_cm, xr.DataArray)
        depth_to_bedrock_cm = depth_to_bedrock_cm.isel(
            get_window(
                depth_to_bedrock_cm.x,
                depth_to_bedrock_cm.y,
                self.bounds,
                buffer=10,
            ),
        ).astype(np.float32)
        depth_to_bedrock_cm: xr.DataArray = convert_nodata(depth_to_bedrock_cm, np.nan)
        depth_to_bedrock_m: xr.DataArray = (
            depth_to_bedrock_cm / 100
        )  # convert from cm to m
        depth_to_bedrock_m: xr.DataArray = interpolate_na_2d(depth_to_bedrock_m)
        depth_to_bedrock_m: xr.DataArray = resample_like(
            depth_to_bedrock_m, soilgrids["silt"]
        )
        self.set_subgrid(depth_to_bedrock_m, name="soil/depth_to_bedrock_m")

    @build_method(depends_on=[], required=True)
    def setup_vegetation(
        self,
    ) -> None:
        """Sets up the vegetation parameters for the model."""
        for vegetation_type in ("forest", "grassland_like"):
            crop_group_number: xr.DataArray = self.data_catalog.fetch(
                f"lisflood_crop_group_number_{vegetation_type}"
            ).read()
            crop_group_number = crop_group_number.isel(
                get_window(
                    crop_group_number.x,
                    crop_group_number.y,
                    self.bounds,
                    buffer=10,
                ),
            )

            crop_group_number = crop_group_number.astype(np.float32)
            crop_group_number = convert_nodata(crop_group_number, np.nan)
            crop_group_number = interpolate_na_2d(crop_group_number)
            crop_group_number = resample_like(
                crop_group_number,
                self.grid["mask"],
                method="nearest",
            )
            self.set_grid(
                crop_group_number,
                name=f"vegetation/crop_group_number_{vegetation_type}",
            )

            leaf_area_index: xr.DataArray = self.data_catalog.fetch(
                f"lisflood_leaf_area_index_{vegetation_type}"
            ).read()
            leaf_area_index = leaf_area_index.isel(
                get_window(
                    leaf_area_index.x,
                    leaf_area_index.y,
                    self.bounds,
                    buffer=10,
                ),
            )

            leaf_area_index = leaf_area_index.astype(np.float32)
            leaf_area_index = convert_nodata(leaf_area_index, np.nan)
            leaf_area_index = interpolate_na_along_dim(leaf_area_index, dim="time")
            leaf_area_index = resample_like(
                leaf_area_index,
                self.grid["mask"],
                method="nearest",
            ).compute()
            self.set_other(
                leaf_area_index, name=f"vegetation/leaf_area_index_{vegetation_type}"
            )

    @build_method(depends_on=[], required=False)
    def setup_forest_restoration_potential(self) -> None:
        """Sets up the forest restoration potential data for the model.

        Source data is in percentage, which is converted to ratio.
        """
        forest_restoration_potential_percentage = self.data_catalog.fetch(
            "forest_restoration_potential"
        ).read()
        assert isinstance(forest_restoration_potential_percentage, xr.DataArray)
        forest_restoration_potential_percentage = (
            forest_restoration_potential_percentage.isel(
                get_window(
                    forest_restoration_potential_percentage.x,
                    forest_restoration_potential_percentage.y,
                    self.bounds,
                    buffer=2,
                ),
            ).compute()
        )

        forest_restoration_potential_percentage = interpolate_na_2d(
            forest_restoration_potential_percentage
        )
        forest_restoration_potential_percentage = resample_like(
            forest_restoration_potential_percentage, self.grid["mask"]
        )
        forest_restoration_potential_ratio = (
            forest_restoration_potential_percentage / 100
        )  # convert from percentage to ratio
        self.set_grid(
            forest_restoration_potential_ratio,
            name="landsurface/forest_restoration_potential_ratio",
        )
