"""Mfusg module."""

import os.path
from inspect import getfullargspec
from os import PathLike, curdir
from typing import Union

import flopy

from ..mbase import PackageLoadException
from ..modflow import Modflow
from ..utils import mfreadnam


class MfUsg(Modflow):
    """MODFLOW-USG Model Class.

    Parameters
    ----------
    modelname : str or PathLike, default "modflowusgtest"
        Name of model.  This string will be used to name the MODFLOW input
        that are created with write_model.
    namefile_ext : str, default "nam"
        Extension for the namefile.
    exe_name : str, default "mfusg"
        The name of the executable to use.
    structured : bool, default True
        Specify if model grid is structured (default) or unstructured.
    listunit : int, default 2
        Unit number for the list file.
    model_ws : str, default "." (curdir)
        Model workspace.  Directory name to create model data sets.
    external_path : str, optional
        Location for external files.
    verbose : bool, default False
        Print additional information to the screen.


    Attributes
    ----------

    Methods
    -------

    See Also
    --------

    Notes
    -----

    Examples
    --------
    >>> import flopy
    >>> usg = flopy.mfusg.MfUsg()
    """

    def __init__(
        self,
        modelname="mfusgtest",
        namefile_ext="nam",
        version="mfusg",
        exe_name: Union[str, PathLike] = "mfusg",
        structured=True,
        listunit=2,
        model_ws: Union[str, PathLike] = curdir,
        external_path=None,
        verbose=False,
        **kwargs,
    ):
        super().__init__(
            modelname=modelname,
            namefile_ext=namefile_ext,
            version=version,
            exe_name=exe_name,
            structured=structured,
            listunit=listunit,
            model_ws=model_ws,
            external_path=external_path,
            verbose=verbose,
            **kwargs,
        )

        self.itrnsp = 0  # transport simulation flag
        self.mcomp = 0  # number of chemical components
        self.iheat = 0  # flag for heat transport
        self.idpf = 0  # flag for dual porosity flow
        self.icln = 0  # flag for CLN package

        # Create a dictionary to map package with package object.
        # This is used for loading models.
        self.mfnam_packages = {
            "zone": flopy.modflow.ModflowZon,
            "mult": flopy.modflow.ModflowMlt,
            "pval": flopy.modflow.ModflowPval,
            "bas6": flopy.mfusg.MfUsgBas,
            "dis": flopy.mfusg.MfUsgDis,
            "hfb6": flopy.modflow.ModflowHfb,
            "fhb": flopy.modflow.ModflowFhb,
            "drn": flopy.modflow.ModflowDrn,
            "drt": flopy.modflow.ModflowDrt,
            "ghb": flopy.modflow.ModflowGhb,
            "riv": flopy.modflow.ModflowRiv,
            "str": flopy.modflow.ModflowStr,
            "sfr": flopy.modflow.ModflowSfr2,
            "gage": flopy.modflow.ModflowGage,
            "sub": flopy.modflow.ModflowSub,
            "swt": flopy.modflow.ModflowSwt,
            "chd": flopy.modflow.ModflowChd,
            "disu": flopy.mfusg.MfUsgDisU,
            "sms": flopy.mfusg.MfUsgSms,
            "wel": flopy.mfusg.MfUsgWel,
            "bcf6": flopy.mfusg.MfUsgBcf,
            "lpf": flopy.mfusg.MfUsgLpf,
            "cln": flopy.mfusg.MfUsgCln,
            "gnc": flopy.mfusg.MfUsgGnc,
            "bct": flopy.mfusg.MfUsgBct,
            "pcb": flopy.mfusg.MfUsgPcb,
            "ddf": flopy.mfusg.MfUsgDdf,
            "mdt": flopy.mfusg.MfUsgMdt,
            "dpf": flopy.mfusg.MfUsgDpf,
            "dpt": flopy.mfusg.MfUsgDpt,
            "rch": flopy.mfusg.MfUsgRch,
            "oc": flopy.mfusg.MfUsgOc,
            "lak": flopy.mfusg.MfUsgLak,
            "evt": flopy.mfusg.MfUsgEvt,
        }

    def __repr__(self):
        """Returns a representation of the MfUsg object."""
        nrow, ncol, nlay, nper = self.get_nrow_ncol_nlay_nper()
        if nrow is not None:
            # structured case
            msg = (
                f"MODFLOW {nlay} layer(s) {nrow} row(s) {ncol} column(s) "
                f"{nper} stress period(s)"
            )
        else:
            # unstructured case
            msg = (
                "MODFLOW unstructured\n"
                f"  nodes = {ncol.sum()}\n"
                f"  layers = {nlay}\n"
                f"  stress periods = {nper}\n"
                f"  nodelay = {ncol}\n"
            )
        return msg

    @classmethod
    def load(
        cls,
        f: str,
        version="mfusg",
        exe_name: Union[str, PathLike] = "mfusg",
        verbose=False,
        model_ws: Union[str, PathLike] = curdir,
        load_only=None,
        forgive=False,
        check=True,
    ):
        """Load an existing MODFLOW-USG model.

        Parameters
        ----------
        f : str
            Name of MODFLOW name file to load.
        version : str, default "mfusg"
            MODFLOW version. Must be "mfusg".
        exe_name : str, default "mfusg"
            MODFLOW executable name.
        verbose : bool, default False
            Show messages that can be useful for debugging.
        model_ws : str or PathLike, default "." (curdir)
            Model workspace path. Default is the current directory.
        load_only : list, str or None
            List of case insensitive packages to load, e.g. ["bas6", "lpf"].
            One package can also be specified, e.g. "rch". Default is None,
            which attempts to load all files. An empty list [] will not load
            any additional packages than is necessary. At a minimum, "dis" or
            "disu" is always loaded.
        forgive : bool, optional, default False
            Option to raise exceptions on package load failure, which can be
            useful for debugging. Default False.
        check : boolean, optional
            Check model input for common errors. Default True.

        Returns
        -------
        flopy.mfusg.MfUsg

        Examples
        --------
        >>> import flopy
        >>> ml = flopy.mfusg.MfUsg.load('model.nam')
        """
        if version != "mfusg":
            version = "mfusg"

        # similar to modflow command: if file does not exist , try file.nam
        namefile_path = os.path.join(model_ws, f)
        if not os.path.isfile(namefile_path) and os.path.isfile(f"{namefile_path}.nam"):
            namefile_path += ".nam"
        if not os.path.isfile(namefile_path):
            raise OSError(f"cannot find name file: {namefile_path}")

        # Determine model name from 'f', without any extension or path
        modelname = os.path.splitext(os.path.basename(f))[0]

        if verbose:
            print(f"\nCreating new model with name: {modelname}\n{50 * '-'}\n")

        attribs = mfreadnam.attribs_from_namfile_header(os.path.join(model_ws, f))

        model = cls(
            modelname,
            exe_name=exe_name,
            verbose=verbose,
            model_ws=model_ws,
            **attribs,
        )

        # read name file
        ext_unit_dict = mfreadnam.parsenamefile(
            namefile_path, model.mfnam_packages, verbose=verbose
        )
        if model.verbose:
            print(
                f"\n{50 * '-'}\nExternal unit dictionary:\n"
                f"{ext_unit_dict}\n{50 * '-'}\n"
            )

        # create a dict where key is the package name, value is unitnumber
        ext_pkg_d = {v.filetype: k for (k, v) in ext_unit_dict.items()}

        # reset version based on packages in the name file
        if "DISU" in ext_pkg_d:
            model.structured = False

        if "DPF" in ext_pkg_d:
            model.idpf = 1

        if "CLN" in ext_pkg_d:
            model.icln = 1

        # reset unit number for list file
        if "LIST" in ext_pkg_d:
            unitnumber = ext_pkg_d["LIST"]
            filepth = os.path.basename(ext_unit_dict[unitnumber].filename)
            model.lst.unit_number = [unitnumber]
            model.lst.file_name = [filepth]

        # look for the free format flag in bas6
        bas_key = ext_pkg_d.get("BAS6")
        if bas_key is not None:
            bas = ext_unit_dict[bas_key]
            start = bas.filehandle.tell()
            line = bas.filehandle.readline()
            while line.startswith("#"):
                line = bas.filehandle.readline()
            if "FREE" in line.upper():
                model.free_format_input = True
            bas.filehandle.seek(start)
        if verbose:
            print(f"ModflowBas6 free format:{model.free_format_input}\n")

        files_successfully_loaded, files_not_loaded = cls._load_packages(
            model, ext_unit_dict, ext_pkg_d, load_only, forgive
        )

        # set up binary output / external file units
        cls._set_output_external(model, ext_unit_dict)

        # send messages re: success/failure of loading
        cls._send_load_messages(model, files_successfully_loaded, files_not_loaded)

        if check:
            model.check(f=f"{model.name}.chk", verbose=model.verbose, level=0)

        # return model object
        return model

    @classmethod
    def _load_packages(cls, model, ext_unit_dict, ext_pkg_d, load_only, forgive):
        """
        Method to load packages into the MODFLOW-USG Model Class.
        For internal class use - should not be called by the user.

        Parameters
        ----------
        model : MfUsg model object
        ext_unit_dict : dict
            For each file listed in the name file, a
            :class:`flopy.utils.mfreadnam.NamData` instance.
            Keyed by unit number.
        ext_pkg_d : dict
            key is package name, value is unitnumber
        load_only : list, str or None
            List of case insensitive packages to load, e.g. ["bas6", "lpf"].
            One package can also be specified, e.g. "rch". Default is None,
            which attempts to load all files. An empty list [] will not load
            any additional packages than is necessary. At a minimum, "dis" or
            "disu" is always loaded.
        forgive : bool
            Option to raise exceptions on package load failure.

        Returns
        -------
        files_successfully_loaded : list of loaded files
        files_not_loaded : list of files that were not loaded
        """
        files_successfully_loaded = []
        files_not_loaded = []

        # load dis
        dis_key = ext_pkg_d.get("DIS") or ext_pkg_d.get("DISU")
        if dis_key is None:
            raise KeyError("discretization entry not found in nam file")
        disnamdata = ext_unit_dict[dis_key]
        dis = disnamdata.package.load(
            disnamdata.filehandle,
            model,
            ext_unit_dict=ext_unit_dict,
            check=False,
        )
        files_successfully_loaded.append(disnamdata.filename)
        if model.verbose:
            print(f"   {dis.name[0]:4s} package load...success")
        assert model.pop_key_list.pop() == dis_key
        ext_unit_dict.pop(dis_key).filehandle.close()

        dis.start_datetime = model.start_datetime

        # BCT has to be loaded before other transport packages for MFUSG-TRANSPORT
        bct_key = ext_pkg_d.get("BCT")
        if bct_key is not None:
            bctnamdata = ext_unit_dict[bct_key]
            bct = bctnamdata.package.load(
                bctnamdata.filehandle,
                model,
                ext_unit_dict=ext_unit_dict,
                check=False,
            )
            files_successfully_loaded.append(bctnamdata.filename)
            if model.verbose:
                print(f"   {bct.name[0]:4s} package load...success")
            ext_unit_dict.pop(bct_key).filehandle.close()

        # set mfpar / pval
        cls._set_mfpar_pval(model, ext_unit_dict, ext_pkg_d)

        load_only = cls._prepare_load_only(load_only, ext_pkg_d)

        # try loading packages in ext_unit_dict
        for key, item in ext_unit_dict.items():
            if item.package is not None:
                (files_successfully_loaded, files_not_loaded) = (
                    cls._load_ext_unit_dict_paks(
                        model,
                        ext_unit_dict,
                        load_only,
                        item,
                        forgive,
                        files_successfully_loaded,
                        files_not_loaded,
                    )
                )
            elif "data" not in item.filetype.lower():
                files_not_loaded.append(item.filename)
                if model.verbose:
                    print(f"   {item.filetype:4s} package load...skipped")
            elif "data" in item.filetype.lower():
                cls._prepare_external_files(model, key, item)
            else:
                raise KeyError(f"unhandled case: {key}, {item}")

        return files_successfully_loaded, files_not_loaded

    @staticmethod
    def _prepare_load_only(load_only, ext_pkg_d):
        """Prepare load_only list."""
        if load_only is None:
            # load all packages/files
            load_only = ext_pkg_d.keys()
        else:  # check items in list
            if not isinstance(load_only, list):
                load_only = [load_only]
            not_found = []
            for idx, filetype in enumerate(load_only):
                load_only[idx] = filetype = filetype.upper()
                if filetype not in ext_pkg_d:
                    not_found.append(filetype)
            if not_found:
                raise KeyError(
                    "the following load_only entries were not found "
                    "in the ext_unit_dict: " + str(not_found)
                )
        return load_only

    @classmethod
    def _load_ext_unit_dict_paks(
        cls,
        model,
        ext_unit_dict,
        load_only,
        item,
        forgive,
        files_successfully_loaded,
        files_not_loaded,
    ):
        """Load packages from ext_unit_dict."""
        if item.filetype in load_only:
            if forgive:
                try:
                    cls._ext_unit_d_load(model, ext_unit_dict, item)
                    files_successfully_loaded.append(item.filename)
                    if model.verbose:
                        print(
                            f"   {item.filetype:4s} package \
                            load...success"
                        )
                except PackageLoadException as err:
                    model.load_fail = True
                    if model.verbose:
                        raise PackageLoadException(
                            error=f"   {item.filetype:4s} package \
                            load...failed"
                        ) from err
                    files_not_loaded.append(item.filename)
            else:
                cls._ext_unit_d_load(model, ext_unit_dict, item)
                files_successfully_loaded.append(item.filename)
                if model.verbose:
                    print(f"   {item.filetype:4s} package load...success")
        else:
            if model.verbose:
                print(f"   {item.filetype:4s} package load...skipped")
            files_not_loaded.append(item.filename)

        return files_successfully_loaded, files_not_loaded

    @staticmethod
    def _prepare_external_files(model, key, item):
        """Prepare external files for ext_unit_dict item."""
        if model.verbose:
            print(f"   {item.filetype} package load...skipped")
            print(f"      {os.path.basename(item.filename)}")
        if key not in model.pop_key_list:
            # do not add unit number (key) if it already exists
            if key not in model.external_units:
                model.external_fnames.append(item.filename)
                model.external_units.append(key)
                model.external_binflag.append("binary" in item.filetype.lower())
                model.external_output.append(False)

    @staticmethod
    def _ext_unit_d_load(model, ext_unit_dict, ext_unit_d_item):
        """
        Method to load a package from an ext_unit_dict item into model

        Parameters
        ----------
        model : MfUsg model object for which package in ext_unit_d_item will
            be loaded
        ext_unit_dict : dict
            For each file listed in the name file, a
            :class:`flopy.utils.mfreadnam.NamData` instance.
            Keyed by unit number.
        ext_unit_d_item : :class:`flopy.utils.mfreadnam.NamData` instance.
            Must be an item of ext_unit_dict.
        """
        package_load_args = getfullargspec(ext_unit_d_item.package.load)[0]
        if "check" in package_load_args:
            ext_unit_d_item.package.load(
                ext_unit_d_item.filehandle,
                model,
                ext_unit_dict=ext_unit_dict,
                check=False,
            )
        else:
            ext_unit_d_item.package.load(
                ext_unit_d_item.filehandle,
                model,
                ext_unit_dict=ext_unit_dict,
            )

    @staticmethod
    def _set_mfpar_pval(model, ext_unit_dict, ext_pkg_d):
        """Set mfpar/pval items."""
        # zone, mult, pval
        if "PVAL" in ext_pkg_d:
            model.mfpar.set_pval(model, ext_unit_dict)
            assert model.pop_key_list.pop() == ext_pkg_d.get("PVAL")
        if "ZONE" in ext_pkg_d:
            model.mfpar.set_zone(model, ext_unit_dict)
            assert model.pop_key_list.pop() == ext_pkg_d.get("ZONE")
        if "MULT" in ext_pkg_d:
            model.mfpar.set_mult(model, ext_unit_dict)
            assert model.pop_key_list.pop() == ext_pkg_d.get("MULT")

    @staticmethod
    def _set_output_external(model, ext_unit_dict):
        """Set up binary output / external file units."""
        # pop binary output keys and any external file units that are now
        # internal
        for key in model.pop_key_list:
            try:
                model.remove_external(unit=key)
                item = ext_unit_dict.pop(key)
                if hasattr(item.filehandle, "close"):
                    item.filehandle.close()
            except KeyError:
                if model.verbose:
                    print(
                        f"\nWARNING:\n    External file unit {key} does not "
                        "exist in ext_unit_dict."
                    )

    @staticmethod
    def _send_load_messages(model, files_successfully_loaded, files_not_loaded):
        """Send messages re: success/failure of loading."""
        # write message indicating packages that were successfully loaded
        if model.verbose:
            print("")
            print(
                f"   The following {len(files_successfully_loaded)} packages "
                "were successfully loaded."
            )
            for fname in files_successfully_loaded:
                print(f"      {os.path.basename(fname)}")
            if len(files_not_loaded) > 0:
                print(
                    f"   The following {len(files_not_loaded)} packages "
                    "were not loaded."
                )
                for fname in files_not_loaded:
                    print(f"      {os.path.basename(fname)}")


def fmt_string(array):
    """
    Returns a C-style fmt string for numpy savetxt that corresponds to
    the dtype.

    Parameters
    ----------
    array : numpy array
    """
    fmts = []
    for field in array.dtype.descr:
        vtype = field[1][1].lower()
        if vtype in {"i", "b"}:
            fmts.append("%10d")
        elif vtype == "f":
            fmts.append("%10.2e")
        elif vtype == "o":
            fmts.append("%10s")
        elif vtype == "s":
            msg = (
                "mfusg.fmt_string error: 'str' type found in dtype."
                "This gives unpredictable results when"
                "recarray to file - change to 'object' type"
            )
            raise TypeError(msg)
        else:
            raise TypeError(f"mfusg.fmt_string error: unknown vtype in field: {field}")
    return "".join(fmts)
