"""
This functions in this module extract data from a (specially prepared) netCDF
for a given polygon and compute a number of statistics, which is presented in
table form.

These statistics can also be formatted so they maybe pasted into QGIS, as a
ready made aquifer-aquitard parametrization.
"""
import pathlib
from typing import List, Union

import geopandas as gpd
import numpy as np
import pandas as pd
import shapely.wkt
import xarray as xr


def coord_reference(da_coord):
    """
    Extracts dx, xmin, xmax for a coordinate DataArray, where x is any coordinate.

    If the DataArray coordinates are nonequidistant, dx will be returned as
    1D ndarray instead of float.

    Parameters
    ----------
    a : xarray.DataArray of a coordinate

    Returns
    --------------
    tuple
        (dx, xmin, xmax) for a coordinate x
    """
    x = da_coord.values

    # Possibly non-equidistant
    dx_string = f"d{da_coord.name}"
    if dx_string in da_coord.coords:
        dx = da_coord.coords[dx_string]
        if (dx.shape == x.shape) and (dx.size != 1):
            # choose correctly for decreasing coordinate
            if dx[0] < 0.0:
                end = 0
                start = -1
            else:
                start = 0
                end = -1
            dx = dx.values.astype(np.float64)
            xmin = float(x.min()) - 0.5 * abs(dx[start])
            xmax = float(x.max()) + 0.5 * abs(dx[end])
            # As a single value if equidistant
            if np.allclose(dx, dx[0]):
                dx = dx[0]
        else:
            dx = float(dx)
            xmin = float(x.min()) - 0.5 * abs(dx)
            xmax = float(x.max()) + 0.5 * abs(dx)
    elif x.size == 1:
        raise ValueError(
            f"DataArray has size 1 along {da_coord.name}, so cellsize must be provided"
            f" as a coordinate named d{da_coord.name}."
        )
    else:  # Equidistant
        # TODO: decide on decent criterium for what equidistant means
        # make use of floating point epsilon? E.g:
        # https://github.com/ioam/holoviews/issues/1869#issuecomment-353115449
        dxs = np.diff(x.astype(np.float64))
        dx = dxs[0]
        atolx = abs(1.0e-4 * dx)
        if not np.allclose(dxs, dx, atolx):
            raise ValueError(
                f"DataArray has to be equidistant along {da_coord.name}, or cellsizes"
                f" must be provided as a coordinate named d{da_coord.name}."
            )

        # as xarray uses midpoint coordinates
        xmin = float(x.min()) - 0.5 * abs(dx)
        xmax = float(x.max()) + 0.5 * abs(dx)

    return dx, xmin, xmax


def layer_statistics(path: Union[pathlib.Path, str], gdf: gpd.GeoDataFrame):
    """
    Extract summary statistics of all variables present in a Dataset.

    Variables must contain dimensions ``("layer", "y","x")``.

    Parameters
    ----------
    path: Union[pathlib.Path, str]
        Path to the netCDF dataset containing geohydrologic subsoil properties.
    geodataframe: gpd.GeoDataFrame
        Defines area of interest

    Returns
    -------
    statistics: pandas.DataFrame
        The layers are stored as rows, the variables as columns.
    """
    xmin, ymin, xmax, ymax = gdf.total_bounds
    ds = xr.open_dataset(path).sel(x=slice(xmin, xmax), y=slice(ymax, ymin))
    first_var = next(iter(ds.data_vars.keys()))
    like = ds[first_var].isel(layer=0).drop("layer")
    area_of_interest = rasterize(gdf, like, fill=0) != 0

    stats = pd.DataFrame()
    for variable in ds.data_vars:
        da = ds[variable].where(area_of_interest)
        stats[f"{variable}-min"] = da.min(["y", "x"])  # .to_dataframe(),
        stats[f"{variable}-max"] = da.max(["y", "x"])  # .to_dataframe(),
        stats[f"{variable}-mean"] = da.mean(["y", "x"])  # .to_dataframe(),
        stats[f"{variable}-median"] = da.median(["y", "x"])  # .to_dataframe(),
        stats[f"{variable}-std"] = da.std(["y", "x"])  # .to_dataframe(),
    return stats


def as_aquifer_aquitard(dataframe: pd.DataFrame, statistic: str = "mean"):
    """
    Convert a layer statistics dataframe into a table that can be directly
    copy-pasted into the QGIS plugin Aquifer properties layer.

    Primarily, this means interleaving conductivities and resistances
    in separate rows.

    Parameters
    ----------
    dataframe: pandas.DataFrame
        A dataframe containing layer statistics.
    statistic: str, optional
        Which statistic to use. Default value: ``"mean"``

    Returns
    -------
    aquifer_aquitard_table: pandas.DataFrame
        Reshuffled rows, containing only the data required for the plugin
        Aquifer properties.
    """
    n_layer = len(dataframe)
    out = pd.DataFrame()
    out["fid"] = np.arange(n_layer)
    out["layer"] = np.arange(n_layer)
    out["aquitard_c"] = np.nan
    out["aquitard_c"].values[1:] = dataframe[f"c-{statistic}"].values[:-1]
    out["aquitard_npor"] = np.nan
    out["aquitard_s"] = np.nan
    out["aquifer_k"] = dataframe[f"kh-{statistic}"].values
    out["aquifer_npor"] = np.nan
    out["aquifer_s"] = np.nan
    out["aquifer_top"] = dataframe[f"top-{statistic}"]
    out["aquifer_bottom"] = dataframe[f"bottom-{statistic}"]
    out["semiconf_top"] = np.nan
    out["semiconf_head"] = np.nan
    return out


def netcdf_to_table(inpath: str, outpath: str, wkt_geometry: List[str]) -> None:
    geometries = [shapely.wkt.loads(geom) for geom in wkt_geometry]
    outpath = pathlib.Path(outpath)
    gdf = gpd.GeoDataFrame(geometry=geometries)
    stats = layer_statistics(inpath, gdf)
    stats.to_csv(outpath.parent / f"{outpath.stem}-statistics.csv")
    table = as_aquifer_aquitard(stats)
    table.to_csv(outpath)
