"""
mflak module.  Contains the ModflowLak class. Note that the user can access
the ModflowLak class as `flopy.modflow.ModflowLak`.

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

"""
import os
import sys
import numpy as np
from flopy import version as vs
if vs.__version__ >= '3.2.4':
    from flopy.pakbase import Package
if vs.__version__ == '3.2.3':
    from flopy.mbase import Package
from util_array import Transient3d, Util3d


class ModflowLak(Package):
    """
    MODFLOW Lake Package Class.

    Parameters
    ----------
    model : model object
        The model object (of type :class:`flopy.modflow.mf.Modflow`) to which
        this package will be added.
    options : list of strings
        Package options. (default is None).
    extension : string
        Filename extension (default is 'str')
    unitnumber : int
        File unit number (default is 118).

    Methods
    -------

    See Also
    --------

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

    Examples
    --------

    >>> import flopy
    >>> m = flopy.modflow.Modflow()
    >>> lak = {}
    >>> lak[0] = [[2, 3, 4, 15.6, 1050., -4]]  #this lake boundary will be
    >>>                                        #applied to all stress periods
    >>> lak = flopy.modflow.ModflowLak(m, nstress_period_data=strd)

    """

    def __init__(self, model, nlakes=1, ipakcb=0, theta=-1.,
                 nssitr=0, sscncr=0.0, surfdep=0., stages=1., stage_range=None,
                 tab_files=None, lakarr=None, bdlknc=None,
                 sill_data=None, flux_data=None,
                 extension='lak', unitnumber=119, options=None, lwrt=0, **kwargs):
        """
        Package constructor.

        """
        # Call parent init to set self.parent, extension, name and unit number
        Package.__init__(self, model, extension, 'LAK', unitnumber)

        self.heading = '# {} package for '.format(self.name[0]) + \
                       ' {}, generated by Flopy.'.format(model.version)
        self.url = 'lak.htm'

        if options is None:
            options = []
        self.options = options
        self.nlakes = nlakes
        if ipakcb != 0:
            self.ipakcb = 53
        else:
            self.ipakcb = 0  # 0: no cell by cell terms are written
        self.theta = theta
        self.nssitr = nssitr
        self.sscncr = sscncr
        self.surfdep = surfdep
        self.lwrt = lwrt
        if isinstance(stages, float):
            stages = np.array(self.nlakes, dtype=np.float) * stages
        elif isinstance(stages, list):
            stages = np.array(stages)
        if stages.shape[0] != nlakes:
            err = 'stages shape should be ' + \
                  '({}) but is only ({}).'.format(nlakes, stages.shape[0])
            raise Exception(err)
        self.stages = stages
        if stage_range is None:
            stage_range = np.ones((nlakes, 2), dtype=np.float)
            stage_range[:, 0] = -10000.
            stage_range[:, 1] = 10000.
        else:
            if isinstance(stage_range, list):
                stage_range = np.array(stage_range)
            elif isinstance(stage_range, float):
                err = 'stage_range should be a list or ' + \
                      'array of size ({}, 2)'.format(nlakes)
                raise Exception(err)
        if self.parent.dis.steady[0]:
            if stage_range.shape != (nlakes, 2):
                err = 'stages shape should be ' + \
                      '({},2) but is only {}.'.format(nlakes,
                                                      stage_range.shape)
                raise Exception(err)
        self.stage_range = stage_range

        self.tabdata = False
        iunit_tab = None
        for option in self.options:
            if 'TABLEINPUT' in option.upper():
                if tab_files is None:
                    err = "tab_files must be specified " + \
                          "if 'TABLEIPUT' option specified"
                    raise Exception(err)
                elif isinstance(tab_files, np.ndarray):
                    tab_files = tab_files.aslist()
                elif isinstance(tab_files, str):
                    tab_files = [tab_files]
                if len(tab_files) != self.nlakes:
                    err = "an external filename for tabular lake data" + \
                          "must be specified for each " + \
                          "lake ({}) ".format(self.nlakes) + \
                          "only {} were provided".format(len(tab_files))
                    raise Exception(err)
                self.tabdata = True
                iunit_tab = []
                for tab_file in tab_files:
                    nextunit = self.parent.next_ext_unit()
                    iunit_tab.append(nextunit)
                    self.parent.add_external(tab_file, nextunit)

        self.iunit_tab = iunit_tab

        if lakarr is None and bdlknc is None:
            err = 'lakarr and bdlknc must be specified'
            raise Exception(err)
        nrow, ncol, nlay, nper = self.parent.get_nrow_ncol_nlay_nper()

        self.lakarr = Transient3d(model, (nlay, nrow, ncol), np.int,
                                  lakarr, name='lakarr_')
        self.bdlknc = Transient3d(model, (nlay, nrow, ncol), np.float32,
                                  bdlknc, name='bdlknc_')

        if sill_data is not None:
            #if not isinstance(sill_data, dict):
                try:
                    sill_data = {0: sill_data}
                except:
                    err = 'sill_data must be a dictionary'
                    raise Exception(err)

        if flux_data is not None:
            if not isinstance(flux_data, dict):
                # convert array to a dictionary
                try:
                    flux_data = {0: flux_data}
                except:
                    err = 'flux_data must be a dictionary'
                    raise Exception(err)
            for key, value in flux_data.items():
                if isinstance(value, np.ndarray):
                    td = {}
                    for k in range(value.shape[0]):
                        td[k] = value[k, :].aslist()
                    flux_data[key] = td
                    if len(list(flux_data.keys())) != nlakes:
                        err = 'flux_data dictionary must ' + \
                              'have {} entries'.format(nlakes)
                        raise Exception(err)
                elif isinstance(value, float) or \
                        isinstance(value, int):
                    td = {}
                    for k in range(self.nlakes):
                        td[k] = (np.ones(6, dtype=np.float) * value).aslist()
                    flux_data[key] = td
                elif isinstance(value, dict):
                    try:
                        steady = self.parent.dis.steady[key]
                    except:
                        steady = True
                    nlen = 4
                    if steady and key > 0:
                        nlen = 6
                    for k in range(self.nlakes):
                        td = value[k]
                        if len(td) < nlen:
                            err = 'flux_data entry for stress period'.format(key + 1) + \
                                  'has {} entries but '.format(nlen) + \
                                  'should have {} entries'.format(len(td))
                            raise Exception(err)

        self.flux_data = flux_data
        self.sill_data = sill_data

        self.parent.add_package(self)

        return

    def ncells(self):
        # Return the  maximum number of cells that have a stream
        # (developed for MT3DMS SSM package)
        nrow, ncol, nlay, nper = self.parent.nrow_ncol_nlay_nper
        return (nlay * nrow * ncol)

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

        Returns
        -------
        None

        """
        f = open(self.fn_path, 'w')
        # dataset 0
        self.heading = '# {} package for '.format(self.name[0]) + \
                       '{}, generated by Flopy.'.format(self.parent.version)
        f.write('{0}\n'.format(self.heading))

        # dataset 1a
        if len(self.options) > 0:
            for option in self.options:
                f.write('{} '.format(option))
            f.write('\n')

        # dataset 1b
        f.write(write_fixed_var([self.nlakes, self.ipakcb],
                                free=self.parent.array_free_format, comment="Data Set 1b"))
        # dataset 2
        steady = np.any(self.parent.dis.steady.array)
        t = [self.theta]
        if self.theta < 0. or steady:
            t.append(self.nssitr)
            t.append(self.sscncr)
        if self.theta < 0.:
            t.append(self.surfdep)
        f.write(write_fixed_var(t, free=self.parent.array_free_format, comment="Data Set 2"))

        # dataset 3
        steady = self.parent.dis.steady[0]
        for n in range(self.nlakes):
            ipos = [10]
            t = [self.stages[n]]
            if steady:
                ipos.append(10)
                t.append(self.stage_range[n, 0])
                ipos.append(10)
                t.append(self.stage_range[n, 1])
            if self.tabdata:
                ipos.append(5)
                t.append(self.iunit_tab[n])
            f.write(write_fixed_var(t, ipos=ipos,
                                    free=self.parent.array_free_format, comment="Data Set 3"))

        ds8_keys = list(self.sill_data.keys())
        ds9_keys = list(self.flux_data.keys())
        nper = self.parent.dis.steady.shape[0]
        for kper in range(nper):
            itmp, file_entry_lakarr = self.lakarr.get_kper_entry(kper)
            ibd, file_entry_bdlknc = self.bdlknc.get_kper_entry(kper)

            itmp2 = 0
            if kper in ds9_keys:
                itmp2 = 1

            t = [itmp, itmp2, self.lwrt]
            comment = 'Stress period {} Data Set 4'.format(kper + 1)
            f.write(write_fixed_var(t, free=self.parent.free_format,
                                    comment=comment))

            if itmp > 0:

                f.write(file_entry_lakarr)
                f.write(file_entry_bdlknc)

                nslms = 0
                if kper in ds8_keys:
                    ds8 = self.sill_data[kper]
                    nslms = len(ds8)

                f.write(write_fixed_var([nslms],
                                        free=self.parent.free_format,
                                        comment='Data set 7'))
                if nslms > 0:
                    for n in range(nslms):
                        d1, d2 = ds8[n]
                        s = write_fixed_var(d1,
                                            free=self.parent.free_format,
                                            comment='Data set 8a')
                        f.write(s)
                        s = write_fixed_var(d2,
                                            free=self.parent.free_format,
                                            comment='Data set 8b')
                        f.write(s)

            if itmp2 > 0:
                ds9 = self.flux_data[kper]
                for n in range(self.nlakes):
                    try:
                        steady = self.parent.dis.steady[kper]
                    except:
                        steady = True
                    if kper > 0 and steady:
                        t = ds9[n]
                    else:
                        t = ds9[n][0:4]
                    s = write_fixed_var(t,
                                        free=self.parent.free_format,
                                        length=20,
                                        comment='Data set 9a')
                    f.write(s)

        # close the lak file
        f.close()

    @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.
        nper : int
            The number of stress periods.  If nper is None, then nper will be
            obtained from the model object. (default is None).
        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
        -------
        str : ModflowStr object
            ModflowStr object.

        Examples
        --------

        >>> import flopy
        >>> m = flopy.modflow.Modflow()
        >>> lak = flopy.modflow.ModflowStr.load('test.lak', m)

        """

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

        if not hasattr(f, 'read'):
            filename = f
            if sys.version_info[0] == 2:
                f = open(filename, 'r')
            elif sys.version_info[0] == 3:
                f = open(filename, 'r', errors='replace')

        # dataset 0 -- header
        while True:
            line = f.readline()
            if line[0] != '#':
                break

        options = []
        tabdata = False
        if 'TABLEINPUT' in line.upper():
            if model.verbose:
                print("   reading lak dataset 1a")
            options.append('TABLEINPUT')
            tabdata = True
            line = f.readline()

        # read dataset 1b
        if model.verbose:
            print("   reading lak dataset 1b")
        t = line.strip().split()
        nlakes = int(t[0])
        ipakcb = 0
        try:
            ipakcb1 = int(t[1])
            if ipakcb1 != 0:
                ipakcb = 53
                model.add_pop_key_list(ipakcb1)
        except:
            pass

        # read dataset 2
        line = f.readline().rstrip()
        if model.free_format:
            t = line.split()
        else:
            t = read_fixed_var(line, ncol=4)
        theta = float(t[0])
        nssitr, sscncr = 0, 0.
        try:
            nssitr = int(t[1])
        except:
            pass
        try:
            sscncr = float(t[2])
        except:
            pass
        surfdep = 0.
        if theta < 0.:
            surfdep = float(t[3])

        if nper is None:
            nrow, ncol, nlay, nper = model.get_nrow_ncol_nlay_nper()

        if model.verbose:
            print("   reading lak dataset 3")
        stages = []
        stage_range = []
        if tabdata:
            tab_files = []
        else:
            tab_files = None
        for lake in range(nlakes):
            line = f.readline().rstrip()
            if model.free_format:
                t = line.split()
            else:
                t = read_fixed_var(line, ipos=[10, 10, 10, 5])
            stages.append(t[0])
            ipos = 1
            if model.dis.steady[0]:
                stage_range.append((float(t[ipos]), float(t[ipos + 1])))
                ipos += 2
            if tabdata:
                iu = int(t[ipos])
                for key, value in ext_unit_dict.items():
                    if key == iu:
                        model.add_pop_key_list(iu)
                        relpth = os.path.relpath(value.filename,
                                                 model.model_ws)
                        tab_files.append(relpth)
                        break

        lake_loc = {}
        lake_lknc = {}
        sill_data = {}
        flux_data = {}
        for iper in range(nper):
            if model.verbose:
                print("   reading lak dataset 4 - " +
                      "for stress period {}".format(iper + 1))
            line = f.readline().rstrip()
            if model.free_format:
                t = line.split()
            else:
                t = read_fixed_var(line, ncol=3)
            itmp, itmp1, lwrt = int(t[0]), int(t[1]), int(t[2])

            if itmp > 0:
                if model.verbose:
                    print("   reading lak dataset 5 - " +
                          "for stress period {}".format(iper + 1))
                name = 'LKARR_StressPeriod_{}'.format(iper)
                lakarr = Util3d.load(f, model, (nlay, nrow, ncol), np.int,
                                     name, ext_unit_dict)
                if model.verbose:
                    print("   reading lak dataset 6 - " +
                          "for stress period {}".format(iper + 1))
                name = 'BDLKNC_StressPeriod_{}'.format(iper)
                bdlknc = Util3d.load(f, model, (nlay, nrow, ncol), np.float32,
                                     name, ext_unit_dict)

                lake_loc[iper] = lakarr
                lake_lknc[iper] = bdlknc

                if model.verbose:
                    print("   reading lak dataset 7 - " +
                          "for stress period {}".format(iper + 1))
                line = f.readline().rstrip()
                t = line.split()
                nslms = int(t[0])
                ds8 = []
                if nslms > 0:
                    if model.verbose:
                        print("   reading lak dataset 8 - " +
                              "for stress period {}".format(iper + 1))
                    for i in range(nslms):
                        line = f.readline().rstrip()
                        if model.free_format:
                            t = line.split()
                        else:
                            ic = int(line[0:5])
                            t = read_fixed_var(line, ncol=ic + 1, length=5)
                        ic = int(t[0])
                        ds8a = [ic]
                        for j in range(1, ic + 1):
                            ds8a.append(int(t[j]))
                        line = f.readline().rstrip()
                        if model.free_format:
                            t = line.split()
                        else:
                            t = read_fixed_var(line, ncol=ic - 1)
                        silvt = []
                        for j in range(ic - 1):
                            silvt.append(float(t[j]))
                        ds8.append((ds8a, silvt))
                    sill_data[iper] = ds8
            if itmp1 >= 0:
                if model.verbose:
                    print("   reading lak dataset 9 - " +
                          "for stress period {}".format(iper + 1))
                ds9 = {}
                for n in range(nlakes):
                    line = f.readline().rstrip()
                    if model.free_format:
                        t = line.split()
                    else:
                        t = read_fixed_var(line, ncol=6)
                    tds = []
                    tds.append(float(t[0]))
                    tds.append(float(t[1]))
                    tds.append(float(t[2]))
                    tds.append(float(t[3]))
                    if model.dis.steady[iper]:
                        if iper == 0:
                            tds.append(stage_range[n][0])
                            tds.append(stage_range[n][1])
                        else:
                            tds.append(float(t[4]))
                            tds.append(float(t[5]))
                    else:
                        tds.append(0.)
                        tds.append(0.)
                    ds9[n] = tds
                flux_data[iper] = ds9

        # convert lake data to Transient3d objects
        lake_loc = Transient3d(model, (nlay, nrow, ncol), np.int,
                               lake_loc, name='lakarr_')
        lake_lknc = Transient3d(model, (nlay, nrow, ncol), np.float32,
                                lake_lknc, name='bdlknc_')

        lakpak = ModflowLak(model, options=options, nlakes=nlakes,
                            ipakcb=ipakcb, theta=theta, nssitr=nssitr,
                            surfdep=surfdep, sscncr=sscncr, stages=stages,
                            stage_range=stage_range, tab_files=tab_files,
                            lakarr=lake_loc, bdlknc=lake_lknc,
                            sill_data=sill_data, flux_data=flux_data)
        return lakpak


def write_fixed_var(v, length=10, ipos=None, free=False, comment=None):
    """

    Parameters
    ----------
    v : list, int, float, bool, or numpy array
        list, int, float, bool, or numpy array containing the data to be
        written to a string.
    length : int
        length of each column for fixed column widths. (default is 10)
    ipos : list, int, or numpy array
        user-provided column widths. (default is None)
    free : bool
        boolean indicating if a free format string should be generated.
        length and ipos are not used if free is True. (default is False)
    comment : str
        comment string to add to the end of the string

    Returns
    -------
    out : str
        fixed or free format string generated using user-provided data

    """
    if isinstance(v, np.ndarray):
        v = v.aslist()
    elif isinstance(v, int) or isinstance(v, float) or isinstance(v, bool):
        v = [v]
    ncol = len(v)
    # construct ipos if it was not passed
    if ipos is None:
        ipos = []
        for i in range(ncol):
            ipos.append(length)
    else:
        if isinstance(ipos, np.ndarray):
            ipos = ipos.flatten().aslist()
        elif isinstance(ipos, int):
            ipos = [ipos]
        if len(ipos) < ncol:
            err = 'user provided ipos length ({})'.format(len(ipos)) + \
                  'should be greater than or equal ' + \
                  'to the length of v ({})'.format(ncol)
            raise Exception(err)
    out = ''
    for n in range(ncol):
        if free:
            write_fmt = '{} '
        else:
            write_fmt = '{{:>{}}}'.format(ipos[n])
        out += write_fmt.format(v[n])
    if comment is not None:
        out += '  # {}'.format(comment)
    out += '\n'
    return out


def read_fixed_var(line, ncol=1, length=10, ipos=None, free=False):
    """
    Parse a fixed format line using user provided data

    Parameters
    ----------
    line : str
        text string to parse.
    ncol : int
        number of columns to parse from line. (default is 1)
    length : int
        length of each column for fixed column widths. (default is 10)
    ipos : list, int, or numpy array
        user-provided column widths. (default is None)
    free : bool
        boolean indicating if sting is free format. ncol, length, and
        ipos are not used if free is True. (default is False)

    Returns
    -------
    out : list
        padded list containing data parsed from the passed text string

    """
    if free:
        out = line.rstrip().split()
    else:
        # construct ipos if it was not passed
        if ipos is None:
            ipos = []
            for i in range(ncol):
                ipos.append(length)
        else:
            if isinstance(ipos, np.ndarray):
                ipos = ipos.flatten().aslist()
            elif isinstance(ipos, int):
                ipos = [ipos]
            ncol = len(ipos)
        line = line.rstrip()
        out = []
        istart = 0
        for ivar in range(ncol):
            istop = istart + ipos[ivar]
            try:
                txt = line[istart:istop]
                if len(txt.strip()) > 0:
                    out.append(txt)
                else:
                    out.append(0)
            except:
                break
            istart = istop
    return out
