import os
import copy
import numpy as np
from .binaryfile import CellBudgetFile
from itertools import groupby
from collections import OrderedDict
from ..utils.utils_def import totim_to_datetime


class ZoneBudget(object):
    """
    ZoneBudget class

    Parameters
    ----------
    cbc_file : str or CellBudgetFile object
        The file name or CellBudgetFile object for which budgets will be
        computed.
    z : ndarray
        The array containing to zones to be used.
    kstpkper : tuple of ints
        A tuple containing the time step and stress period (kstp, kper).
        The kstp and kper values are zero based.
    totim : float
        The simulation time.
    aliases : dict
        A dictionary with key, value pairs of zones and aliases. Replaces
        the corresponding record and field names with the aliases provided.
        When using this option in conjunction with a list of zones, the
        zone(s) passed may either be all strings (aliases), all integers,
        or mixed.

    Returns
    -------
    None

    Examples
    --------

    >>> from flopy.utils.zonbud import ZoneBudget, read_zbarray
    >>> zon = read_zbarray('zone_input_file')
    >>> zb = ZoneBudget('zonebudtest.cbc', zon, kstpkper=(0, 0))
    >>> zb.to_csv('zonebudtest.csv')
    >>> zb_mgd = zb * 7.48052 / 1000000
    """

    def __init__(self, cbc_file, z, kstpkper=None, totim=None, aliases=None,
                 verbose=False, **kwargs):

        if isinstance(cbc_file, CellBudgetFile):
            self.cbc = cbc_file
        elif isinstance(cbc_file, str) and os.path.isfile(cbc_file):
            self.cbc = CellBudgetFile(cbc_file)
        else:
            raise Exception(
                'Cannot load cell budget file: {}.'.format(cbc_file))

        if isinstance(z, np.ndarray):
            assert np.issubdtype(z.dtype,
                                 np.integer), 'Zones dtype must be integer'
        else:
            e = 'Please pass zones as a numpy ndarray of (positive)' \
                ' integers. {}'.format(z.dtype)
            raise Exception(e)

        # Check for negative zone values
        if np.any(z < 0):
            raise Exception('Negative zone value(s) found:',
                            np.unique(z[z < 0]))

        self.dis = None
        self.sr = None
        if 'model' in kwargs.keys():
            self.model = kwargs.pop('model')
            self.sr = self.model.sr
            self.dis = self.model.dis
        if 'dis' in kwargs.keys():
            self.dis = kwargs.pop('dis')
            self.sr = self.dis.parent.sr
        if 'sr' in kwargs.keys():
            self.sr = kwargs.pop('sr')
        if len(kwargs.keys()) > 0:
            args = ','.join(kwargs.keys())
            raise Exception('LayerFile error: unrecognized kwargs: ' + args)

        # Check the shape of the cbc budget file arrays
        self.cbc_shape = self.cbc.get_data(idx=0, full3D=True)[0].shape
        self.nlay, self.nrow, self.ncol = self.cbc_shape
        self.cbc_times = self.cbc.get_times()
        self.cbc_kstpkper = self.cbc.get_kstpkper()
        self.kstpkper = None
        self.totim = None

        if kstpkper is not None:
            if isinstance(kstpkper, tuple):
                kstpkper = [kstpkper]
            for kk in kstpkper:
                s = 'The specified time step/stress period ' \
                    'does not exist {}'.format(kk)
                assert kk in self.cbc.get_kstpkper(), s
            self.kstpkper = kstpkper
        elif totim is not None:
            if isinstance(totim, float):
                totim = [totim]
            elif isinstance(totim, int):
                totim = [float(totim)]
            for t in totim:
                s = 'The specified simulation time ' \
                    'does not exist {}'.format(t)
                assert t in self.cbc.get_times(), s
            self.totim = totim
        else:
            # No time step/stress period or simulation time pass
            self.kstpkper = self.cbc.get_kstpkper()

        # Set float and integer types
        self.float_type = np.float32
        self.int_type = np.int32

        # Check dimensions of input zone array
        s = 'Row/col dimensions of zone array {}' \
            ' do not match model row/col dimensions {}'.format(z.shape,
                                                               self.cbc_shape)
        assert z.shape[-2] == self.nrow and \
               z.shape[-1] == self.ncol, s

        if z.shape == self.cbc_shape:
            izone = z.copy()
        elif len(z.shape) == 2:
            izone = np.zeros(self.cbc_shape, self.int_type)
            izone[:] = z[:, :]
        elif len(z.shape) == 3 and z.shape[0] == 1:
            izone = np.zeros(self.cbc_shape, self.int_type)
            izone[:] = z[0, :, :]
        else:
            e = 'Shape of the zone array is not recognized: {}'.format(
                z.shape)
            raise Exception(e)

        self.izone = izone
        self.allzones = np.unique(izone)
        self._zonenamedict = OrderedDict([(z, 'ZONE_{}'.format(z))
                                          for z in self.allzones])

        if aliases is not None:
            s = 'Input aliases not recognized. Please pass a dictionary ' \
                'with key,value pairs of zone/alias.'
            assert isinstance(aliases, dict), s
            # Replace the relevant field names (ignore zone 0)
            seen = []
            for z, a in iter(aliases.items()):
                if z != 0 and z in self._zonenamedict.keys():
                    if z in seen:
                        raise Exception(
                            'Zones may not have more than 1 alias.')
                    self._zonenamedict[z] = '_'.join(a.split())
                    seen.append(z)

        # self._iflow_recnames = self._get_internal_flow_record_names()

        # All record names in the cell-by-cell budget binary file
        self.record_names = [n.strip() for n in
                             self.cbc.get_unique_record_names(decode=True)]

        # Get imeth for each record in the CellBudgetFile record list
        self.imeth = {}
        for record in self.cbc.recordarray:
            self.imeth[record['text'].strip().decode("utf-8")] = record[
                'imeth']

        # INTERNAL FLOW TERMS ARE USED TO CALCULATE FLOW BETWEEN ZONES.
        # CONSTANT-HEAD TERMS ARE USED TO IDENTIFY WHERE CONSTANT-HEAD CELLS
        # ARE AND THEN USE FACE FLOWS TO DETERMINE THE AMOUNT OF FLOW.
        # SWIADDTO--- terms are used by the SWI2 groundwater flow process.
        internal_flow_terms = ['CONSTANT HEAD', 'FLOW RIGHT FACE',
                               'FLOW FRONT FACE', 'FLOW LOWER FACE',
                               'SWIADDTOCH', 'SWIADDTOFRF', 'SWIADDTOFFF',
                               'SWIADDTOFLF']

        # Source/sink/storage term record names
        # These are all of the terms that are not related to constant
        # head cells or face flow terms
        self.ssst_record_names = [n for n in self.record_names
                                  if n not in internal_flow_terms]

        # Initialize budget recordarray
        array_list = []
        if self.kstpkper is not None:
            for kk in self.kstpkper:
                recordarray = self._initialize_budget_recordarray(kstpkper=kk,
                                                                  totim=None)
                array_list.append(recordarray)
        elif self.totim is not None:
            for t in self.totim:
                recordarray = self._initialize_budget_recordarray(
                    kstpkper=None, totim=t)
                array_list.append(recordarray)
        self._budget = np.concatenate(array_list, axis=0)

        # Update budget record array
        if self.kstpkper is not None:
            for kk in self.kstpkper:
                if verbose:
                    s = 'Computing the budget for' \
                        ' time step {} in stress period {}'.format(kk[0] + 1,
                                                                   kk[1] + 1)
                    print(s)
                self._compute_budget(kstpkper=kk)
        elif self.totim is not None:
            for t in self.totim:
                if verbose:
                    s = 'Computing the budget for time {}'.format(t)
                    print(s)
                self._compute_budget(totim=t)

        return

    def get_model_shape(self):
        """

        Returns
        -------
        nlay : int
            Number of layers
        nrow : int
            Number of rows
        ncol : int
            Number of columns

        """
        return self.nlay, self.nrow, self.ncol

    def get_record_names(self, stripped=False):
        """
        Get a list of water budget record names in the file.

        Returns
        -------
        out : list of strings
            List of unique text names in the binary file.

        Examples
        --------

        >>> zb = ZoneBudget('zonebudtest.cbc', zon, kstpkper=(0, 0))
        >>> recnames = zb.get_record_names()

        """
        if not stripped:
            return np.unique(self._budget['name'])
        else:
            seen = []
            for recname in self.get_record_names():
                if recname in ['IN-OUT', 'TOTAL_IN', 'TOTAL_OUT']:
                    continue
                if recname.endswith('_IN'):
                    recname = recname[:-3]
                elif recname.endswith('_OUT'):
                    recname = recname[:-4]
                if recname not in seen:
                    seen.append(recname)
            seen.extend(['IN-OUT', 'TOTAL'])
            return np.array(seen)

    def get_budget(self, names=None, zones=None, net=False):
        """
        Get a list of zonebudget record arrays.

        Parameters
        ----------

        names : list of strings
            A list of strings containing the names of the records desired.
        zones : list of ints or strings
            A list of integer zone numbers or zone names desired.
        net : boolean
            If True, returns net IN-OUT for each record.

        Returns
        -------
        budget_list : list of record arrays
            A list of the zonebudget record arrays.

        Examples
        --------

        >>> names = ['FROM_CONSTANT_HEAD', 'RIVER_LEAKAGE_OUT']
        >>> zones = ['ZONE_1', 'ZONE_2']
        >>> zb = ZoneBudget('zonebudtest.cbc', zon, kstpkper=(0, 0))
        >>> bud = zb.get_budget(names=names, zones=zones)

        """
        if isinstance(names, str):
            names = [names]
        if isinstance(zones, str):
            zones = [zones]
        elif isinstance(zones, int):
            zones = [zones]
        select_fields = ['totim', 'time_step', 'stress_period',
                         'name'] + list(self._zonenamedict.values())
        select_records = np.where(
            (self._budget['name'] == self._budget['name']))
        if zones is not None:
            for idx, z in enumerate(zones):
                if isinstance(z, int):
                    zones[idx] = self._zonenamedict[z]
            select_fields = ['totim', 'time_step', 'stress_period',
                             'name'] + zones
        if names is not None:
            names = self._clean_budget_names(names)
            select_records = np.in1d(self._budget['name'], names)
        if net:
            if names is None:
                names = self._clean_budget_names(self.get_record_names())
            net_budget = self._compute_net_budget()
            seen = []
            net_names = []
            for name in names:
                iname = '_'.join(name.split('_')[1:])
                if iname not in seen:
                    seen.append(iname)
                else:
                    net_names.append(iname)
            select_records = np.in1d(net_budget['name'], net_names)
            return net_budget[select_fields][select_records]
        else:
            return self._budget[select_fields][select_records]

    def to_csv(self, fname):
        """
        Saves the budget record arrays to a formatted
        comma-separated values file.

        Parameters
        ----------
        fname : str
            The name of the output comma-separated values file.

        Returns
        -------
        None

        """
        # Needs updating to handle the new budget list structure. Write out
        # budgets for all kstpkper if kstpkper is None or pass list of
        # kstpkper/totim to save particular budgets.
        with open(fname, 'w') as f:
            # Write header
            f.write(','.join(self._budget.dtype.names) + '\n')
            # Write rows
            for rowidx in range(self._budget.shape[0]):
                s = ','.join(
                    [str(i) for i in list(self._budget[:][rowidx])]) + '\n'
                f.write(s)
        return

    def get_dataframes(self, start_datetime=None, timeunit='D',
                       index_key='totim', names=None, zones=None, net=False):
        """
        Get pandas dataframes.

        Parameters
        ----------

        start_datetime : str
            Datetime string indicating the time at which the simulation starts.
        timeunit : str
            String that indicates the time units used in the model.
        index_key : str
            Indicates the fields to be used (in addition to "record") in the
            resulting DataFrame multi-index.
        names : list of strings
            A list of strings containing the names of the records desired.
        zones : list of ints or strings
            A list of integer zone numbers or zone names desired.
        net : boolean
            If True, returns net IN-OUT for each record.

        Returns
        -------
        df : Pandas DataFrame
            Pandas DataFrame with the budget information.

        Examples
        --------
        >>> from flopy.utils.zonbud import ZoneBudget, read_zbarray
        >>> zon = read_zbarray('zone_input_file')
        >>> zb = ZoneBudget('zonebudtest.cbc', zon, kstpkper=(0, 0))
        >>> df = zb.get_dataframes()

        """
        try:
            import pandas as pd
        except Exception as e:
            msg = "ZoneBudget.get_dataframes() error import pandas: " + str(e)
            raise ImportError(msg)

        valid_index_keys = ['totim', 'kstpkper']
        s = 'index_key "{}" is not valid.'.format(index_key)
        assert index_key in valid_index_keys, s

        valid_timeunit = ['S', 'M', 'H', 'D', 'Y']

        if timeunit.upper() == 'SECONDS':
            timeunit = 'S'
        elif timeunit.upper() == 'MINUTES':
            timeunit = 'M'
        elif timeunit.upper() == 'HOURS':
            timeunit = 'H'
        elif timeunit.upper() == 'DAYS':
            timeunit = 'D'
        elif timeunit.upper() == 'YEARS':
            timeunit = 'Y'

        errmsg = 'Specified time units ({}) not recognized. ' \
                 'Please use one of '.format(timeunit)
        assert timeunit in valid_timeunit, errmsg + ', '.join(
            valid_timeunit) + '.'

        df = pd.DataFrame().from_records(self.get_budget(names, zones, net))
        if start_datetime is not None:
            totim = totim_to_datetime(df.totim,
                                      start=pd.to_datetime(start_datetime),
                                      timeunit=timeunit)
            df['datetime'] = totim
            index_cols = ['datetime', 'name']
        else:
            if index_key == 'totim':
                index_cols = ['totim', 'name']
            elif index_key == 'kstpkper':
                index_cols = ['time_step', 'stress_period', 'name']
        df = df.set_index(index_cols)  # .sort_index(level=0)
        if zones is not None:
            keep_cols = zones
        else:
            keep_cols = self._zonenamedict.values()
        return df.loc[:, keep_cols]

    def copy(self):
        """
        Return a deepcopy of the object.
        """
        return copy.deepcopy(self)

    def __deepcopy__(self, memo):
        """
        Over-rides the default deepcopy behavior. Copy all attributes except
        the CellBudgetFile object which does not copy nicely.
        """
        cls = self.__class__
        result = cls.__new__(cls)
        memo[id(self)] = result
        ignore_attrs = ['cbc']
        for k, v in self.__dict__.items():
            if k not in ignore_attrs:
                setattr(result, k, copy.deepcopy(v, memo))

        # Set CellBudgetFile object attribute manually. This is object
        # read-only so should not be problems with pointers from
        # multiple objects.
        result.cbc = self.cbc
        return result

    def _compute_budget(self, kstpkper=None, totim=None):
        """
        Creates a budget for the specified zone array. This function only
        supports the use of a single time step/stress period or time.

        Parameters
        ----------
        kstpkper : tuple
            Tuple of kstp and kper to compute budget for (default is None).
        totim : float
            Totim to compute budget for (default is None).

        Returns
        -------
        None

        """
        # Initialize an array to track where the constant head cells
        # are located.
        ich = np.zeros(self.cbc_shape, self.int_type)
        swiich = np.zeros(self.cbc_shape, self.int_type)

        if 'CONSTANT HEAD' in self.record_names:
            """
            C-----CONSTANT-HEAD FLOW -- DON'T ACCUMULATE THE CELL-BY-CELL VALUES FOR
            C-----CONSTANT-HEAD FLOW BECAUSE THEY MAY INCLUDE PARTIALLY CANCELING
            C-----INS AND OUTS.  USE CONSTANT-HEAD TERM TO IDENTIFY WHERE CONSTANT-
            C-----HEAD CELLS ARE AND THEN USE FACE FLOWS TO DETERMINE THE AMOUNT OF
            C-----FLOW.  STORE CONSTANT-HEAD LOCATIONS IN ICH ARRAY.
            """
            chd = self.cbc.get_data(text='CONSTANT HEAD', full3D=True,
                                    kstpkper=kstpkper, totim=totim)[0]
            ich[np.ma.where(chd != 0.)] = 1
        if 'FLOW RIGHT FACE' in self.record_names:
            self._accumulate_flow_frf('FLOW RIGHT FACE', ich, kstpkper, totim)
        if 'FLOW FRONT FACE' in self.record_names:
            self._accumulate_flow_fff('FLOW FRONT FACE', ich, kstpkper, totim)
        if 'FLOW LOWER FACE' in self.record_names:
            self._accumulate_flow_flf('FLOW LOWER FACE', ich, kstpkper, totim)
        if 'SWIADDTOCH' in self.record_names:
            swichd = self.cbc.get_data(text='SWIADDTOCH', full3D=True,
                                       kstpkper=kstpkper, totim=totim)[0]
            swiich[swichd != 0] = 1
        if 'SWIADDTOFRF' in self.record_names:
            self._accumulate_flow_frf('SWIADDTOFRF', swiich, kstpkper, totim)
        if 'SWIADDTOFFF' in self.record_names:
            self._accumulate_flow_fff('SWIADDTOFFF', swiich, kstpkper, totim)
        if 'SWIADDTOFLF' in self.record_names:
            self._accumulate_flow_flf('SWIADDTOFLF', swiich, kstpkper, totim)

        # NOT AN INTERNAL FLOW TERM, SO MUST BE A SOURCE TERM OR STORAGE
        # ACCUMULATE THE FLOW BY ZONE
        # iterate over remaining items in the list
        for recname in self.ssst_record_names:
            self._accumulate_flow_ssst(recname, kstpkper, totim)

        # Compute mass balance terms
        self._compute_mass_balance(kstpkper, totim)

        return

    # def _get_internal_flow_record_names(self):
    #     """
    #     Get internal flow record names
    #
    #     Returns
    #     -------
    #     iflow_recnames : np.recarray
    #         recarray of internal flow terms
    #
    #     """
    #     iflow_recnames = OrderedDict()
    #     for z, a in iter(self._zonenamedict.items()):
    #         iflow_recnames[z] = '{}'.format(a)
    #     dtype = np.dtype([('zone', '<i4'), ('name', (str, 50))])
    #     iflow_recnames = np.array(list(iflow_recnames.items()), dtype=dtype)
    #     return iflow_recnames

    def _add_empty_record(self, recordarray, recname, kstpkper=None,
                          totim=None):
        """
        Build an empty records based on the specified flow direction and
        record name for the given list of zones.

        Parameters
        ----------
        recordarray :
        recname :
        kstpkper : tuple
            Tuple of kstp and kper to compute budget for (default is None).
        totim : float
            Totim to compute budget for (default is None).

        Returns
        -------
        recordarray : np.recarray

        """
        if kstpkper is not None:
            if len(self.cbc_times) > 0:
                totim = self.cbc_times[self.cbc_kstpkper.index(kstpkper)]
            else:
                totim = 0.
        elif totim is not None:
            if len(self.cbc_times) > 0:
                kstpkper = self.cbc_kstpkper[self.cbc_times.index(totim)]
            else:
                kstpkper = (0, 0)

        row = [totim, kstpkper[0], kstpkper[1], recname]
        row += [0. for _ in self._zonenamedict.values()]
        recs = np.array(tuple(row), dtype=recordarray.dtype)
        recordarray = np.append(recordarray, recs)
        return recordarray

    def _initialize_budget_recordarray(self, kstpkper=None, totim=None):
        """
        Initialize the budget record array which will store all of the
        fluxes in the cell-budget file.

        Parameters
        ----------
        kstpkper : tuple
            Tuple of kstp and kper to compute budget for (default is None).
        totim : float
            Totim to compute budget for (default is None).

        Returns
        -------

        """

        # Create empty array for the budget terms.
        dtype_list = [('totim', '<f4'), ('time_step', '<i4'),
                      ('stress_period', '<i4'), ('name', (str, 50))]
        dtype_list += [(n, self.float_type) for n in
                       self._zonenamedict.values()]
        dtype = np.dtype(dtype_list)
        recordarray = np.array([], dtype=dtype)

        # Add "from" records
        if 'STORAGE' in self.record_names:
            recordarray = self._add_empty_record(recordarray, 'FROM_STORAGE',
                                                 kstpkper, totim)
        if 'CONSTANT HEAD' in self.record_names:
            recordarray = self._add_empty_record(recordarray,
                                                 'FROM_CONSTANT_HEAD',
                                                 kstpkper, totim)
        for recname in self.ssst_record_names:
            if recname != 'STORAGE':
                recordarray = self._add_empty_record(recordarray,
                                                     'FROM_' + '_'.join(
                                                         recname.split()),
                                                     kstpkper, totim)

        for z, n in self._zonenamedict.items():
            if z == 0 and 0 not in self.allzones:
                continue
            else:
                recordarray = self._add_empty_record(recordarray,
                                                     'FROM_' + '_'.join(
                                                         n.split()),
                                                     kstpkper, totim)
        recordarray = self._add_empty_record(recordarray, 'TOTAL_IN',
                                             kstpkper, totim)

        # Add "out" records
        if 'STORAGE' in self.record_names:
            recordarray = self._add_empty_record(recordarray, 'TO_STORAGE',
                                                 kstpkper, totim)
        if 'CONSTANT HEAD' in self.record_names:
            recordarray = self._add_empty_record(recordarray,
                                                 'TO_CONSTANT_HEAD',
                                                 kstpkper, totim)
        for recname in self.ssst_record_names:
            if recname != 'STORAGE':
                recordarray = self._add_empty_record(recordarray,
                                                     'TO_' + '_'.join(
                                                         recname.split()),
                                                     kstpkper, totim)

        for z, n in self._zonenamedict.items():
            if z == 0 and 0 not in self.allzones:
                continue
            else:
                recordarray = self._add_empty_record(recordarray,
                                                     'TO_' + '_'.join(
                                                         n.split()), kstpkper,
                                                     totim)
        recordarray = self._add_empty_record(recordarray, 'TOTAL_OUT',
                                             kstpkper, totim)

        recordarray = self._add_empty_record(recordarray, 'IN-OUT', kstpkper,
                                             totim)
        recordarray = self._add_empty_record(recordarray,
                                             'PERCENT_DISCREPANCY', kstpkper,
                                             totim)
        return recordarray

    @staticmethod
    def _filter_circular_flow(fz, tz, f):
        """

        Parameters
        ----------
        fz
        tz
        f

        Returns
        -------

        """
        e = np.equal(fz, tz)
        fz = fz[np.logical_not(e)]
        tz = tz[np.logical_not(e)]
        f = f[np.logical_not(e)]
        return fz, tz, f

    def _update_budget_fromfaceflow(self, fz, tz, f,
                                    kstpkper=None, totim=None):
        """

        Parameters
        ----------
        fz
        tz
        f
        kstpkper
        totim

        Returns
        -------

        """

        # No circular flow within zones
        fz, tz, f = self._filter_circular_flow(fz, tz, f)

        if len(f) == 0:
            return

        # Inflows
        idx = tz != 0
        fzi = fz[idx]
        tzi = tz[idx]
        rownames = ['FROM_' + self._zonenamedict[z] for z in fzi]
        colnames = [self._zonenamedict[z] for z in tzi]
        fluxes = f[idx]
        self._update_budget_recordarray(rownames, colnames, fluxes, kstpkper,
                                        totim)

        # Outflows
        idx = fz != 0
        fzi = fz[idx]
        tzi = tz[idx]
        rownames = ['TO_' + self._zonenamedict[z] for z in tzi]
        colnames = [self._zonenamedict[z] for z in fzi]
        fluxes = f[idx]
        self._update_budget_recordarray(rownames, colnames, fluxes, kstpkper,
                                        totim)
        return

    def _update_budget_fromssst(self, fz, tz, f, kstpkper=None, totim=None):
        """

        Parameters
        ----------
        fz
        tz
        f
        kstpkper
        totim

        Returns
        -------

        """
        if len(f) == 0:
            return
        self._update_budget_recordarray(fz, tz, f, kstpkper, totim)
        return

    def _update_budget_recordarray(self, rownames, colnames, fluxes,
                                   kstpkper=None, totim=None):
        """
        Update the budget record array with the flux for the specified
        flow direction (in/out), record name, and column.

        Parameters
        ----------
        rownames
        colnames
        fluxes
        kstpkper
        totim

        Returns
        -------
        None

        """
        try:

            if kstpkper is not None:
                for rn, cn, flux in zip(rownames, colnames, fluxes):
                    rowidx = np.where(
                        (self._budget['time_step'] == kstpkper[0]) &
                        (self._budget['stress_period'] == kstpkper[1]) &
                        (self._budget['name'] == rn))
                    self._budget[cn][rowidx] += flux
            elif totim is not None:
                for rn, cn, flux in zip(rownames, colnames, fluxes):
                    rowidx = np.where((self._budget['totim'] == totim) &
                                      (self._budget['name'] == rn))
                    self._budget[cn][rowidx] += flux

        except Exception as e:
            print(e)
            raise
        return

    def _accumulate_flow_frf(self, recname, ich, kstpkper, totim):
        """

        Parameters
        ----------
        recname
        ich
        kstpkper
        totim

        Returns
        -------

        """
        try:
            if self.ncol >= 2:
                data = \
                    self.cbc.get_data(text=recname, kstpkper=kstpkper,
                                      totim=totim)[0]

                # "FLOW RIGHT FACE"  COMPUTE FLOW BETWEEN ZONES ACROSS COLUMNS.
                # COMPUTE FLOW ONLY BETWEEN A ZONE AND A HIGHER ZONE -- FLOW FROM
                # ZONE 4 TO 3 IS THE NEGATIVE OF FLOW FROM 3 TO 4.
                # 1ST, CALCULATE FLOW BETWEEN NODE J,I,K AND J-1,I,K

                k, i, j = np.where(
                    self.izone[:, :, 1:] > self.izone[:, :, :-1])

                # Adjust column values to account for the starting position of "nz"
                j += 1

                # Define the zone to which flow is going
                nz = self.izone[k, i, j]

                # Define the zone from which flow is coming
                jl = j - 1
                nzl = self.izone[k, i, jl]

                # Get the face flow
                q = data[k, i, jl]

                # Get indices where flow face values are positive (flow out of higher zone)
                # Don't include CH to CH flow (can occur if CHTOCH option is used)
                # Create an iterable tuple of (from zone, to zone, flux)
                # Then group tuple by (from_zone, to_zone) and sum the flux values
                idx = np.where(
                    (q > 0) & ((ich[k, i, j] != 1) | (ich[k, i, jl] != 1)))
                fzi, tzi, fi = sum_flux_tuples(nzl[idx],
                                               nz[idx],
                                               q[idx])
                self._update_budget_fromfaceflow(fzi, tzi, np.abs(fi),
                                                 kstpkper, totim)

                # Get indices where flow face values are negative (flow into higher zone)
                # Don't include CH to CH flow (can occur if CHTOCH option is used)
                # Create an iterable tuple of (from zone, to zone, flux)
                # Then group tuple by (from_zone, to_zone) and sum the flux values
                idx = np.where(
                    (q < 0) & ((ich[k, i, j] != 1) | (ich[k, i, jl] != 1)))
                fzi, tzi, fi = sum_flux_tuples(nz[idx],
                                               nzl[idx],
                                               q[idx])
                self._update_budget_fromfaceflow(fzi, tzi, np.abs(fi),
                                                 kstpkper, totim)

                # FLOW BETWEEN NODE J,I,K AND J+1,I,K
                k, i, j = np.where(
                    self.izone[:, :, :-1] > self.izone[:, :, 1:])

                # Define the zone from which flow is coming
                nz = self.izone[k, i, j]

                # Define the zone to which flow is going
                jr = j + 1
                nzr = self.izone[k, i, jr]

                # Get the face flow
                q = data[k, i, j]

                # Get indices where flow face values are positive (flow out of higher zone)
                # Don't include CH to CH flow (can occur if CHTOCH option is used)
                # Create an iterable tuple of (from zone, to zone, flux)
                # Then group tuple by (from_zone, to_zone) and sum the flux values
                idx = np.where(
                    (q > 0) & ((ich[k, i, j] != 1) | (ich[k, i, jr] != 1)))
                fzi, tzi, fi = sum_flux_tuples(nz[idx],
                                               nzr[idx],
                                               q[idx])
                self._update_budget_fromfaceflow(fzi, tzi, np.abs(fi),
                                                 kstpkper, totim)

                # Get indices where flow face values are negative (flow into higher zone)
                # Don't include CH to CH flow (can occur if CHTOCH option is used)
                # Create an iterable tuple of (from zone, to zone, flux)
                # Then group tuple by (from_zone, to_zone) and sum the flux values
                idx = np.where(
                    (q < 0) & ((ich[k, i, j] != 1) | (ich[k, i, jr] != 1)))
                fzi, tzi, fi = sum_flux_tuples(nzr[idx],
                                               nz[idx],
                                               q[idx])
                self._update_budget_fromfaceflow(fzi, tzi, np.abs(fi),
                                                 kstpkper, totim)

                # CALCULATE FLOW TO CONSTANT-HEAD CELLS IN THIS DIRECTION
                k, i, j = np.where(ich == 1)
                k, i, j = k[j > 0], i[j > 0], j[j > 0]
                jl = j - 1
                nzl = self.izone[k, i, jl]
                nz = self.izone[k, i, j]
                q = data[k, i, jl]
                idx = np.where(
                    (q > 0) & ((ich[k, i, j] != 1) | (ich[k, i, jl] != 1)))
                fzi, tzi, f = sum_flux_tuples(nzl[idx],
                                              nz[idx],
                                              q[idx])
                fz = ['TO_CONSTANT_HEAD'] * len(tzi)
                tz = [self._zonenamedict[z] for z in tzi]
                self._update_budget_fromssst(fz, tz, np.abs(f), kstpkper,
                                             totim)

                idx = np.where(
                    (q < 0) & ((ich[k, i, j] != 1) | (ich[k, i, jl] != 1)))
                fzi, tzi, f = sum_flux_tuples(nzl[idx],
                                              nz[idx],
                                              q[idx])
                fz = ['FROM_CONSTANT_HEAD'] * len(fzi)
                tz = [self._zonenamedict[z] for z in tzi[tzi != 0]]
                self._update_budget_fromssst(fz, tz, np.abs(f), kstpkper,
                                             totim)

                k, i, j = np.where(ich == 1)
                k, i, j = k[j < self.ncol - 1], i[j < self.ncol - 1], j[
                    j < self.ncol - 1]
                nz = self.izone[k, i, j]
                jr = j + 1
                nzr = self.izone[k, i, jr]
                q = data[k, i, j]
                idx = np.where(
                    (q > 0) & ((ich[k, i, j] != 1) | (ich[k, i, jr] != 1)))
                fzi, tzi, f = sum_flux_tuples(nzr[idx],
                                              nz[idx],
                                              q[idx])
                fz = ['FROM_CONSTANT_HEAD'] * len(tzi)
                tz = [self._zonenamedict[z] for z in tzi]
                self._update_budget_fromssst(fz, tz, np.abs(f), kstpkper,
                                             totim)

                idx = np.where(
                    (q < 0) & ((ich[k, i, j] != 1) | (ich[k, i, jr] != 1)))
                fzi, tzi, f = sum_flux_tuples(nzr[idx],
                                              nz[idx],
                                              q[idx])
                fz = ['TO_CONSTANT_HEAD'] * len(fzi)
                tz = [self._zonenamedict[z] for z in tzi]
                self._update_budget_fromssst(fz, tz, np.abs(f), kstpkper,
                                             totim)

        except Exception as e:
            print(e)
            raise
        return

    def _accumulate_flow_fff(self, recname, ich, kstpkper, totim):
        """

        Parameters
        ----------
        recname
        ich
        kstpkper
        totim

        Returns
        -------

        """
        try:
            if self.nrow >= 2:
                data = \
                    self.cbc.get_data(text=recname, kstpkper=kstpkper,
                                      totim=totim)[0]

                # "FLOW FRONT FACE"
                # CALCULATE FLOW BETWEEN NODE J,I,K AND J,I-1,K
                k, i, j = np.where(
                    self.izone[:, 1:, :] < self.izone[:, :-1, :])
                i += 1
                ia = i - 1
                nza = self.izone[k, ia, j]
                nz = self.izone[k, i, j]
                q = data[k, ia, j]
                idx = np.where(
                    (q > 0) & ((ich[k, i, j] != 1) | (ich[k, ia, j] != 1)))
                fzi, tzi, fi = sum_flux_tuples(nza[idx],
                                               nz[idx],
                                               q[idx])
                self._update_budget_fromfaceflow(fzi, tzi, np.abs(fi),
                                                 kstpkper, totim)

                idx = np.where(
                    (q < 0) & ((ich[k, i, j] != 1) | (ich[k, ia, j] != 1)))
                fzi, tzi, fi = sum_flux_tuples(nz[idx],
                                               nza[idx],
                                               q[idx])
                self._update_budget_fromfaceflow(fzi, tzi, np.abs(fi),
                                                 kstpkper, totim)

                # CALCULATE FLOW BETWEEN NODE J,I,K AND J,I+1,K.
                k, i, j = np.where(
                    self.izone[:, :-1, :] < self.izone[:, 1:, :])
                nz = self.izone[k, i, j]
                ib = i + 1
                nzb = self.izone[k, ib, j]
                q = data[k, i, j]
                idx = np.where(
                    (q > 0) & ((ich[k, i, j] != 1) | (ich[k, ib, j] != 1)))
                fzi, tzi, fi = sum_flux_tuples(nz[idx],
                                               nzb[idx],
                                               q[idx])
                self._update_budget_fromfaceflow(fzi, tzi, np.abs(fi),
                                                 kstpkper, totim)

                idx = np.where(
                    (q < 0) & ((ich[k, i, j] != 1) | (ich[k, ib, j] != 1)))
                fzi, tzi, fi = sum_flux_tuples(nzb[idx],
                                               nz[idx],
                                               q[idx])
                self._update_budget_fromfaceflow(fzi, tzi, np.abs(fi),
                                                 kstpkper, totim)

                # CALCULATE FLOW TO CONSTANT-HEAD CELLS IN THIS DIRECTION
                k, i, j = np.where(ich == 1)
                k, i, j = k[i > 0], i[i > 0], j[i > 0]
                ia = i - 1
                nza = self.izone[k, ia, j]
                nz = self.izone[k, i, j]
                q = data[k, ia, j]
                idx = np.where(
                    (q > 0) & ((ich[k, i, j] != 1) | (ich[k, ia, j] != 1)))
                fzi, tzi, f = sum_flux_tuples(nza[idx],
                                              nz[idx],
                                              q[idx])
                fz = ['TO_CONSTANT_HEAD'] * len(tzi)
                tz = [self._zonenamedict[z] for z in tzi]
                self._update_budget_fromssst(fz, tz, np.abs(f), kstpkper,
                                             totim)

                idx = np.where(
                    (q < 0) & ((ich[k, i, j] != 1) | (ich[k, ia, j] != 1)))
                fzi, tzi, f = sum_flux_tuples(nza[idx],
                                              nz[idx],
                                              q[idx])
                fz = ['FROM_CONSTANT_HEAD'] * len(fzi)
                tz = [self._zonenamedict[z] for z in tzi]
                self._update_budget_fromssst(fz, tz, np.abs(f), kstpkper,
                                             totim)

                k, i, j = np.where(ich == 1)
                k, i, j = k[i < self.nrow - 1], i[i < self.nrow - 1], j[
                    i < self.nrow - 1]
                nz = self.izone[k, i, j]
                ib = i + 1
                nzb = self.izone[k, ib, j]
                q = data[k, i, j]
                idx = np.where(
                    (q > 0) & ((ich[k, i, j] != 1) | (ich[k, ib, j] != 1)))
                fzi, tzi, f = sum_flux_tuples(nzb[idx],
                                              nz[idx],
                                              q[idx])
                fz = ['FROM_CONSTANT_HEAD'] * len(tzi)
                tz = [self._zonenamedict[z] for z in tzi]
                self._update_budget_fromssst(fz, tz, np.abs(f), kstpkper,
                                             totim)

                idx = np.where(
                    (q < 0) & ((ich[k, i, j] != 1) | (ich[k, ib, j] != 1)))
                fzi, tzi, f = sum_flux_tuples(nzb[idx],
                                              nz[idx],
                                              q[idx])
                fz = ['TO_CONSTANT_HEAD'] * len(fzi)
                tz = [self._zonenamedict[z] for z in tzi]
                self._update_budget_fromssst(fz, tz, np.abs(f), kstpkper,
                                             totim)

        except Exception as e:
            print(e)
            raise
        return

    def _accumulate_flow_flf(self, recname, ich, kstpkper, totim):
        """

        Parameters
        ----------
        recname
        ich
        kstpkper
        totim

        Returns
        -------

        """
        try:
            if self.nlay >= 2:
                data = \
                    self.cbc.get_data(text=recname, kstpkper=kstpkper,
                                      totim=totim)[0]

                # "FLOW LOWER FACE"
                # CALCULATE FLOW BETWEEN NODE J,I,K AND J,I,K-1
                k, i, j = np.where(
                    self.izone[1:, :, :] < self.izone[:-1, :, :])
                k += 1
                ka = k - 1
                nza = self.izone[ka, i, j]
                nz = self.izone[k, i, j]
                q = data[ka, i, j]
                idx = np.where(
                    (q > 0) & ((ich[k, i, j] != 1) | (ich[ka, i, j] != 1)))
                fzi, tzi, fi = sum_flux_tuples(nza[idx],
                                               nz[idx],
                                               q[idx])
                self._update_budget_fromfaceflow(fzi, tzi, np.abs(fi),
                                                 kstpkper, totim)

                idx = np.where(
                    (q < 0) & ((ich[k, i, j] != 1) | (ich[ka, i, j] != 1)))
                fzi, tzi, fi = sum_flux_tuples(nz[idx],
                                               nza[idx],
                                               q[idx])
                self._update_budget_fromfaceflow(fzi, tzi, np.abs(fi),
                                                 kstpkper, totim)

                # CALCULATE FLOW BETWEEN NODE J,I,K AND J,I,K+1
                k, i, j = np.where(
                    self.izone[:-1, :, :] < self.izone[1:, :, :])
                nz = self.izone[k, i, j]
                kb = k + 1
                nzb = self.izone[kb, i, j]
                q = data[k, i, j]
                idx = np.where(
                    (q > 0) & ((ich[k, i, j] != 1) | (ich[kb, i, j] != 1)))
                fzi, tzi, fi = sum_flux_tuples(nz[idx],
                                               nzb[idx],
                                               q[idx])
                self._update_budget_fromfaceflow(fzi, tzi, np.abs(fi),
                                                 kstpkper, totim)

                idx = np.where(
                    (q < 0) & ((ich[k, i, j] != 1) | (ich[kb, i, j] != 1)))
                fzi, tzi, fi = sum_flux_tuples(nzb[idx],
                                               nz[idx],
                                               q[idx])
                self._update_budget_fromfaceflow(fzi, tzi, np.abs(fi),
                                                 kstpkper, totim)

                # CALCULATE FLOW TO CONSTANT-HEAD CELLS IN THIS DIRECTION
                k, i, j = np.where(ich == 1)
                k, i, j = k[k > 0], i[k > 0], j[k > 0]
                ka = k - 1
                nza = self.izone[ka, i, j]
                nz = self.izone[k, i, j]
                q = data[ka, i, j]
                idx = np.where(
                    (q > 0) & ((ich[k, i, j] != 1) | (ich[ka, i, j] != 1)))
                fzi, tzi, f = sum_flux_tuples(nza[idx],
                                              nz[idx],
                                              q[idx])
                fz = ['TO_CONSTANT_HEAD'] * len(tzi)
                tz = [self._zonenamedict[z] for z in tzi]
                self._update_budget_fromssst(fz, tz, np.abs(f), kstpkper,
                                             totim)

                idx = np.where(
                    (q < 0) & ((ich[k, i, j] != 1) | (ich[ka, i, j] != 1)))
                fzi, tzi, f = sum_flux_tuples(nza[idx],
                                              nz[idx],
                                              q[idx])
                fz = ['FROM_CONSTANT_HEAD'] * len(fzi)
                tz = [self._zonenamedict[z] for z in tzi]
                self._update_budget_fromssst(fz, tz, np.abs(f), kstpkper,
                                             totim)

                k, i, j = np.where(ich == 1)
                k, i, j = k[k < self.nlay - 1], i[k < self.nlay - 1], j[
                    k < self.nlay - 1]
                nz = self.izone[k, i, j]
                kb = k + 1
                nzb = self.izone[kb, i, j]
                q = data[k, i, j]
                idx = np.where(
                    (q > 0) & ((ich[k, i, j] != 1) | (ich[kb, i, j] != 1)))
                fzi, tzi, f = sum_flux_tuples(nzb[idx],
                                              nz[idx],
                                              q[idx])
                fz = ['FROM_CONSTANT_HEAD'] * len(tzi)
                tz = [self._zonenamedict[z] for z in tzi]
                self._update_budget_fromssst(fz, tz, np.abs(f), kstpkper,
                                             totim)

                idx = np.where(
                    (q < 0) & ((ich[k, i, j] != 1) | (ich[kb, i, j] != 1)))
                fzi, tzi, f = sum_flux_tuples(nzb[idx],
                                              nz[idx],
                                              q[idx])
                fz = ['TO_CONSTANT_HEAD'] * len(fzi)
                tz = [self._zonenamedict[z] for z in tzi]
                self._update_budget_fromssst(fz, tz, np.abs(f), kstpkper,
                                             totim)

        except Exception as e:
            print(e)
            raise
        return

    def _accumulate_flow_ssst(self, recname, kstpkper, totim):

        # NOT AN INTERNAL FLOW TERM, SO MUST BE A SOURCE TERM OR STORAGE
        # ACCUMULATE THE FLOW BY ZONE

        imeth = self.imeth[recname]

        data = self.cbc.get_data(text=recname, kstpkper=kstpkper,
                                 totim=totim)
        if len(data) == 0:
            # Empty data, can occur during the first time step of a transient
            # model when storage terms are zero and not in the cell-budget
            # file.
            return
        else:
            data = data[0]

        if imeth == 2 or imeth == 5:
            # LIST
            qin = np.ma.zeros((self.nlay * self.nrow * self.ncol),
                              self.float_type)
            qout = np.ma.zeros((self.nlay * self.nrow * self.ncol),
                               self.float_type)
            for [node, q] in zip(data['node'], data['q']):
                idx = node - 1
                if q > 0:
                    qin.data[idx] += q
                elif q < 0:
                    qout.data[idx] += q
            qin = np.ma.reshape(qin, (self.nlay, self.nrow, self.ncol))
            qout = np.ma.reshape(qout, (self.nlay, self.nrow, self.ncol))
        elif imeth == 0 or imeth == 1:
            # FULL 3-D ARRAY
            qin = np.ma.zeros(self.cbc_shape, self.float_type)
            qout = np.ma.zeros(self.cbc_shape, self.float_type)
            qin[data > 0] = data[data > 0]
            qout[data < 0] = data[data < 0]
        elif imeth == 3:
            # 1-LAYER ARRAY WITH LAYER INDICATOR ARRAY
            rlay, rdata = data[0], data[1]
            data = np.ma.zeros(self.cbc_shape, self.float_type)
            for (r, c), l in np.ndenumerate(rlay):
                data[l - 1, r, c] = rdata[r, c]
            qin = np.ma.zeros(self.cbc_shape, self.float_type)
            qout = np.ma.zeros(self.cbc_shape, self.float_type)
            qin[data > 0] = data[data > 0]
            qout[data < 0] = data[data < 0]
        elif imeth == 4:
            # 1-LAYER ARRAY THAT DEFINES LAYER 1
            qin = np.ma.zeros(self.cbc_shape, self.float_type)
            qout = np.ma.zeros(self.cbc_shape, self.float_type)
            r, c = np.where(data > 0)
            qin[0, r, c] = data[r, c]
            r, c = np.where(data < 0)
            qout[0, r, c] = data[r, c]
        else:
            # Should not happen
            raise Exception(
                'Unrecognized "imeth" for {} record: {}'.format(recname,
                                                                imeth))

        # Inflows
        fz = []
        tz = []
        f = []
        for z in self.allzones:
            if z != 0:
                flux = qin[(self.izone == z)].sum()
                if type(flux) == np.ma.core.MaskedConstant:
                    flux = 0.
                fz.append('FROM_' + '_'.join(recname.split()))
                tz.append(self._zonenamedict[z])
                f.append(flux)
        fz = np.array(fz)
        tz = np.array(tz)
        f = np.array(f)
        self._update_budget_fromssst(fz, tz, np.abs(f), kstpkper, totim)

        # Outflows
        fz = []
        tz = []
        f = []
        for z in self.allzones:
            if z != 0:
                flux = qout[(self.izone == z)].sum()
                if type(flux) == np.ma.core.MaskedConstant:
                    flux = 0.
                fz.append('TO_' + '_'.join(recname.split()))
                tz.append(self._zonenamedict[z])
                f.append(flux)
        fz = np.array(fz)
        tz = np.array(tz)
        f = np.array(f)
        self._update_budget_fromssst(fz, tz, np.abs(f), kstpkper, totim)

        return

    def _compute_mass_balance(self, kstpkper, totim):
        # Returns a record array with total inflow, total outflow,
        # and percent error summed by column.
        skipcols = ['time_step', 'stress_period', 'totim', 'name']

        # Compute inflows
        recnames = self.get_record_names()
        innames = [n for n in recnames if n.startswith('FROM_')]
        outnames = [n for n in recnames if n.startswith('TO_')]
        if kstpkper is not None:
            rowidx = np.where((self._budget['time_step'] == kstpkper[0]) &
                              (self._budget['stress_period'] == kstpkper[1]) &
                              np.in1d(self._budget['name'], innames))
        elif totim is not None:
            rowidx = np.where((self._budget['totim'] == totim) &
                              np.in1d(self._budget['name'], innames))
        a = _numpyvoid2numeric(
            self._budget[list(self._zonenamedict.values())][rowidx])
        intot = np.array(a.sum(axis=0))
        tz = np.array(
            list([n for n in self._budget.dtype.names if n not in skipcols]))
        fz = np.array(['TOTAL_IN'] * len(tz))
        self._update_budget_fromssst(fz, tz, intot, kstpkper, totim)

        # Compute outflows
        if kstpkper is not None:
            rowidx = np.where((self._budget['time_step'] == kstpkper[0]) &
                              (self._budget['stress_period'] == kstpkper[1]) &
                              np.in1d(self._budget['name'], outnames))
        elif totim is not None:
            rowidx = np.where((self._budget['totim'] == totim) &
                              np.in1d(self._budget['name'], outnames))
        a = _numpyvoid2numeric(
            self._budget[list(self._zonenamedict.values())][rowidx])
        outot = np.array(a.sum(axis=0))
        tz = np.array(
            list([n for n in self._budget.dtype.names if n not in skipcols]))
        fz = np.array(['TOTAL_OUT'] * len(tz))
        self._update_budget_fromssst(fz, tz, outot, kstpkper, totim)

        # Compute IN-OUT
        tz = np.array(
            list([n for n in self._budget.dtype.names if n not in skipcols]))
        f = intot - outot
        fz = np.array(['IN-OUT'] * len(tz))
        self._update_budget_fromssst(fz, tz, np.abs(f), kstpkper, totim)

        # Compute percent discrepancy
        tz = np.array(
            list([n for n in self._budget.dtype.names if n not in skipcols]))
        fz = np.array(['PERCENT_DISCREPANCY'] * len(tz))
        in_minus_out = intot - outot
        in_plus_out = intot + outot
        f = 100 * in_minus_out / (in_plus_out / 2.)
        self._update_budget_fromssst(fz, tz, np.abs(f), kstpkper, totim)

        return

    def _clean_budget_names(self, names):
        newnames = []
        mbnames = ['TOTAL_IN', 'TOTAL_OUT',
                   'IN-OUT', 'PERCENT_DISCREPANCY']
        for name in names:
            if name in mbnames:
                newnames.append(name)
            elif not name.startswith('FROM_') and not name.startswith('TO_'):
                newname_in = 'FROM_' + name.upper()
                newname_out = 'TO_' + name.upper()
                if newname_in in self._budget['name']:
                    newnames.append(newname_in)
                if newname_out in self._budget['name']:
                    newnames.append(newname_out)
            else:
                if name in self._budget['name']:
                    newnames.append(name)
        return newnames

    def _compute_net_budget(self):
        recnames = self.get_record_names()
        innames = [n for n in recnames if n.startswith('FROM_')]
        outnames = [n for n in recnames if n.startswith('TO_')]
        select_fields = ['totim', 'time_step', 'stress_period',
                         'name'] + list(self._zonenamedict.values())
        select_records_in = np.in1d(self._budget['name'], innames)
        select_records_out = np.in1d(self._budget['name'], outnames)
        in_budget = self._budget[select_fields][select_records_in]
        out_budget = self._budget[select_fields][select_records_out]
        net_budget = in_budget.copy()
        for f in [n for n in self._zonenamedict.values() if
                  n in select_fields]:
            net_budget[f] = np.array([r for r in in_budget[f]]) - np.array(
                [r for r in out_budget[f]])
        newnames = ['_'.join(n.split('_')[1:]) for n in net_budget['name']]
        net_budget['name'] = newnames
        return net_budget

    def __mul__(self, other):
        newbud = self._budget.copy()
        for f in self._zonenamedict.values():
            newbud[f] = np.array([r for r in newbud[f]]) * other
        idx = np.in1d(self._budget['name'], 'PERCENT_DISCREPANCY')
        newbud[:][idx] = self._budget[:][idx]
        newobj = self.copy()
        newobj._budget = newbud
        return newobj

    def __truediv__(self, other):
        newbud = self._budget.copy()
        for f in self._zonenamedict.values():
            newbud[f] = np.array([r for r in newbud[f]]) / float(other)
        idx = np.in1d(self._budget['name'], 'PERCENT_DISCREPANCY')
        newbud[:][idx] = self._budget[:][idx]
        newobj = self.copy()
        newobj._budget = newbud
        return newobj

    def __div__(self, other):
        newbud = self._budget.copy()
        for f in self._zonenamedict.values():
            newbud[f] = np.array([r for r in newbud[f]]) / float(other)
        idx = np.in1d(self._budget['name'], 'PERCENT_DISCREPANCY')
        newbud[:][idx] = self._budget[:][idx]
        newobj = self.copy()
        newobj._budget = newbud
        return newobj

    def __add__(self, other):
        newbud = self._budget.copy()
        for f in self._zonenamedict.values():
            newbud[f] = np.array([r for r in newbud[f]]) + other
        idx = np.in1d(self._budget['name'], 'PERCENT_DISCREPANCY')
        newbud[:][idx] = self._budget[:][idx]
        newobj = self.copy()
        newobj._budget = newbud
        return newobj

    def __sub__(self, other):
        newbud = self._budget.copy()
        for f in self._zonenamedict.values():
            newbud[f] = np.array([r for r in newbud[f]]) - other
        idx = np.in1d(self._budget['name'], 'PERCENT_DISCREPANCY')
        newbud[:][idx] = self._budget[:][idx]
        newobj = self.copy()
        newobj._budget = newbud
        return newobj


def _numpyvoid2numeric(a):
    # The budget record array has multiple dtypes and a slice returns
    # the flexible-type numpy.void which must be converted to a numeric
    # type prior to performing reducing functions such as sum() or
    # mean()
    return np.array([list(r) for r in a])


def write_zbarray(fname, X, fmtin=None, iprn=None):
    """
    Saves a numpy array in a format readable by the zonebudget program
    executable.

    File format:
    line 1: nlay, nrow, ncol
    line 2: INTERNAL (format)
    line 3: begin data
    .
    .
    .

    example from NACP:
    19 250 500
    INTERNAL      (10I8)
    199     199     199     199     199     199     199     199     199     199
    199     199     199     199     199     199     199     199     199     199
    ...
    INTERNAL      (10I8)
    199     199     199     199     199     199     199     199     199     199
    199     199     199     199     199     199     199     199     199     199
    ...

    Parameters
    ----------
    X : array
        The array of zones to be written.
    fname :  str
        The path and name of the file to be written.
    fmtin : int
        The number of values to write to each line.
    iprn : int
        Padding space to add between each value.

    Returns
    -------

    """
    if len(X.shape) == 2:
        b = np.zeros((1, X.shape[0], X.shape[1]), dtype=np.int32)
        b[0, :, :] = X[:, :]
        X = b.copy()
    elif len(X.shape) < 2 or len(X.shape) > 3:
        raise Exception(
            'Shape of the input array is not recognized: {}'.format(X.shape))
    if np.ma.is_masked(X):
        X = np.ma.filled(X, 0)

    nlay, nrow, ncol = X.shape

    if fmtin is not None:
        assert fmtin < ncol, 'The specified width is greater than the ' \
                             'number of columns in the array.'
    else:
        fmtin = ncol

    iprnmin = len(str(X.max()))
    if iprn is None or iprn <= iprnmin:
        iprn = iprnmin + 1

    formatter_str = '{{:>{iprn}}}'.format(iprn=iprn)
    formatter = formatter_str.format

    with open(fname, 'w') as f:
        header = '{nlay} {nrow} {ncol}\n'.format(nlay=nlay,
                                                 nrow=nrow,
                                                 ncol=ncol)
        f.write(header)
        for lay in range(nlay):
            record_2 = 'INTERNAL\t({fmtin}I{iprn})\n'.format(fmtin=fmtin,
                                                             iprn=iprn)
            f.write(record_2)
            if fmtin < ncol:
                for row in range(nrow):
                    rowvals = X[lay, row, :].ravel()
                    start = 0
                    end = start + fmtin
                    vals = rowvals[start:end]
                    while len(vals) > 0:
                        s = ''.join(
                            [formatter(int(val)) for val in vals]) + '\n'
                        f.write(s)
                        start = end
                        end = start + fmtin
                        vals = rowvals[start:end]

            elif fmtin == ncol:
                for row in range(nrow):
                    vals = X[lay, row, :].ravel()
                    f.write(
                        ''.join([formatter(int(val)) for val in vals]) + '\n')
    return


def read_zbarray(fname):
    """
    Reads an ascii array in a format readable by the zonebudget program
    executable.

    Parameters
    ----------
    fname :  str
        The path and name of the file to be written.

    Returns
    -------
    zones : numpy ndarray
        An integer array of the zones.
    """
    with open(fname, 'r') as f:
        lines = f.readlines()

    # Initialize layer
    lay = 0

    # Initialize data counter
    totlen = 0
    i = 0

    # First line contains array dimensions
    dimstring = lines.pop(0).strip().split()
    nlay, nrow, ncol = [int(v) for v in dimstring]
    zones = np.zeros((nlay, nrow, ncol), dtype=np.int32)

    # The number of values to read before placing
    # them into the zone array
    datalen = nrow * ncol

    # List of valid values for LOCAT
    locats = ['CONSTANT', 'INTERNAL', 'EXTERNAL']

    # ITERATE OVER THE ROWS
    for line in lines:
        rowitems = line.strip().split()

        # Skip blank lines
        if len(rowitems) == 0:
            continue

        # HEADER
        if rowitems[0].upper() in locats:
            vals = []
            locat = rowitems[0].upper()

            if locat == 'CONSTANT':
                iconst = int(rowitems[1])
            else:
                fmt = rowitems[1].strip('()')
                fmtin, iprn = [int(v) for v in fmt.split('I')]

        # ZONE DATA
        else:
            if locat == 'CONSTANT':
                vals = np.ones((nrow, ncol), dtype=np.int32) * iconst
                lay += 1
            elif locat == 'INTERNAL':
                # READ ZONES
                rowvals = [int(v) for v in rowitems]
                s = 'Too many values encountered on this line.'
                assert len(rowvals) <= fmtin, s
                vals.extend(rowvals)

            elif locat == 'EXTERNAL':
                # READ EXTERNAL FILE
                fname = rowitems[0]
                if not os.path.isfile(fname):
                    errmsg = 'Could not find external file "{}"'.format(fname)
                    raise Exception(errmsg)
                with open(fname, 'r') as ext_f:
                    ext_flines = ext_f.readlines()
                for ext_frow in ext_flines:
                    ext_frowitems = ext_frow.strip().split()
                    rowvals = [int(v) for v in ext_frowitems]
                    vals.extend(rowvals)
                if len(vals) != datalen:
                    errmsg = 'The number of values read from external ' \
                             'file "{}" does not match the expected ' \
                             'number.'.format(len(vals))
                    raise Exception(errmsg)
            else:
                # Should not get here
                raise Exception('Locat not recognized: {}'.format(locat))

                # IGNORE COMPOSITE ZONES

            if len(vals) == datalen:
                # place values for the previous layer into the zone array
                vals = np.array(vals, dtype=np.int32).reshape((nrow, ncol))
                zones[lay, :, :] = vals[:, :]
                lay += 1
            totlen += len(rowitems)
        i += 1
    s = 'The number of values read ({:,.0f})' \
        ' does not match the number expected' \
        ' ({:,.0f})'.format(totlen,
                            nlay * nrow * ncol)
    assert totlen == nlay * nrow * ncol, s
    return zones


def sum_flux_tuples(fromzones, tozones, fluxes):
    tup = zip(fromzones, tozones, fluxes)
    sorted_tups = sort_tuple(tup)

    # Group the sorted tuples by (from zone, to zone)
    # itertools.groupby() returns the index (from zone, to zone) and
    # a list of the tuples with that index
    from_zones = []
    to_zones = []
    fluxes = []
    for (fz, tz), ftup in groupby(sorted_tups, lambda tup: tup[:2]):
        f = np.sum([tup[-1] for tup in list(ftup)])
        from_zones.append(fz)
        to_zones.append(tz)
        fluxes.append(f)
    return np.array(from_zones), np.array(to_zones), np.array(fluxes)


def sort_tuple(tup, n=2):
    """
    Sort a tuple by the first n values
    :param tup:
    :param n:
    :return:
    """
    return tuple(sorted(tup, key=lambda t: t[:n]))


def get_totim_modflow6(tdis):
    """
    Create a totim array from the tdis file in modflow 6

    Parameters
    ----------
    tdis : ModflowTdis object

    Returns
    -------
        totim : np.ndarray

    """
    recarray = tdis.perioddata.array
    delt = []
    for record in recarray:
        perlen = record.perlen
        nstp = record.nstp
        tsmult = record.tsmult
        for stp in range(nstp):
            if stp == 0:
                if tsmult != 1.0:
                    dt = perlen * (tsmult - 1) / ((tsmult ** nstp) - 1)
                else:
                    dt = perlen / nstp
            else:
                dt = delt[-1] * tsmult

            delt.append(dt)

    totim = np.add.accumulate(delt)

    return totim


class ZBNetOutput(object):
    """
    Class that holds zonebudget netcdf output and allows export utilities
    to recognize the output data type.

    Parameters
    ----------
    zones : np.ndarray
        array of zone numbers
    time : np.ndarray
        array of totim
    arrays : dict
        dictionary of budget term arrays.
        axis 0 is totim,
        axis 1 is zones
    flux : bool
        boolean flag to indicate if budget data is a flux "L^3/T"(True,
        default) or if the data have been processed to
        volumetric values "L^3" (False)
    """

    def __init__(self, zones, time, arrays, zone_array, flux=True):
        self.zones = zones
        self.time = time
        self.arrays = arrays
        self.zone_array = zone_array
        self.flux = flux


class ZoneBudgetOutput(object):
    """
    Class method to process zonebudget output into volumetric budgets

    Parameters
    ----------
    f : str
        zonebudget output file path
    dis : flopy.modflow.ModflowDis object
    zones : np.ndarray
        numpy array of zones

    """

    def __init__(self, f, dis, zones):
        import pandas as pd
        from ..modflow import ModflowDis

        self._filename = f
        self._otype = None
        self._zones = zones
        self.__pd = pd

        if isinstance(dis, ModflowDis):
            self._totim = dis.get_totim()
            self._nstp = dis.nstp.array
            self._steady = dis.steady.array

        else:
            self._totim = get_totim_modflow6(dis)
            self._nstp = np.array(dis.perioddata.array.nstp)
            # self._steady is a placeholder, data not used for ZB6 read
            self._steady = [False for _ in dis.perioddata.array]

        self._tslen = None
        self._date_time = None
        self._data = None

        if self._otype is None:
            self._get_otype()

        self._calculate_tslen()
        self._read_file()

    def __repr__(self):
        """
        String representation of the ZoneBudgetOutput class

        """
        zones = ", ".join([str(i) for i in self.zones])
        l = ["ZoneBudgetOutput Class",
             "----------------------\n",
             "Number of zones: {}".format(len(self.zones)),
             "Unique zones: {}".format(zones),
             "Number of buget records: {}".format(len(self.dataframe))]

        return "\n".join(l)

    @property
    def zone_array(self):
        """
        Property method to get the zone array

        """
        return np.asarray(self._zones, dtype=int)

    @property
    def zones(self):
        """
        Get a unique list of zones

        """
        return np.unique(self.zone_array)

    @property
    def dataframe(self):
        """
        Returns a net flux dataframe of the zonebudget output

        """
        return self.__pd.DataFrame.from_dict(self._data)

    def _calculate_tslen(self):
        """
        Method to calculate each timestep length from totim
        and reset totim to a dictionary of {(kstp, kper): totim}

        """
        n = 0
        totim = {}
        for ix, stp in enumerate(self._nstp):
            for i in range(stp):
                if self._tslen is None:
                    tslen = self._totim[n]
                    self._tslen = {(i, ix): tslen}
                else:
                    tslen = self._totim[n] - self._totim[n - 1]
                    self._tslen[(i, ix)] = tslen

                totim[(i, ix)] = self._totim[n]
                n += 1

        self._totim = totim

    def _read_file(self):
        """
        Delegator method for reading zonebudget outputs

        """
        if self._otype == 1:
            self._read_file1()
        elif self._otype == 2:
            self._read_file2()
        elif self._otype == 3:
            self._read_file3()
        else:
            raise AssertionError("Invalid otype supplied: {}"
                                 .format(self._otype))

    def _read_file1(self):
        """
        Read original style zonebudget output file

        """

        with open(self._filename) as foo:

            data_in = {}
            data_out = {}
            read_in = False
            read_out = False
            flow_budget = False
            empty = 0
            while True:
                line = foo.readline().strip().lower()

                if "flow budget for zone" in line:
                    flow_budget = True
                    read_in = False
                    read_out = False
                    empty = 0
                    t = line.split()
                    zone = int(t[4])
                    if len(t[7]) > 4:
                        t.insert(8, t[7][4:])
                    kstp = int(t[8]) - 1
                    if len(t[11]) > 6:
                        t.append(t[11][6:])
                    kper = int(t[12]) - 1
                    if "zone" not in data_in:
                        data_in["zone"] = [zone]
                        data_in['kstp'] = [kstp]
                        data_in['kper'] = [kper]
                    else:
                        data_in['zone'].append(zone)
                        data_in['kstp'].append(kstp)
                        data_in['kper'].append(kper)

                    if self._steady[kper]:
                        try:
                            data_in['storage'].append(0.)
                            data_out['storage'].append(0.)
                        except KeyError:
                            data_in['storage'] = [0.]
                            data_out['storage'] = [0.]

                elif line in ("", " "):
                    empty += 1

                elif read_in:
                    if "=" in line:
                        t = line.split("=")
                        label = t[0].strip()
                        if "zone" in line:
                            # currently we do not support zone to zone
                            # flow for option 1
                            pass
                        else:
                            if "total" in line:
                                label = "total"

                            if label in data_in:
                                data_in[label].append(float(t[1]))
                            else:
                                data_in[label] = [float(t[1])]

                    elif "out:" in line:
                        read_out = True
                        read_in = False

                    else:
                        pass

                elif read_out:
                    if "=" in line:
                        t = line.split("=")
                        label = t[0].strip()
                        if "zone" in line:
                            # currently we do not support zone to zone
                            # flow for option 1
                            pass

                        elif "in - out" in line:
                            pass

                        elif "percent discrepancy" in line:
                            pass

                        else:
                            if "total" in line:
                                label = "total"

                            if label in data_out:
                                data_out[label].append(float(t[1]))
                            else:
                                data_out[label] = [float(t[1])]
                    else:
                        pass

                elif flow_budget:
                    if "in:" in line:
                        read_in = True
                        flow_budget = False

                else:
                    pass

                if empty >= 30:
                    break

        data = self._net_flux(data_in, data_out)

        self._data = data

    def _read_file2(self):
        """
        Method to read csv output type 1

        """
        with open(self._filename) as foo:
            data_in = {}
            data_out = {}
            zone_header = False
            read_in = False
            read_out = False
            empty = 0
            while True:
                line = foo.readline().strip().lower()

                if "time step" in line:
                    t = line.split(",")
                    kstp = int(t[1]) - 1
                    kper = int(t[3]) - 1
                    if "kstp" not in data_in:
                        data_in['kstp'] = []
                        data_in['kper'] = []
                        data_in["zone"] = []

                    zone_header = True
                    empty = 0

                elif zone_header:
                    t = line.split(",")
                    zones = [int(i.split()[-1]) for i in t[1:]
                             if i not in ('',)]

                    for zone in zones:
                        data_in['kstp'].append(kstp)
                        data_in['kper'].append(kper)
                        data_in['zone'].append(zone)
                        if self._steady[kper]:
                            try:
                                data_in['storage'].append(0.)
                                data_out['storage'].append(0.)
                            except KeyError:
                                data_in['storage'] = [0.]
                                data_out['storage'] = [0.]

                    zone_header = False
                    read_in = True

                elif read_in:
                    t = line.split(",")
                    if "in" in t[1]:
                        pass

                    elif "out" in t[1]:
                        read_in = False
                        read_out = True

                    else:
                        if 'zone' in t[0]:
                            label = " ".join(t[0].split()[1:])

                        elif 'total' in t[0]:
                            label = "total"

                        else:
                            label = t[0]

                        if label not in data_in:
                            data_in[label] = []

                        for val in t[1:]:
                            if val in ('',):
                                continue

                            data_in[label].append(float(val))

                elif read_out:
                    t = line.split(",")

                    if "percent error" in line:
                        read_out = False

                    elif "in-out" in line:
                        pass

                    else:
                        if 'zone' in t[0]:
                            label = " ".join(t[0].split()[1:])

                        elif 'total' in t[0]:
                            label = "total"

                        else:
                            label = t[0]

                        if label not in data_out:
                            data_out[label] = []

                        for val in t[1:]:
                            if val in ('',):
                                continue

                            data_out[label].append(float(val))

                elif line in ("", " "):
                    empty += 1

                else:
                    pass

                if empty >= 25:
                    break

        data = self._net_flux(data_in, data_out)

        self._data = data

    def _read_file3(self):
        """
        Method to read CSV2 output from zonebudget and CSV output
        from Zonebudget6

        """
        with open(self._filename) as foo:
            data_in = {}
            data_out = {}
            read_in = True
            read_out = False
            # read the header
            header = foo.readline().lower().strip().split(',')
            header = [i.strip() for i in header]

            array = np.genfromtxt(foo, delimiter=",").T

        for ix, label in enumerate(header):
            if label in ('totim', 'in-out', 'percent error'):
                continue

            elif label == 'percent error':
                continue

            elif label == "step":
                label = "kstp"

            elif label == "period":
                label = "kper"

            elif "other zones" in label:
                label = "other zones"

            elif "from zone" in label or "to zone" in label:
                if "from" in label:
                    read_in = True
                    read_out = False
                else:
                    read_out = True
                    read_in = False
                label = " ".join(label.split()[1:])

            elif 'total' in label:
                label = 'total'

            elif label.split("-")[-1] == "in":
                label = "-".join(label.split("-")[:-1])
                read_in = True
                read_out = False

            elif label.split("-")[-1] == "out":
                label = "-".join(label.split("-")[:-1])
                read_in = False
                read_out = True

            else:
                pass

            if read_in:

                if label in ('kstp', 'kper'):
                    data_in[label] = np.asarray(array[ix], dtype=int) - 1

                elif label == "zone":
                    data_in[label] = np.asarray(array[ix], dtype=int)

                else:
                    data_in[label] = array[ix]

                if label == 'total':
                    read_in = False
                    read_out = True

            elif read_out:
                data_out[label] = array[ix]

            else:
                pass

        data = self._net_flux(data_in, data_out)

        self._data = data

    def _net_flux(self, data_in, data_out):
        """
        Method to create a single dictionary of net flux data

        data_in : dict
            inputs to the zone
        data_out : dict
            outputs from the zone

        Returns
        -------
        dict : dictionary of netflux data to feed into a pandas dataframe
        """
        data = {}
        # calculate net storage flux (subroutine this?)
        for key, value in data_in.items():
            if key in ("zone", "kstp", "kper"):
                data[key] = np.asarray(value, dtype=int)
            else:
                arrayin = np.asarray(value)
                arrayout = np.asarray(data_out[key])

                data[key] = arrayin - arrayout

        kstp = data['kstp']
        kper = data['kper']
        tslen = np.array([self._tslen[(stp, kper[ix])]
                          for ix, stp in enumerate(kstp)])
        totim = np.array([self._totim[(stp, kper[ix])]
                          for ix, stp in enumerate(kstp)])

        data["tslen"] = tslen
        data['totim'] = totim

        return data

    def _get_otype(self):
        """
        Method to automatically distinguish output type based on the
        zonebudget header

        """
        with open(self._filename) as foo:
            line = foo.readline()
            if "zonebudget version" in line.lower():
                self._otype = 1
            elif "time step" in line.lower():
                self._otype = 2
            elif "totim" in line.lower():
                self._otype = 3
            else:
                raise AssertionError("Cant distinguish output type")

    def export(self, f, ml, **kwargs):
        """
        Method to export a netcdf file, or add zonebudget output to
        an open netcdf file instance

        Parameters
        ----------
        f : str or flopy.export.netcdf.NetCdf object
        ml : flopy.modflow.Modflow or flopy.mf6.ModflowGwf object
        **kwargs :
            logger : flopy.export.netcdf.Logger instance
            masked_vals : list
                list of values to mask

        Returns
        -------
            flopy.export.netcdf.NetCdf object

        """
        from flopy.export.utils import output_helper

        if isinstance(f, str):
            if not f.endswith(".nc"):
                raise AssertionError("File extension must end with .nc to "
                                     "export a netcdf file")

        zbncfobj = self.dataframe_to_netcdf_fmt(self.dataframe)
        oudic = {"zbud": zbncfobj}
        return output_helper(f, ml, oudic, **kwargs)

    def volumetric_flux(self, extrapolate_kper=False):
        """
        Method to generate a volumetric budget table based on flux information

        Parameters
        ----------
        extrapolate_kper : bool
            flag to determine if we fill in data gaps with other
            timestep information from the same stress period.
            if True, we assume that flux is constant throughout a stress period
            and the pandas dataframe returned contains a
            volumetric budget per stress period

            if False, calculates volumes from available flux data

        Returns
        -------
            pd.DataFrame

        """
        nper = len(self._nstp)
        volumetric_data = {}

        for key in self._data:
            volumetric_data[key] = []

        if extrapolate_kper:
            volumetric_data.pop("tslen")
            volumetric_data.pop("kstp")
            volumetric_data['perlen'] = []

            perlen = []
            for per in range(nper):
                tslen = 0
                for stp in range(self._nstp[per]):
                    tslen += self._tslen[(stp, per)]

                perlen.append(tslen)

            totim = np.add.accumulate(perlen)

            for per in range(nper):
                idx = np.where(self._data["kper"] == per)[0]

                if len(idx) == 0:
                    continue

                temp = self._data["zone"][idx]

                for zone in self.zones:
                    if zone == 0:
                        continue

                    zix = np.where(temp == zone)[0]

                    if len(zix) == 0:
                        raise Exception

                    for key, value in self._data.items():
                        if key == "totim":
                            volumetric_data[key].append(totim[per])

                        elif key == "tslen":
                            volumetric_data["perlen"].append(perlen[per])

                        elif key == "kstp":
                            continue

                        elif key == "kper":
                            volumetric_data[key].append(per)

                        elif key == "zone":
                            volumetric_data[key].append(zone)

                        else:
                            tv = value[idx]
                            zv = tv[zix]
                            for i in zv:
                                vol = i * perlen[per]
                                volumetric_data[key].append(vol)
                                break

        else:

            for key, value in self._data.items():
                if key in ("zone", "kstp", "kper", "tslen"):
                    volumetric_data[key] = value
                else:
                    volumetric_data[key] = value * self._data['tslen']

        return self.__pd.DataFrame.from_dict(volumetric_data)

    def dataframe_to_netcdf_fmt(self, df, flux=True):
        """
        Method to transform a volumetric zonebudget dataframe into
        array format for netcdf.

        time is on axis 0
        zone is on axis 1

        Parameters
        ----------
        df : pd.DataFrame
        flux : bool
            boolean flag to indicate if budget data is a flux "L^3/T" (True,
            default) or if the data have been processed to
            volumetric values "L^3" (False)
        zone_array : np.ndarray
            zonebudget zones array

        Returns
        -------
            ZBNetOutput object

        """
        zones = np.sort(np.unique(df.zone.values))
        totim = np.sort(np.unique(df.totim.values))

        data = {}
        for col in df.columns:
            if col in ('totim', 'zone', 'kper', 'perlen'):
                pass
            else:
                data[col] = np.zeros((totim.size, zones.size), dtype=float)

        for i, time in enumerate(totim):
            tdf = df.loc[df.totim.isin([time, ])]
            tdf = tdf.sort_values(by=['zone'])

            for col in df.columns:
                if col in ('totim', 'zone', 'kper', 'perlen'):
                    pass
                else:
                    data[col][i, :] = tdf[col].values

        return ZBNetOutput(zones, totim, data, self.zone_array, flux=flux)
