"""
Most of the functionality here attempts to replicate what iMOD does with
project files.
"""

from __future__ import annotations

import itertools
import pickle
from collections import Counter
from datetime import datetime
from typing import Dict, List, Optional, Tuple, cast

import numpy as np
import pandas as pd
import xarray as xr
import xugrid as xu

import imod
from imod.mf6.model import Modflow6Model
from imod.mf6.utilities.package import get_repeat_stress
from imod.prepare.layer import get_upper_active_grid_cells
from imod.typing import GridDataArray
from imod.util.imports import MissingOptionalModule

try:
    import geopandas as gpd
except ImportError:
    gpd = MissingOptionalModule("geopandas")


def hash_xy(da: xr.DataArray) -> Tuple[int]:
    """
    Create a unique identifier based on the x and y coordinates of a DataArray.
    """
    x = hash(pickle.dumps(da["x"].values))
    y = hash(pickle.dumps(da["y"].values))
    return x, y


class SingularTargetRegridderWeightsCache:
    """
    Create a mapping of (source_coords, regridder_cls) => regridding weights.

    Allows re-use of the regridding weights, as computing the weights is the
    most costly step.

    The regridder only processes x and y coordinates: we can hash these,
    and get a unique identifier. The target is assumed to be constant.
    """

    def __init__(self, projectfile_data, target, cache_size: int):
        # Collect the x-y coordinates of all x-y dimensioned DataArrays.
        # Determine which regridding method to use.
        # Count occurrences of both.
        # Create and cache weights of the most common ones.
        keys = []
        sources = {}
        methods = {}
        for pkgdict in projectfile_data.values():
            for variable, da in pkgdict.items():
                xydims = {"x", "y"}

                if isinstance(da, xr.DataArray) and xydims.issubset(da.dims):
                    # for initial condition, constant head, general head boundary
                    if variable == "head":
                        cls = xu.BarycentricInterpolator
                        method = None
                    elif variable == "conductance":
                        cls = xu.RelativeOverlapRegridder
                        method = "conductance"
                    else:
                        cls = xu.OverlapRegridder
                        method = "mean"

                    x, y = hash_xy(da)
                    key = (x, y, cls)
                    keys.append(key)
                    sources[key] = da
                    methods[key] = method

        counter = Counter(keys)
        self.target = target
        self.weights = {}
        for key, _ in counter.most_common(cache_size):
            cls = key[2]
            ugrid_source = xu.UgridDataArray.from_structured(sources[key])
            kwargs = {"source": ugrid_source, "target": target}
            method = methods[key]
            if method is not None:
                kwargs["method"] = method
            regridder = cls(**kwargs)
            self.weights[key] = regridder.weights

    def regrid(
        self,
        source: xr.DataArray,
        method: str = "mean",
        original2d: Optional[xr.DataArray] = None,
    ):
        if source.dims[-2:] != ("y", "x"):  # So it's a constant
            if original2d is None:
                raise ValueError("original2d must be provided for constant values")
            source = source * xr.ones_like(original2d)

        kwargs = {"target": self.target}
        if method == "barycentric":
            cls = xu.BarycentricInterpolator
        elif method == "conductance":
            cls = xu.RelativeOverlapRegridder
            kwargs["method"] = method
        else:
            cls = xu.OverlapRegridder
            kwargs["method"] = method

        x, y = hash_xy(source)
        key = (x, y, cls)
        if key in self.weights:
            kwargs["weights"] = self.weights[key]
            regridder = cls.from_weights(**kwargs)
            # Avoid creation of a UgridDataArray here
            dims = source.dims[:-2]
            coords = {k: source.coords[k] for k in dims}
            facedim = self.target.face_dimension
            face_source = xr.DataArray(
                source.data.reshape(*source.shape[:-2], -1),
                coords=coords,
                dims=[*dims, facedim],
                name=source.name,
            )
            return xu.UgridDataArray(
                regridder.regrid_dataarray(face_source, (facedim,)),
                regridder._target.ugrid_topology,
            )
        else:
            ugrid_source = xu.UgridDataArray.from_structured(source)
            kwargs["source"] = ugrid_source
            regridder = cls(**kwargs)
            return regridder.regrid(ugrid_source)


def raise_on_layer(value, variable: str):
    da = value[variable]
    if "layer" in da.dims:
        raise ValueError(f"{variable} should not be assigned a layer")
    return da


def finish(uda):
    """
    Set dimension order, and drop empty layers.
    """
    facedim = uda.ugrid.grid.face_dimension
    dims = ("time", "layer", facedim)
    return uda.transpose(*dims, missing_dims="ignore").dropna("layer", how="all")


def create_idomain(thickness):
    """
    Find cells that should get a passthrough value: IDOMAIN = -1.

    We may find them by forward and back-filling: if they are filled in both, they
    contain active cells in both directions, and the value should be set to -1.
    """
    active = (
        thickness > 0
    ).compute()  # TODO larger than a specific value for round off?
    ones = xu.ones_like(active, dtype=float).where(active)
    passthrough = ones.ffill("layer").notnull() & ones.bfill("layer").notnull()
    idomain = ones.combine_first(
        xu.full_like(active, -1.0, dtype=float).where(passthrough)
    )
    return idomain.fillna(0).astype(int)


def create_disv(
    cache,
    top,
    bottom,
    ibound,
):
    if top.dims == ("layer",):
        if ibound.dims != ("layer", "y", "x"):
            raise ValueError(
                "Either ibound or top must have dimensions (layer, y, x) to "
                "derive model extent. Both may not be provided as constants."
            )
        top = top * xr.ones_like(ibound)
        original2d = ibound.isel(layer=0, drop=True)
    else:
        original2d = top.isel(layer=0, drop=True)

    if bottom.dims == ("layer",):
        bottom = bottom * xr.ones_like(ibound)

    top = top.compute()
    bottom = bottom.compute()
    disv_top = cache.regrid(top).compute()
    disv_bottom = cache.regrid(bottom).compute()
    thickness = disv_top - disv_bottom
    idomain = create_idomain(thickness)
    disv = imod.mf6.VerticesDiscretization(
        top=disv_top.sel(layer=1),
        bottom=disv_bottom,
        idomain=idomain,
    )
    active = idomain > 0
    return disv, disv_top, disv_bottom, active, original2d


def create_npf(
    cache,
    k,
    vertical_anisotropy,
    active,
    original2d,
):
    disv_k = cache.regrid(k, method="geometric_mean", original2d=original2d).where(
        active
    )
    k33 = k * vertical_anisotropy
    disv_k33 = cache.regrid(k33, method="harmonic_mean", original2d=original2d).where(
        active
    )
    return imod.mf6.NodePropertyFlow(
        icelltype=0,
        k=disv_k,
        k33=disv_k33,
    )


def create_chd(
    cache,
    model,
    key,
    value,
    ibound,
    active,
    original2d,
    repeat,
    **kwargs,
):
    head = value["head"]

    if "layer" in head.coords:
        layer = head.layer
        ibound = ibound.sel(layer=layer)
        active = active.sel(layer=layer)

    disv_head = cache.regrid(
        head,
        method="barycentric",
        original2d=original2d,
    )
    valid = (ibound < 0) & active

    if not valid.any():
        return

    chd = imod.mf6.ConstantHead(head=disv_head.where(valid))
    if repeat is not None:
        chd.dataset["repeat_stress"] = get_repeat_stress(repeat)
    model[key] = chd
    return


def create_drn(
    cache,
    model,
    key,
    value,
    active,
    top,
    bottom,
    original2d,
    repeat,
    **kwargs,
):
    conductance = raise_on_layer(value, "conductance")
    elevation = raise_on_layer(value, "elevation")

    disv_cond = cache.regrid(
        conductance,
        method="conductance",
        original2d=original2d,
    )
    disv_elev = cache.regrid(elevation, original2d=original2d)
    valid = (disv_cond > 0) & disv_elev.notnull() & active
    location = xu.ones_like(active, dtype=float)
    location = location.where((disv_elev > bottom) & (disv_elev <= top)).where(valid)
    disv_cond = finish(location * disv_cond)
    disv_elev = finish(location * disv_elev)

    if disv_cond.isnull().all():
        return

    drn = imod.mf6.Drainage(
        elevation=disv_elev,
        conductance=disv_cond,
    )
    if repeat is not None:
        drn.dataset["repeat_stress"] = get_repeat_stress(repeat)
    model[key] = drn
    return


def create_ghb(
    cache,
    model,
    key,
    value,
    active,
    original2d,
    repeat,
    **kwargs,
):
    conductance = value["conductance"]
    head = value["head"]

    disv_cond = cache.regrid(
        conductance,
        method="conductance",
        original2d=original2d,
    )
    disv_head = cache.regrid(
        head,
        method="barycentric",
        original2d=original2d,
    )
    valid = (disv_cond > 0.0) & disv_head.notnull() & active

    ghb = imod.mf6.GeneralHeadBoundary(
        conductance=disv_cond.where(valid),
        head=disv_head.where(valid),
    )
    if repeat is not None:
        ghb.dataset["repeat_stress"] = get_repeat_stress(repeat)
    model[key] = ghb
    return


def create_riv(
    cache,
    model,
    key,
    value,
    active,
    original2d,
    top,
    bottom,
    repeat,
    **kwargs,
):
    def assign_to_layer(
        conductance, stage, elevation, infiltration_factor, top, bottom, active
    ):
        """
        Assign river boundary to multiple layers. Distribute the conductance based
        on the vertical degree of overlap.

        Parameters
        ----------
        conductance
        stage:
            water stage
        elevation:
            bottom elevation of the river
        infiltration_factor
            factor (generally <1) to reduce infiltration conductance compared
            to drainage conductance.
        top:
            layer model top elevation
        bottom:
            layer model bottom elevation
        active:
            active or inactive cells (idomain > 0)
        """
        valid = conductance > 0.0
        conductance = conductance.where(valid)
        stage = stage.where(valid)
        elevation = elevation.where(valid)
        elevation = elevation.where(elevation <= stage, other=stage)

        # TODO: this removes too much when the stage is higher than the top...
        # Instead: just cut through all layers until the bottom elevation.
        # Then, assign a transmissivity weighted conductance.
        water_top = stage.where(stage <= top)
        water_bottom = elevation.where(elevation > bottom)
        layer_height = top - bottom
        layer_height = layer_height.where(active)  # avoid 0 thickness layers
        fraction = (water_top - water_bottom) / layer_height
        # Set values of 0.0 to 1.0, but do not change NaN values:
        fraction = fraction.where(~(fraction == 0.0), other=1.0)
        location = xu.ones_like(fraction).where(fraction.notnull() & active)

        layered_conductance = finish(conductance * fraction)
        layered_stage = finish(stage * location)
        layered_elevation = finish(elevation * location)
        infiltration_factor = finish(infiltration_factor * location)

        return (
            layered_conductance,
            layered_stage,
            layered_elevation,
            infiltration_factor,
        )

    conductance = raise_on_layer(value, "conductance")
    stage = raise_on_layer(value, "stage")
    bottom_elevation = raise_on_layer(value, "bottom_elevation")
    infiltration_factor = raise_on_layer(value, "infiltration_factor")

    disv_cond_2d = cache.regrid(
        conductance,
        method="conductance",
        original2d=original2d,
    ).compute()
    disv_elev_2d = cache.regrid(bottom_elevation, original2d=original2d)
    disv_stage_2d = cache.regrid(stage, original2d=original2d)
    disv_inff_2d = cache.regrid(infiltration_factor, original2d=original2d)

    disv_cond, disv_stage, disv_elev, disv_inff = assign_to_layer(
        conductance=disv_cond_2d,
        stage=disv_stage_2d,
        elevation=disv_elev_2d,
        infiltration_factor=disv_inff_2d,
        top=top,
        bottom=bottom,
        active=active,
    )

    if disv_cond.isnull().all():
        return

    # The infiltration factor may be 0. In that case, we need only create a DRN
    # package.
    drn = imod.mf6.Drainage(
        conductance=(1.0 - disv_inff) * disv_cond,
        elevation=disv_stage,
    )
    if repeat is not None:
        drn.dataset["repeat_stress"] = get_repeat_stress(repeat)
    model[f"{key}-drn"] = drn

    riv_cond = disv_cond * disv_inff
    riv_valid = riv_cond > 0.0
    if not riv_valid.any():
        return

    riv = imod.mf6.River(
        stage=disv_stage.where(riv_valid),
        conductance=riv_cond.where(riv_valid),
        bottom_elevation=disv_elev.where(riv_valid),
    )
    if repeat is not None:
        riv.dataset["repeat_stress"] = get_repeat_stress(repeat)
    model[key] = riv

    return


def create_rch(
    cache,
    model,
    key,
    value,
    active,
    original2d,
    repeat,
    **kwargs,
):
    rate = raise_on_layer(value, "rate") * 0.001
    disv_rate = cache.regrid(rate, original2d=original2d).where(active)
    # Find highest active layer
    location = get_upper_active_grid_cells(active)
    disv_rate = finish(disv_rate.where(location))

    # Skip if there's no data
    if disv_rate.isnull().all():
        return

    rch = imod.mf6.Recharge(rate=disv_rate)
    if repeat is not None:
        rch.dataset["repeat_stress"] = get_repeat_stress(repeat)
    model[key] = rch
    return


def create_evt(
    cache,
    model,
    key,
    value,
    active,
    original2d,
    repeat,
    **kwargs,
):
    surface = raise_on_layer(value, "surface")
    rate = raise_on_layer(value, "rate") * 0.001
    depth = raise_on_layer(value, "depth")

    # Find highest active layer
    highest = active["layer"] == active["layer"].where(active).min()
    location = highest.where(highest)

    disv_surface = cache.regrid(surface, original2d=original2d).where(active)
    disv_surface = finish(disv_surface * location)

    disv_rate = cache.regrid(rate, original2d=original2d).where(active)
    disv_rate = finish(disv_rate * location)

    disv_depth = cache.regrid(depth, original2d=original2d).where(active)
    disv_depth = finish(disv_depth * location)

    # At depth 1.0, the rate is 0.0.
    proportion_depth = xu.ones_like(disv_surface).where(disv_surface.notnull())
    proportion_rate = xu.zeros_like(disv_surface).where(disv_surface.notnull())

    evt = imod.mf6.Evapotranspiration(
        surface=disv_surface,
        rate=disv_rate,
        depth=disv_depth,
        proportion_rate=proportion_rate,
        proportion_depth=proportion_depth,
    )
    if repeat is not None:
        evt.dataset["repeat_stress"] = get_repeat_stress(repeat)
    model[key] = evt
    return


def create_sto(
    cache,
    storage_coefficient,
    active,
    original2d,
    transient,
):
    if storage_coefficient is None:
        disv_coef = 0.0
    else:
        disv_coef = cache.regrid(storage_coefficient, original2d=original2d).where(
            active
        )

    sto = imod.mf6.StorageCoefficient(
        storage_coefficient=disv_coef,
        specific_yield=0.0,
        transient=transient,
        convertible=0,
    )
    return sto


def create_wel(
    cache,
    model,
    key,
    value,
    active,
    top,
    bottom,
    k,
    repeat,
    **kwargs,
):
    target = cache.target
    dataframe = value["dataframe"]
    layer = value["layer"]

    if layer <= 0:
        dataframe = imod.prepare.assign_wells(
            wells=dataframe,
            top=top,
            bottom=bottom,
            k=k,
            minimum_thickness=0.01,
            minimum_k=1.0,
        )
    else:
        dataframe["index"] = np.arange(len(dataframe))
        dataframe["layer"] = layer

    first = dataframe.groupby("index").first()
    well_layer = first["layer"].values
    xy = np.column_stack([first["x"], first["y"]])
    cell2d = target.locate_points(xy)
    valid = (cell2d >= 0) & active.values[well_layer - 1, cell2d]

    cell2d = cell2d[valid] + 1
    # Skip if no wells are located inside cells
    if not valid.any():
        return

    if "time" in dataframe.columns:
        # Ensure the well data is rectangular.
        time = np.unique(dataframe["time"].values)
        dataframe = dataframe.set_index("time")
        # First ffill, then bfill!
        dfs = [df.reindex(time).ffill().bfill() for _, df in dataframe.groupby("index")]
        rate = (
            pd.concat(dfs)
            .reset_index()
            .set_index(["time", "index"])["rate"]
            .to_xarray()
        )
    else:
        rate = xr.DataArray(
            dataframe["rate"], coords={"index": dataframe["index"]}, dims=["index"]
        )

    # Don't forget to remove the out-of-bounds points.
    rate = rate.where(xr.DataArray(valid, dims=["index"]), drop=True)

    wel = imod.mf6.WellDisVertices(
        layer=well_layer,
        cell2d=cell2d,
        rate=rate,
    )
    if repeat is not None:
        wel.dataset["repeat_stress"] = get_repeat_stress(repeat)

    model[key] = wel
    return


def create_ic(cache, model, key, value, active, **kwargs):
    start = value["head"]
    disv_start = cache.regrid(source=start, method="barycentric").where(active)
    model[key] = imod.mf6.InitialConditions(start=disv_start)
    return


def create_hfb(
    model: Modflow6Model,
    key: str,
    value: Dict,
    top: GridDataArray,
    bottom: GridDataArray,
    **kwargs,
) -> None:
    dataframe = value["geodataframe"]

    barrier_gdf = gpd.GeoDataFrame(
        geometry=dataframe["geometry"].values,
        data={
            "resistance": dataframe["resistance"].values,
            "ztop": np.ones_like(dataframe["geometry"].values) * top.max().values,
            "zbottom": np.ones_like(dataframe["geometry"].values) * bottom.min().values,
        },
    )

    model[key] = imod.mf6.HorizontalFlowBarrierResistance(barrier_gdf)


def merge_hfbs(
    horizontal_flow_barriers: List[imod.mf6.HorizontalFlowBarrierResistance],
):
    datasets = []
    for horizontal_flow_barrier in horizontal_flow_barriers:
        datasets.append(horizontal_flow_barrier.dataset)

    combined_dataset = xr.concat(datasets, "index")
    combined_dataset.coords["index"] = np.arange(combined_dataset.sizes["index"])

    combined_dataframe = cast(gpd.GeoDataFrame, combined_dataset.to_dataframe())
    combined_dataframe.drop("print_input", axis=1, inplace=True)  # noqa: PD002

    return imod.mf6.HorizontalFlowBarrierResistance(combined_dataframe)


PKG_CONVERSION = {
    "chd": create_chd,
    "drn": create_drn,
    "evt": create_evt,
    "ghb": create_ghb,
    "hfb": create_hfb,
    "shd": create_ic,
    "rch": create_rch,
    "riv": create_riv,
    "wel": create_wel,
}


def expand_repetitions(
    repeat_stress: List[datetime], time_min: datetime, time_max: datetime
) -> Dict[datetime, datetime]:
    expanded = {}
    for year, date in itertools.product(
        range(time_min.year, time_max.year + 1),
        repeat_stress,
    ):
        newdate = date.replace(year=year)
        if newdate < time_max:
            expanded[newdate] = date
    return expanded


def convert_to_disv(
    projectfile_data, target, time_min=None, time_max=None, repeat_stress=None
):
    """
    Convert the contents of a project file to a MODFLOW6 DISV model.

    The ``time_min`` and ``time_max`` are **both** required when
    ``repeat_stress`` is given. The entries in the Periods section of the
    project file will be expanded to yearly repeats between ``time_min`` and
    ``time_max``.

    Additionally, ``time_min`` and ``time_max`` may be used to slice the input
    to a specific time domain.

    The returned model is steady-state if none of the packages contain a time
    dimension. The model is transient if any of the packages contain a time
    dimension. This can be changed by setting the "transient" value in the
    storage package of the returned model. Storage coefficient input is
    required for a transient model.

    Parameters
    ----------
    projectfile_data: dict
        Dictionary with the projectfile topics as keys, and the data
        as xarray.DataArray, pandas.DataFrame, or geopandas.GeoDataFrame.
    target: xu.Ugrid2d
        The unstructured target topology. All data is transformed to match this
        topology.
    time_min: datetime, optional
        Minimum starting time of a stress.
        Required when ``repeat_stress`` is provided.
    time_max: datetime, optional
        Maximum starting time of a stress.
        Required when ``repeat_stress`` is provided.
    repeat_stress: dict of dict of string to datetime, optional
        This dict contains contains, per topic, the period alias (a string) to
        its datetime.

    Returns
    -------
    disv_model: imod.mf6.GroundwaterFlowModel

    """
    if repeat_stress is not None:
        if time_min is None or time_max is None:
            raise ValueError(
                "time_min and time_max are required when repeat_stress is given"
            )

    for arg in (time_min, time_max):
        if arg is not None and not isinstance(arg, datetime):
            raise TypeError(
                "time_min and time_max must be datetime.datetime. "
                f"Received: {type(arg).__name__}"
            )

    data = projectfile_data.copy()
    model = imod.mf6.GroundwaterFlowModel()

    # Setup the regridding weights cache.
    weights_cache = SingularTargetRegridderWeightsCache(data, target, cache_size=5)

    # Mandatory packages first.
    ibound = data["bnd"]["ibound"].compute()
    disv, top, bottom, active, original2d = create_disv(
        cache=weights_cache,
        top=data["top"]["top"],
        bottom=data["bot"]["bottom"],
        ibound=ibound,
    )

    npf = create_npf(
        cache=weights_cache,
        k=data["khv"]["kh"],
        vertical_anisotropy=data["kva"]["vertical_anisotropy"],
        active=active,
        original2d=original2d,
    )

    model["npf"] = npf
    model["disv"] = disv
    model["oc"] = imod.mf6.OutputControl(save_head="all")

    # Used in other package construction:
    k = npf["k"].compute()
    new_ibound = weights_cache.regrid(source=ibound, method="minimum").compute()

    # Boundary conditions, one by one.
    for key, value in data.items():
        pkg = key.split("-")[0]
        convert = PKG_CONVERSION.get(pkg)
        # Skip unsupported packages
        if convert is None:
            continue

        if repeat_stress is None:
            repeat = None
        else:
            repeat = repeat_stress.get(key)
            if repeat is not None:
                repeat = expand_repetitions(repeat, time_min, time_max)

        try:
            # conversion will update model instance
            convert(
                cache=weights_cache,
                model=model,
                key=key,
                value=value,
                ibound=new_ibound,
                active=active,
                original2d=original2d,
                top=top,
                bottom=bottom,
                k=k,
                repeat=repeat,
            )
        except Exception as e:
            raise type(e)(f"{e}\nduring conversion of {key}")

    # Treat hfb's separately: they must be merged into one,
    # as MODFLOW6 only supports a single HFB.
    hfb_keys = [key for key in model.keys() if key.split("-")[0] == "hfb"]
    hfbs = [model.pop(key) for key in hfb_keys]
    if hfbs:
        model["hfb"] = merge_hfbs(hfbs)

    transient = any("time" in pkg.dataset.dims for pkg in model.values())
    if transient and (time_min is not None or time_max is not None):
        model = model.clip_box(time_min=time_min, time_max=time_max)

    sto_entry = data.get("sto")
    if sto_entry is None:
        if transient:
            raise ValueError("storage input is required for a transient run")
        storage_coefficient = None
    else:
        storage_coefficient = sto_entry["storage_coefficient"]

    model["sto"] = create_sto(
        cache=weights_cache,
        storage_coefficient=storage_coefficient,
        active=active,
        original2d=original2d,
        transient=transient,
    )

    return model
