"""
Functions to load sample data.
"""

import importlib
from pathlib import Path
from typing import Union
from zipfile import ZipFile

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

from imod.formats.prj import open_projectfile_data
from imod.mf6 import Modflow6Simulation
from imod.util.imports import MissingOptionalModule

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


def create_pooch_registry() -> pooch.core.Pooch:
    registry = pooch.create(
        path=pooch.os_cache("imod"),
        base_url="https://github.com/Deltares/imod-artifacts/raw/main/",
        version=None,
        version_dev="main",
        env="IMOD_DATA_DIR",
    )
    return registry


def load_pooch_registry(registry: pooch.core.Pooch) -> pooch.core.Pooch:
    with importlib.resources.files("imod.data") as pkg_dir:
        registry.load_registry(pkg_dir / "registry.txt")
    return registry


REGISTRY = create_pooch_registry()
REGISTRY = load_pooch_registry(REGISTRY)


def twri_output(path: Union[str, Path]) -> None:
    fname_twri = REGISTRY.fetch("ex01-twri-output.zip")
    with ZipFile(fname_twri) as archive:
        archive.extractall(path)


def hondsrug_initial() -> xr.Dataset:
    fname = REGISTRY.fetch("hondsrug-initial.nc")
    return xr.open_dataset(fname)


def hondsrug_layermodel() -> xr.Dataset:
    fname = REGISTRY.fetch("hondsrug-layermodel.nc")
    return xr.open_dataset(fname)


def hondsrug_meteorology() -> xr.Dataset:
    fname = REGISTRY.fetch("hondsrug-meteorology.nc")
    return xr.open_dataset(fname)


def hondsrug_river() -> xr.Dataset:
    fname = REGISTRY.fetch("hondsrug-river.nc")
    return xr.open_dataset(fname)


def hondsrug_drainage() -> xr.Dataset:
    fname = REGISTRY.fetch("hondsrug-drainage.nc")
    return xr.open_dataset(fname)


def head_observations() -> pd.DataFrame:
    fname = REGISTRY.fetch("head-observations.csv")
    df = pd.read_csv(fname)
    # Manually convert time column to datetime type because pandas >2.0 doesn't
    # do this automatically anymore upon reading.
    df["time"] = pd.to_datetime(df["time"])
    return df


def fluxes() -> xr.Dataset:
    fname = REGISTRY.fetch("fluxes.nc")
    return xr.open_dataset(fname)


def ahn() -> xr.Dataset:
    fname = REGISTRY.fetch("ahn.nc")
    return xr.open_dataset(fname)


def lakes_shp(path: Union[str, Path]) -> "geopandas.GeoDataFrame":  # type: ignore # noqa
    fname_lakes_shp = REGISTRY.fetch("lakes_shp.zip")
    with ZipFile(fname_lakes_shp) as archive:
        archive.extractall(path)
    return gpd.read_file(Path(path) / "lakes.shp")


def circle() -> xu.Ugrid2d:
    fname_nodes = REGISTRY.fetch("circle-nodes.txt")
    fname_triangles = REGISTRY.fetch("circle-triangles.txt")

    nodes = np.loadtxt(fname_nodes)
    triangles = np.loadtxt(fname_triangles).astype(np.int32)

    return xu.Ugrid2d(*nodes.T, -1, triangles)


def imod5_projectfile_data(path: Union[str, Path]) -> dict:
    fname_model = REGISTRY.fetch("iMOD5_model.zip")

    with ZipFile(fname_model) as archive:
        archive.extractall(path)

    return open_projectfile_data(Path(path) / "iMOD5_model_pooch" / "iMOD5_model.prj")


def hondsrug_simulation(path: Union[str, Path]) -> Modflow6Simulation:
    fname_simulation = REGISTRY.fetch("hondsrug-simulation.zip")

    with ZipFile(fname_simulation) as archive:
        archive.extractall(path)

    simulation = Modflow6Simulation.from_file(Path(path) / "mf6-hondsrug-example.toml")
    # The model was written before the xt3d_option and rhs_option arguments were
    # added to iMOD Python. Set missing options to False.
    simulation["GWF"]["npf"].set_xt3d_option(is_xt3d_used=False, is_rhs=False)

    return simulation


def hondsrug_crosssection(path: Union[str, Path]) -> "geopandas.GeoDataFrame":  # type: ignore # noqa
    fname_simulation = REGISTRY.fetch("hondsrug-crosssection.zip")

    with ZipFile(fname_simulation) as archive:
        archive.extractall(path)

    return gpd.read_file(Path(path) / "crosssection.shp")


def hondsrug_layermodel_topsystem() -> xr.Dataset:
    """
    This is a modified version of the hondsrug_layermodel, used for the
    topsystem example in the user guide. n_max_old original layers are
    taken and subdivided into n_new layers. This makes for more layers around
    the topsystem.
    """
    fname = REGISTRY.fetch("hondsrug-layermodel.nc")
    layer_model = xr.open_dataset(fname)
    # Make layer model more interesting for this example by subdividing layers
    # into n_new layers.
    n_new = 4
    n_max_old = 5

    # Loop over original layers until n_max_old and subdivide each into n_new
    # layers.
    new_ds_ls = []
    for i in range(n_max_old):
        sub_iter = np.arange(n_new) + 1
        layer_coord = sub_iter + i * (n_max_old - 1)
        distribution_factors = 1 / n_new * sub_iter
        da_distribution = xr.DataArray(
            distribution_factors, coords={"layer": layer_coord}, dims=("layer",)
        )
        layer_model_sel = layer_model.sel(layer=i + 1, drop=True)
        # Compute thickness
        D = layer_model_sel["top"] - layer_model_sel["bottom"]

        new_ds = xr.Dataset()
        new_ds["k"] = xr.ones_like(da_distribution) * layer_model_sel["k"]
        new_ds["idomain"] = xr.ones_like(da_distribution) * layer_model_sel["idomain"]
        # Put da_distribution in front of equation to enforce dims as (layer, y, x)
        new_ds["top"] = (da_distribution - 1 / n_new) * -D + layer_model_sel["top"]
        new_ds["bottom"] = da_distribution * -D + layer_model_sel["top"]

        new_ds_ls.append(new_ds)

    return xr.concat(new_ds_ls, dim="layer")


def hondsrug_planar_river() -> xr.Dataset:
    """
    This is the hondsrug river dataset with the following modifications:

    1) Aggregated over layer dimension to create planar grid.
    2) Stages raised towards the top of the model, as the original stages are
       most of the time laying at bottom elevation, making for boring examples.

    """
    fname = REGISTRY.fetch("hondsrug-river.nc")
    planar_river = xr.open_dataset(fname).max(dim="layer")

    fname = REGISTRY.fetch("hondsrug-layermodel.nc")
    top = xr.open_dataset(fname)["top"].sel(layer=1)

    planar_river["stage"] = (top - planar_river["stage"]) / 2 + planar_river["stage"]

    return planar_river
