"""
mfdrn module.  Contains the ModflowDrn class. Note that the user can access
the ModflowDrn class as `flopy.modflow.ModflowDrn`.

Additional information for this MODFLOW package can be found at the `Online
MODFLOW Guide
<http://water.usgs.gov/ogw/modflow/MODFLOW-2005-Guide/index.html?drn.htm>`_.

"""
import sys
import numpy as np
from flopy.mbase import Package
from flopy.utils.util_list import mflist


class ModflowDrn(Package):
    """
    MODFLOW Drain Package Class.

    Parameters
    ----------
    model : model object
        The model object (of type :class:`flopy.modflow.mf.Modflow`) to which
        this package will be added.
    ipakcb : int
        is a flag and a unit number. (default is 0).
    stress_period_data : list of boundaries or
                         recarray of boundaries or
                         dictionary of boundaries
        Each drain cell is defined through definition of
        layer(int), row(int), column(int), elevation(float), conductance(float)
        The simplest form is a dictionary with a lists of boundaries for each
        stress period, where each list of boundaries itself is a list of
        boundaries. Indices of the dictionary are the numbers of the stress
        period. This gives the form of
            stress_period_data =
            {0: [
                [lay, row, col, stage, cond],
                [lay, row, col, stage, cond],
                [lay, row, col, stage, cond],
                ],
            1:  [
                [lay, row, col, stage, cond],
                [lay, row, col, stage, cond],
                [lay, row, col, stage, cond],
                ], ...
            kper:
                [
                [lay, row, col, stage, cond],
                [lay, row, col, stage, cond],
                [lay, row, col, stage, cond],
                ]
            }
        Note that if no values are specified for a certain stress period, then
        the list of boundaries for the previous stress period for which values
        were defined is used. Full details of all options to specify
        stress_period_data can be found in the flopy3boundaries Notebook in
        the basic subdirectory of the examples directory
    dtype : dtype definition
        if data type is different from default
    options : list of strings
        Package options. (default is None).
    extension : string
        Filename extension (default is 'drn')
    unitnumber : int
        File unit number (default is 21).

    Attributes
    ----------

    Methods
    -------

    See Also
    --------

    Notes
    -----
    Parameters are not supported in FloPy.

    Examples
    --------

    >>> import flopy
    >>> ml = flopy.modflow.Modflow()
    >>> lrcec = {0:[2, 3, 4, 10., 100.]}  #this drain will be applied to all
    >>>                                   #stress periods
    >>> drn = flopy.modflow.ModflowDrn(ml, stress_period_data=lrcec)

    """

    def __init__(self, model, ipakcb=0, stress_period_data=None, dtype=None,
                 extension='drn', unitnumber=21, options=None, **kwargs):
        """
        Package constructor
        """
        Package.__init__(self, model, extension, 'DRN',
                         unitnumber)  # Call ancestor's init to set self.parent, extension, name and unit number
        self.heading = '# DRN for MODFLOW, generated by Flopy.'
        self.url = 'drn.htm'
        self.ipakcb = ipakcb  # 0: no cell by cell terms are written

        self.np = 0
        if options is None:
            options = []
        self.options = options
        if dtype is not None:
            self.dtype = dtype
        else:
            self.dtype = self.get_default_dtype(structured=self.parent.structured)
        self.stress_period_data = mflist(self, stress_period_data)
        self.parent.add_package(self)

    def __repr__(self):
        return 'Drain class'

    @staticmethod
    def get_default_dtype(structured=True):
        if structured:
            dtype = np.dtype([("k", np.int), ("i", np.int),
                              ("j", np.int), ("elev", np.float32),
                              ("cond", np.float32)])
        else:
            dtype = np.dtype([("node", np.int), ("elev", np.float32),
                              ("cond", np.float32)])
        return dtype

    def ncells(self):
        # Returns the  maximum number of cells that have drains (developed for MT3DMS SSM package)
        # print 'Function must be implemented properly for drn package'
        return self.stress_period_data.mxact

    def write_file(self):
        """
        Write the file.

        """
        f_drn = open(self.fn_path, 'w')
        f_drn.write('{0}\n'.format(self.heading))
        # f_drn.write('%10i%10i\n' % (self.mxactd, self.idrncb))
        line = '{0:10d}{1:10d}'.format(self.stress_period_data.mxact, self.ipakcb)
        for opt in self.options:
            line += ' ' + str(opt)
        line += '\n'
        f_drn.write(line)
        self.stress_period_data.write_transient(f_drn)
        f_drn.close()

    def add_record(self, kper, index, values):
        try:
            self.stress_period_data.add_record(kper, index, values)
        except Exception as e:
            raise Exception("mfdrn error adding record to list: " + str(e))


    @staticmethod
    def get_empty(ncells=0, aux_names=None, structured=True):
        # get an empty recaray that correponds to dtype
        dtype = ModflowDrn.get_default_dtype(structured=structured)
        if aux_names is not None:
            dtype = Package.add_to_dtype(dtype, aux_names, np.float32)
        d = np.zeros((ncells, len(dtype)), dtype=dtype)
        d[:, :] = -1.0E+10
        return np.core.records.fromarrays(d.transpose(), dtype=dtype)

    @staticmethod
    def load(f, model, nper=None, ext_unit_dict=None):
        """
        Load an existing package.

        Parameters
        ----------
        f : filename or file handle
            File to load.
        model : model object
            The model object (of type :class:`flopy.modflow.mf.Modflow`) to
            which this package will be added.
        ext_unit_dict : dictionary, optional
            If the arrays in the file are specified using EXTERNAL,
            or older style array control records, then `f` should be a file
            handle.  In this case ext_unit_dict is required, which can be
            constructed using the function
            :class:`flopy.utils.mfreadnam.parsenamefile`.

        Returns
        -------
        drn : ModflowDrn object
            ModflowDrn object.

        Examples
        --------

        >>> import flopy
        >>> m = flopy.modflow.Modflow()
        >>> drn = flopy.modflow.ModflowDrn.load('test.drn', m)

        """

        if model.verbose:
            sys.stdout.write('loading drn package file...\n')

        return Package.load(model, ModflowDrn, f, nper)
