import tempfile
from pathlib import Path

import numpy as np

from pywatemsedem.grasstrips import scale_ktc_with_grass_strip_width, scale_ktc_zhang


def create_ktc(
    composite_landuse,
    cfactor,
    mask,
    ktc_low,
    ktc_high,
    ktc_limit,
    grass=None,
    correction_width=True,
):
    """Create ktc raster based on C-factor raster

    The ktc raster is generated by classifying low and high erosion potential based
    on the C-factor and ktc_limit (i.e. C_factor < ktc_high -> ktc_low, C_factor
    > ktc_limit -> ktc_high).

    The ktc value for landuse rivers, infrastructure and ponds is set to 9999
    (all sediment is routed downwards).

    The ktc values for grass strips are scaled according to their width (see
    :func:`pywatemsedem.ktc.scale_ktc_gdf_grass_strips`).

    Parameters
    ----------
    composite_landuse: numpy.ndarray
        WaTEM/SEDEM composite landuse. (see :ref:`here <watemsedem:prcmap>`).
    cfactor: numpy.ndarray
        C-factor.
    mask: pywatemsedem.geo.rasters.AbstractRaster
        Mask
    ktc_low: float
        Transport coefficient for land covers with low erosion potential
    ktc_high: float
        Transport coefficient for land covers with high erosion potential
    ktc_limit: float
        C-factor to make distinction between ktc_low and ktc_high
    grass: pywatemsedem.geo.vectors.AbstractVector
    correction_width: bool, default True
        Scale kTC according to width of grass strip.

    Returns
    -------
    arr_ktc: numpy.ndarray
        kTC values.
    grass: pywatemsedem.geo.vectors.AbstractVector
        Updated grass strips
    """

    tiff_temp = Path(
        tempfile.NamedTemporaryFile(
            suffix=".tif", prefix="pywatemsedem", delete=False
        ).name
    )
    mask.write(tiff_temp, format="tiff")

    # reclass based on C-factor
    arr_ktc = np.where(cfactor <= ktc_limit, ktc_low, ktc_high)

    # give certain landuse classes an extremely high ktc-value
    arr_ktc = np.where(np.isin(composite_landuse.arr, [-2, -1, -5]), 9999, arr_ktc)

    # set KTC outside modeldomain to zero
    # KTC_arr = np.where(self.catchm.binarr == 1, KTC_arr, 0)

    if correction_width is not None and not grass.is_empty():

        grass._geodata = scale_ktc_gdf_grass_strips(
            grass.geodata,
            ktc_low,
            ktc_high,
        )
        arr_grass = grass.rasterize(
            tiff_temp, composite_landuse.rp.epsg, col="KTC", gdal=True
        )
        arr_ktc = np.where(
            np.logical_and(composite_landuse.arr == -6, arr_grass != mask.rp.nodata),
            arr_grass,
            arr_ktc,
        )
        arr_ktc = np.where(
            np.logical_and(composite_landuse.arr == -6, arr_grass == mask.rp.nodata),
            ktc_low,
            arr_ktc,
        )
    else:
        arr_ktc = np.where(composite_landuse.arr == -6, ktc_low, arr_ktc)

    return arr_ktc, grass


def scale_ktc_gdf_grass_strips(gdf_grass_strips, ktc_low, ktc_high):
    """Scale the ktc values for grass strips in the grass strips dataframe

    Parameters
    ----------
    gdf_grass_strips: geopandas.GeoDataFrame
        Grass strips geodataframe, 'width' and 'scale_ktc' should be present as
        columns.
    ktc_low: float
        Transport coefficient for land covers with low erosion potential
    ktc_high: float
        Transport coefficient for land covers with high erosion potential

    Returns
    -------
    gdf_grass_strips: geopandas.GeoDataFrame
        Updated with the KTC value and the linear sediment trapping efficiency
        (STE_linear).
        For definition STE_linear,
        :func:`pywatemsedem.grasstrips.scale_ktc_with_grass_strip_width`.

    Notes
    -----
    Whether the kTC is scaled for the grass strips depends on the grass strip
    attribute "scale_ktc". If it is not scaled, the ktc_low-value is assigned.
    """
    if ("width" not in gdf_grass_strips.columns) or (
        "scale_ktc" not in gdf_grass_strips.columns
    ):
        msg = "Columns 'BREEDTE' and 'scale_ktc' should be in grass strips dataframe"
        return KeyError(msg)

    arr_ktc_grass, arr_ste = scale_ktc_with_grass_strip_width(
        gdf_grass_strips["width"].values,
        scale_ktc_zhang,
        ktc_high=ktc_high,
    )

    try:
        arr_scale_mask = np.array(gdf_grass_strips["scale_ktc"], dtype=bool)
    except TypeError:
        msg = (
            "Type of 'scale_ktc' column cannot be converted to boolean. Cannot "
            "mask 'KTC' column."
        )
        raise TypeError(msg)

    arr_ktc_grass[~arr_scale_mask] = ktc_low
    arr_ste[~arr_scale_mask] = np.nan

    gdf_grass_strips["KTC"] = arr_ktc_grass
    gdf_grass_strips["STE_linear"] = arr_ste

    return gdf_grass_strips
