from datetime import date
from pathlib import Path
from typing import List, Union

import numpy as np
import pandas as pd
import xarray
from affine import Affine
from geopandas import GeoSeries
from numpy import ndarray
from rasterstats import zonal_stats
from typing import Union


def flatten_stats(stats_dict: List[str], stats: List[str]) -> List[float]:
    return np.array([[item[stat] for stat in stats] for item in stats_dict]).flatten()


def sample_geoseries(
    values: ndarray,
    geometries: Union[List, GeoSeries],
    affine: Affine,
    nodata: float,
    stats: Union[str, List[str]] = "mean",
) -> List[float]:

    # we work with lists to flatten_stats
    if isinstance(stats, str):
        stats = [stats]

    stats_dict = zonal_stats(
        geometries,
        values,
        affine=affine,
        nodata=nodata,
        stats=" ".join(stats),
        boundless=True,
    )

    return flatten_stats(stats_dict, stats)


def sample_netcdf(
    nc_file: Union[Path, str],
    variable_code: str,
    geometries: Union[List, GeoSeries],
    stats: Union[str, List[str]] = "mean",
    start_date: Union[date, None] = None,
    end_date: Union[date, None] = None,
    unlink: bool = False,
) -> pd.DataFrame:
    """Sample a set of geometries over a netcdf file

    Parameters
    ----------
    nc_file : Path or str
        path to NetCDF file
    variable_code : str
        Variable in NetCDF file to sample
    geometries : Union[List, GeoSeries]
        geometries to sample
    stats : List[str]
        statistics to sample
    start_date : Union[date, None]
        start date for selection, by default None
    end_date: Union[date, None]
        end date for selection, by default None
    unlink : bool, optional
        option to delete netcdf-file after sampling, by default False

    Returns
    -------
    pd.DataFrame
        Pandas DataFrame with statistics per timestamp per geometry
    """
    if isinstance(nc_file, str):
        nc_file = Path(str)
    assert nc_file.is_file(), f"nc_file {nc_file} does not exist"

    # read temp-source for sampling
    with xarray.open_dataset(nc_file, engine="netcdf4") as ds:
        if (start_date is not None) and (end_date is not None):
            ds = ds.sel(time=slice(start_date, end_date))
        nodata =  ds[variable_code].encoding.get("_FillValue", None)
        affine = ds.rio.transform()
        data = {
            time: sample_geoseries(
                values=ds[variable_code].sel(time=time).values,
                geometries=geometries,
                affine=affine,
                nodata=nodata,
                stats=stats,
            )
            for time in ds["time"].values
        }

    # delete temp-file
    if unlink:
        if nc_file.exists():
            nc_file.unlink()

    # create columns
    if isinstance(stats, str):
        stats = [stats]

    if len(stats) == 1:
        columns = geometries.index
    else:
        geom_index = geometries.index
        columns = pd.MultiIndex.from_product(
            iterables=[geom_index, stats], names=["index", "stats"]
        )

    return pd.DataFrame.from_dict(data, orient="index", columns=columns)


def sample_netcdfs(
    nc_files: list[Path],
    variable_code: str,
    geometries: Union[List, GeoSeries],
    stats: Union[str, List[str]] = "mean",
    start_date: Union[date, None] = None,
    end_date: Union[date, None] = None,
):
    """Sample over a set of netcdf-files

    Parameters
    ----------
    nc_files : list[Path]
        A list of netcdf-files
    variable_code : str
        Variable in NetCDF file to sample
    geometries : Union[List, GeoSeries]
        geometries to sample
    stats : List[str]
        statistics to sample
    start_date : Union[date, None]
        start date for selection, by default None
    end_date: Union[date, None]
        end date for selection, by default None

    Returns
    ------
    pd.DataFrame
        Pandas DataFrame with statistics per timestamp per geometry
    """
    assert nc_files, f"no NetCDF files to sample"
    for nc_file in nc_files:
        assert nc_file.is_file(), f"nc_file {nc_file} does not exist"

    # read all transforms to see if dataset is consistent
    transforms = list(
        set(xarray.open_dataset(file).rio.transform() for file in nc_files)
    )
    if len(transforms) == 1:
        with xarray.open_dataset(nc_files[0], decode_coords="all") as ds:
            geometries = geometries.to_crs(ds.rio.crs)
    else:
        raise ValueError(
            f"Files do not have one consistent transform. Got {transforms}"
        )

    dfs = [
        sample_netcdf(
            nc_file,
            variable_code,
            geometries,
            stats=stats,
            start_date=start_date,
            end_date=end_date,
            unlink=False,
        )
        for nc_file in nc_files
    ]
    dfs = [i for i in dfs if not i.empty]
    df = pd.concat(dfs).sort_index()

    return df


def sample_nc_dir(
    dir_path: Union[Path, str],
    variable_code: str,
    geometries: Union[List, GeoSeries],
    stats: Union[str, List[str]] = "mean",
    start_date: Union[date, None] = None,
    end_date: Union[date, None] = None,
):
    """Sample over a directory of netcdf-files

    Parameters
    ----------
    dir_path : Union[Path, str]
        Directory with netcdf files
    variable_code : str
        Variable in NetCDF file to sample
    geometries : Union[List, GeoSeries]
        geometries to sample
    stats : List[str]
        statistics to sample
    start_date : Union[date, None]
        start date for selection, by default None
    end_date: Union[date, None]
        end date for selection, by default None

    Returns
    ------
    pd.DataFrame
        Pandas DataFrame with statistics per timestamp per geometry
    """
    if isinstance(dir_path, str):
        dir_path = Path(dir_path)
    assert dir_path.is_dir(), f"dir_path {dir_path} does not exist"

    nc_files = [x for x in dir_path.iterdir() if x.suffix == ".nc"]

    df = sample_netcdfs(
        nc_files,
        variable_code=variable_code,
        geometries=geometries,
        stats=stats,
        start_date=start_date,
        end_date=end_date,
    )
    return df
