"""Interception functions."""

# --------------------------------------------------------------------------------
# This file contains code that has been adapted from an original source available
# in a public repository under the GNU General Public License. The original code
# has been modified to fit the specific needs of this project.
#
# Original source repository: https://github.com/iiasa/CWatM
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
# --------------------------------------------------------------------------------

import numpy as np
import numpy.typing as npt
from numba import njit

from geb.geb_types import Shape

from ..landcovers import (
    FOREST,
    GRASSLAND_LIKE,
    NON_PADDY_IRRIGATED,
    OPEN_WATER,
    PADDY_IRRIGATED,
    SEALED,
)


def leaf_area_index_to_interception_capacity_m(
    leaf_area_index: np.ndarray[Shape, np.dtype[np.float32]],
) -> np.ndarray[Shape, np.dtype[np.float32]]:
    """Convert leaf area index to interception capacity in meters.

    Args:
        leaf_area_index: Leaf area index (LAI) array.

    Returns:
        interception_capacity_m: Interception capacity in meters.
    """
    interception_capacity_m = (
        np.float32(0.935)
        + np.float32(0.498) * leaf_area_index
        - 0.00575 * (leaf_area_index) ** 2
    ) / np.float32(1000.0)  # convert from mm to m
    interception_capacity_m[leaf_area_index <= np.float32(0.1)] = np.float32(0.0)
    return interception_capacity_m


def get_interception_capacity(
    land_use_type: npt.NDArray[np.int32],
    interception_capacity_m_forest_HRU: npt.NDArray[np.float32],
    interception_capacity_m_grassland_HRU: npt.NDArray[np.float32],
) -> npt.NDArray[np.float32]:
    """Get interception capacity based on land use type.

    Args:
        land_use_type: Array of land use types.
        interception_capacity_m_forest_HRU: Interception capacity for forest land use type.
        interception_capacity_m_grassland_HRU: Interception capacity for grassland land use type

    Returns:
        interception_capacity_m: Array of interception capacities corresponding to land use types.
    """
    interception_capacity_m = np.full(land_use_type.shape, np.nan, dtype=np.float32)
    interception_capacity_m[land_use_type == OPEN_WATER] = 0.0
    interception_capacity_m[land_use_type == SEALED] = 0.0
    interception_capacity_m[land_use_type == PADDY_IRRIGATED] = 0.001  # 1 mm
    interception_capacity_m[land_use_type == NON_PADDY_IRRIGATED] = 0.001  # 1 mm
    interception_capacity_m[land_use_type == FOREST] = (
        interception_capacity_m_forest_HRU[land_use_type == FOREST]
    )
    interception_capacity_m[land_use_type == GRASSLAND_LIKE] = (
        interception_capacity_m_grassland_HRU[land_use_type == GRASSLAND_LIKE]
    )
    assert not np.isnan(interception_capacity_m).any()
    return interception_capacity_m


def get_leaf_area_index(
    land_use_type: npt.NDArray[np.int32],
    leaf_area_index_forest_HRU: npt.NDArray[np.float32],
    leaf_area_index_grassland_HRU: npt.NDArray[np.float32],
    crop_map: npt.NDArray[np.int32],
) -> npt.NDArray[np.float32]:
    """Get Leaf Area Index (LAI) based on land use type.

    Args:
        land_use_type: Array of land use types.
        leaf_area_index_forest_HRU: LAI for forest land use type.
        leaf_area_index_grassland_HRU: LAI for grassland land use type
        crop_map: Array of crop types for each cell.

    Returns:
        leaf_area_index: Array of LAI corresponding to land use types.
    """
    leaf_area_index = np.zeros(land_use_type.shape, dtype=np.float32)

    mask_forest = land_use_type == FOREST
    leaf_area_index[mask_forest] = leaf_area_index_forest_HRU[mask_forest]

    mask_grassland = (land_use_type == GRASSLAND_LIKE) & (
        crop_map == -1
    )  # Only assign grassland LAI to non-crop grasslands
    leaf_area_index[mask_grassland] = leaf_area_index_grassland_HRU[mask_grassland]

    mask_cropland = (
        (land_use_type == GRASSLAND_LIKE)
        | (land_use_type == PADDY_IRRIGATED)
        | (land_use_type == NON_PADDY_IRRIGATED)
    ) & (crop_map != -1)
    # TODO: refine this
    leaf_area_index[mask_cropland] = 3.0  # Assign a default LAI for croplands

    return leaf_area_index


@njit(cache=True, inline="always")
def interception(
    rainfall_m: np.float32,
    storage_m: np.float32,
    capacity_m: np.float32,
    potential_interception_evaporation_m: np.float32,
    potential_transpiration_m: np.float32,
    potential_direct_evaporation_m: np.float32,
) -> tuple[np.float32, np.float32, np.float32, np.float32, np.float32]:
    """Calculate interception storage, throughfall, and evaporation.

    The potential transpiration and potential direct evaporation are reduced by
    the amount of evaporation from the interception storage, with priority given
    to reducing transpiration.

    Args:
        rainfall_m: Precipitation (rain) (m).
        storage_m: Current interception storage (m).
        capacity_m: Interception capacity of vegetation (m).
        potential_interception_evaporation_m: Potential evaporation from a wet surface (m).
        potential_transpiration_m: Potential transpiration (m).
        potential_direct_evaporation_m: Potential direct evaporation (soil/water) (m).

    Returns:
        new_storage: Updated interception storage (m).
        throughfall: Water reaching the ground after interception (m).
        evaporation: Evaporation from intercepted water (m).
        potential_transpiration_m: Updated potential transpiration (m).
        potential_direct_evaporation_m: Updated potential direct evaporation (m).
    """
    # Calculate throughfall
    throughfall = max(np.float32(0.0), rainfall_m + storage_m - capacity_m)

    # Update interception storage after throughfall
    new_storage = storage_m + rainfall_m - throughfall

    # Calculate evaporation from intercepted water
    evaporation = min(
        new_storage,
        potential_interception_evaporation_m
        * (new_storage / capacity_m) ** np.float32(2.0 / 3.0)
        if capacity_m > np.float32(0.0)
        else np.float32(0.0),
    )

    # Update interception storage after evaporation
    new_storage -= evaporation

    # Subtract evaporation from transpiration potential first.
    evaporation_from_transpiration = min(evaporation, potential_transpiration_m)
    potential_transpiration_m -= evaporation_from_transpiration

    # Subtract remaining evaporation from direct (soil/water) potential.
    evaporation_from_direct = max(
        np.float32(0.0), evaporation - evaporation_from_transpiration
    )
    potential_direct_evaporation_m -= evaporation_from_direct
    potential_direct_evaporation_m = max(
        np.float32(0.0), potential_direct_evaporation_m
    )

    return (
        new_storage,
        throughfall,
        evaporation,
        potential_transpiration_m,
        potential_direct_evaporation_m,
    )
