import sys
import collections
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 flopy.utils import Util2d
import subprocess as sub

class ModflowZbd(Package):
    """
    MODFLOW Zone Budget Package Class.

    Parameters
    ----------
    model : model object
        The model object (of type :class:`flopy.modflow.mf.Modflow`) to which
        this package will be added.
    zone_dict : dict
        Dictionary with zone data for the model.
        Keywords are the layer number (0-based)
        Argument are 2D numpy array shape = (nrow, ncol)
    extension : string
        Filename extension (default is 'zon')
    unitnumber : int
        File unit number (default is 2001).


    Attributes
    ----------

    Methods
    -------

    See Also
    --------

    Notes
    -----
    Parameters are supported in Flopy only when reading in existing models.
    Parameter values are converted to native values in Flopy and the
    connection to "parameters" is thus nonexistent.


    """

    def __init__(self, model, nzn = 1, zone_dict = None, exe_name = None, extension='zon', unitnumber=2001):
        """
        Package constructor.

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


        self.nzn = nzn
        self.parent.add_package(self)
        self.zone_dict = zone_dict
        self.exe_name = exe_name


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


        """
        nrow, ncol, nlay, nper = self.parent.nrow_ncol_nlay_nper
        # Open file for writing
        f_zbd = open(self.fn_path, 'w')

        # Item 1: NLAY, NROW, NCOL
        f_zbd.write('{0:3d}{1:10d}{2:10d} \n'.format(nlay,nrow,ncol))
        # Iem 2: IZONE(NCOL,NROW), for each layer of the model
        #        this array is stored in the dictionary zone_dict, where keywords are layer number
        for kl in range(nlay):
            f_zbd.write('INTERNAL (20I2) \n')
            ar = self.zone_dict[kl]
            for nr in range(nrow):
                for nc in range(ncol):
                    f_zbd.write('{0:2d}'.format(int(ar[nr,nc])))
                f_zbd.write('\n')
        # TO DO: maybe the above part can be implemented more effciently if using util_2d by Flopy


    def run_zonebudget(self):

        modelpath = self.parent.model_ws
        import os
        modelname = [os.path.basename(self.parent.name).replace('', '')][0]

        outputfile = os.path.join(modelpath, modelname + '.zbud')
        budgetfile = os.path.join(modelpath, modelname + '.cbc')
        zonefile = os.path.join(modelpath, modelname + '.zon')
        # outputfile =   modelpath + '\\' + modelname + '.zbud'
        # budgetfile =   modelpath + '\\' + modelname + '.cbc'
        # zonefile   =   modelpath + '\\' + modelname + '.zon'


        p = sub.Popen([self.exe_name], stdout=sub.PIPE, stdin=sub.PIPE, stderr = sub.PIPE)
        p.stdin.write(outputfile + '\n')
        p.stdin.write(budgetfile + '\n')
        p.stdin.write('Zone Budget for Model ' + self.parent.name + ' \n')
        p.stdin.write(zonefile + '\n')
        p.stdin.write('A\n')
