from typing import List, Optional, Tuple

import numpy as np
import pandas as pd

from imod.logging import init_log_decorator
from imod.mf6.interfaces.imaskingsettings import IMaskingSettings
from imod.mf6.interfaces.iregridpackage import IRegridPackage
from imod.mf6.package import Package
from imod.mf6.utilities.regrid import RegridderType
from imod.mf6.validation import DisBottomSchema
from imod.mf6.write_context import WriteContext
from imod.schemata import (
    AllValueSchema,
    AnyValueSchema,
    DimsSchema,
    DTypeSchema,
    IdentityNoDataSchema,
    IndexesSchema,
)


class VerticesDiscretization(Package, IRegridPackage, IMaskingSettings):
    """
    Discretization by Vertices (DISV).

    Parameters
    ----------
    top: array of floats (xu.UgridDataArray)
    bottom: array of floats (xu.UgridDataArray)
    idomain: array of integers (xu.UgridDataArray)
    validate: {True, False}
        Flag to indicate whether the package should be validated upon
        initialization. This raises a ValidationError if package input is
        provided in the wrong manner. Defaults to True.
    """

    _pkg_id = "disv"

    _init_schemata = {
        "top": [
            DTypeSchema(np.floating),
            DimsSchema("{face_dim}") | DimsSchema(),
            IndexesSchema(),
        ],
        "bottom": [
            DTypeSchema(np.floating),
            DimsSchema("layer", "{face_dim}") | DimsSchema("layer"),
            IndexesSchema(),
        ],
        "idomain": [
            DTypeSchema(np.integer),
            DimsSchema("layer", "{face_dim}"),
            IndexesSchema(),
        ],
    }
    _write_schemata = {
        "idomain": (AnyValueSchema(">", 0),),
        "top": (
            AllValueSchema(">", "bottom", ignore=("idomain", "==", -1)),
            IdentityNoDataSchema(other="idomain", is_other_notnull=(">", 0)),
            # No need to check coords: dataset ensures they align with idomain.
        ),
        "bottom": (DisBottomSchema(other="idomain", is_other_notnull=(">", 0)),),
    }

    _grid_data = {"top": np.float64, "bottom": np.float64, "idomain": np.int32}
    _keyword_map = {"bottom": "botm"}
    _template = Package._initialize_template(_pkg_id)

    _regrid_method = {
        "top": (RegridderType.OVERLAP, "mean"),
        "bottom": (RegridderType.OVERLAP, "mean"),
        "idomain": (RegridderType.OVERLAP, "mode"),
    }

    @property
    def skip_variables(self) -> List[str]:
        return ["bottom"]

    @init_log_decorator()
    def __init__(self, top, bottom, idomain, validate: bool = True):
        dict_dataset = {
            "idomain": idomain,
            "top": top,
            "bottom": bottom,
        }
        super().__init__(dict_dataset)
        self._validate_init_schemata(validate)

    def render(self, directory, pkgname, binary):
        disdirectory = directory / pkgname
        d = {}
        grid = self.dataset.ugrid.grid
        d["xorigin"] = grid.node_x.min()
        d["yorigin"] = grid.node_y.min()
        d["nlay"] = self.dataset["idomain"].coords["layer"].size
        facedim = grid.face_dimension
        d["ncpl"] = self.dataset["idomain"].coords[facedim].size
        d["nvert"] = grid.node_x.size

        _, d["top"] = self._compose_values(
            self.dataset["top"], disdirectory, "top", binary=binary
        )
        d["botm_layered"], d["botm"] = self._compose_values(
            self["bottom"], disdirectory, "botm", binary=binary
        )
        d["idomain_layered"], d["idomain"] = self._compose_values(
            self["idomain"], disdirectory, "idomain", binary=binary
        )
        return self._template.render(d)

    def _verts_dataframe(self) -> pd.DataFrame:
        grid = self.dataset.ugrid.grid
        df = pd.DataFrame(grid.node_coordinates)
        df.index += 1
        return df

    def _cell2d_dataframe(self) -> pd.DataFrame:
        grid = self.dataset.ugrid.grid
        df = pd.DataFrame(grid.face_coordinates)
        df.index += 1
        # modflow requires clockwise; ugrid requires ccw
        face_nodes = grid.face_node_connectivity[:, ::-1]
        df[2] = (face_nodes != grid.fill_value).sum(axis=1)
        for i, column in enumerate(face_nodes.T):
            # Use extension array to write empty values
            # Should be more efficient than mixed column?
            df[3 + i] = pd.arrays.IntegerArray(
                values=column + 1,
                mask=(column == grid.fill_value),
            )
        return df

    def write_blockfile(self, pkgname, globaltimes, write_context: WriteContext):
        dir_for_render = write_context.get_formatted_write_directory()
        content = self.render(dir_for_render, pkgname, write_context.use_binary)
        filename = write_context.write_directory / f"{pkgname}.{self._pkg_id}"
        with open(filename, "w") as f:
            f.write(content)
            f.write("\n\n")

            f.write("begin vertices\n")
            self._verts_dataframe().to_csv(
                f, header=False, sep=" ", lineterminator="\n"
            )
            f.write("end vertices\n\n")

            f.write("begin cell2d\n")
            self._cell2d_dataframe().to_csv(
                f, header=False, sep=" ", lineterminator="\n"
            )
            f.write("end cell2d\n")
        return

    def _validate(self, schemata, **kwargs):
        # Insert additional kwargs
        kwargs["bottom"] = self["bottom"]
        errors = super()._validate(schemata, **kwargs)

        return errors

    def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]:
        return self._regrid_method
