import geopandas
import pandas as pd
from shapely import wkt
import cdsapi
import xarray as xr

from shapes.models import Shape

from datalayers.datasources.base_layer import LayerValueType, LayerTimeResolution
from datalayers.datasources.base_layer import BaseLayer


class Era5T2m(BaseLayer):
    def __init__(self):
        super().__init__()

        self.time_col = LayerTimeResolution.DAY
        self.value_type = LayerValueType.VALUE

    def download(self):
        self._create_data_dir_if_not_exists()

        c = cdsapi.Client()

        c.retrieve(
            "reanalysis-era5-single-levels",
            {
                "product_type": "reanalysis",
                "format": "netcdf",
                "variable": "2m_temperature",
                "year": "2023",
                "month": [
                    "01",
                    "02",
                    "03",
                    "04",
                    "05",
                    "06",
                    "07",
                    "08",
                    "09",
                    "10",
                    "11",
                    "12",
                ],
                "day": [
                    "01",
                    "02",
                    "03",
                    "04",
                    "05",
                    "06",
                    "07",
                    "08",
                    "09",
                    "10",
                    "11",
                    "12",
                    "13",
                    "14",
                    "15",
                    "16",
                    "17",
                    "18",
                    "19",
                    "20",
                    "21",
                    "22",
                    "23",
                    "24",
                    "25",
                    "26",
                    "27",
                    "28",
                    "29",
                    "30",
                    "31",
                ],
                "time": [
                    "00:00",
                    "01:00",
                    "02:00",
                    "03:00",
                    "04:00",
                    "05:00",
                    "06:00",
                    "07:00",
                    "08:00",
                    "09:00",
                    "10:00",
                    "11:00",
                    "12:00",
                    "13:00",
                    "14:00",
                    "15:00",
                    "16:00",
                    "17:00",
                    "18:00",
                    "19:00",
                    "20:00",
                    "21:00",
                    "22:00",
                    "23:00",
                ],
                "area": [
                    11.17,
                    -3.26,
                    4.74,
                    1.2,
                ],
            },
            self.get_data_path() / "2023 2m_temperature.nz",
        )

    def process(self, shapes=None, save_output=False, param_dir=None):
        # get shapes
        if shapes is None:
            shapes = Shape.objects.all()

        # netCFD
        ds = xr.open_mfdataset(self.get_data_path() / "2023 2m_temperature.nz")

        # t2m is provided in Kelvin [K], transform to °C by subtracting 273.15
        # as specified in the Copernicus documentation
        ds["t2m"] -= 273.15
        ds["t2m"].attrs["units"] = "°C"

        # get points inside netCFD
        df = ds.isel(time=0).to_dataframe().reset_index()
        gdf_df_points = geopandas.GeoDataFrame(
            df,
            geometry=geopandas.points_from_xy(df.longitude, df.latitude),
            crs="EPSG:4326",
        )

        for shape in shapes:
            shape_geom = wkt.loads(shape.geometry.wkt)
            shape_points = gdf_df_points[gdf_df_points.within(shape_geom)]

            if len(shape_points) == 0:
                self.layer.warning("No points inside %s", {"shape_id": shape.id})
                continue

            # extract row for each point in the netCFD that is inside the
            # current shape
            datetime_index = ds.time.values
            point_series = []

            for i, row in shape_points.iterrows():
                lat = row["latitude"]
                lng = row["longitude"]

                point_series.append(ds.t2m.sel(longitude=lng, latitude=lat).values)

            # combine to data frame
            datetime_index = pd.to_datetime(datetime_index)
            df = pd.DataFrame(index=datetime_index)

            i = 1
            for s in point_series:
                df[f"p{i}"] = s
                i += 1

            # mean per time step
            df["mean"] = df.max(axis=1)

            # group to mean per day
            daily_df = df.groupby(pd.Grouper(freq="D")).max()

            for i, row in daily_df.iterrows():
                self.rows.append(
                    {
                        "date": i,
                        "value": row["mean"],
                        "shape_id": shape.id,
                    }
                )

        self.save()