<reponame>kwilcox/flopy
#from numpy import ones, empty
import numpy as np
from flopy.mbase import Package
from flopy.utils import util_2d

class ModflowEvt(Package):
    'Evapotranspiration class'
    def __init__(self, model, nevtop=3, ievtcb=0, surf=0., evtr=1e-3, exdp=1., ievt=1, 
                 extension ='evt', unitnumber=22, external=True):
        '''
        external flag is used to control writing external arrays of constant value
        since this package has the potential to create a lot of external arrays
        '''
        Package.__init__(self, model, extension, 'EVT', unitnumber) # Call ancestor's init to set self.parent, extension, name and unit number
        nrow, ncol, nlay, nper = self.parent.nrow_ncol_nlay_nper
        self.heading = '# EVT for MODFLOW, generated by Flopy.'
        self.url = 'evt.htm'
        self.nevtop = nevtop
        self.ievtcb = ievtcb
        self.surf = []
        self.evtr = []
        self.exdp = []
        self.ievt = []        
        self.external = external
        if self.external is False:
            load = True
        else:
            load = model.load            
        if (not isinstance(surf, list)):
            surf = [surf]
        for i,a in enumerate(surf):
            s = util_2d(model,(nrow,ncol),np.float32,a,name='surf_'+str(i+1))
            self.surf = self.surf + [s]
        if (not isinstance(evtr, list)):
            evtr = [evtr]
        for i,a in enumerate(evtr):
            e = util_2d(model,(nrow,ncol),np.float32,a,name='etvr_'+str(i+1))
            self.evtr = self.evtr + [e]
        if (not isinstance(exdp, list)):
            exdp = [exdp]
        for i,a in enumerate(exdp):
            e = util_2d(model,(nrow,ncol),np.float32,a,name='exdp_'+str(i+1))
            self.exdp = self.exdp + [e]
        if (not isinstance(ievt, list)):
            ievt = [ievt]
        for i,a in enumerate(ievt):
            iv = util_2d(model,(nrow,ncol),np.int,a,name='ievt_'+str(i+1))
            self.ievt = self.ievt + [iv]
        self.np = 0
        self.parent.add_package(self)
    def __repr__( self ):
        return 'Evapotranspiration class'
    def ncells( self):
        # Returns the  maximum number of cells that have 
        # evapotranspiration (developped for MT3DMS SSM package)
        nrow, ncol, nlay, nper = self.parent.nrow_ncol_nlay_nper
        return (nrow * ncol)
    def write_file(self):
        nrow, ncol, nlay, nper = self.parent.nrow_ncol_nlay_nper
        f_evt = open(self.fn_path, 'w')        
        f_evt.write('{0:s}\n'.format(self.heading))        
        f_evt.write('{0:10d}{1:10d}\n'.format(self.nevtop,self.ievtcb))
        for n in range(nper):
            #comment = 'Evapotranspiration array for stress period ' + str(n + 1)
            if (n < len(self.surf)):
                insurf = 1
            else:
                insurf = -1
            if (n < len(self.evtr)):
                inevtr = 1
            else:
                inevtr = -1
            if (n < len(self.exdp)):
                inexdp = 1
            else:
                inexdp = -1
            if (n < len(self.ievt)):
                inievt = 1
            else:
                inievt = -1
            comment = 'Evapotranspiration surface array for stress period ' + str(n + 1)
            f_evt.write('{0:10d}{1:10d}{2:10d}{3:10d} # {4:s}\n'\
                .format(insurf,inevtr,inexdp,inievt,comment))
            
            if (n < len(self.surf)):
                f_evt.write(self.surf[n].get_file_entry())
            comment = 'Evapotranspiration rate array for stress period ' + str(n + 1)
            if (n < len(self.evtr)):
                f_evt.write(self.evtr[n].get_file_entry())
            comment = 'Evapotranspiration extinction depth array for stress period ' + str(n + 1)
            if (n < len(self.exdp)):
                f_evt.write(self.exdp[n].get_file_entry())
            comment = 'Evapotranspiration layers for stress period ' + str(n + 1)
            if ((n < len(self.ievt)) and (self.nevtop == 2)):
                f_evt.write(self.ievt[n].get_file_entry())
        f_evt.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
        -------
        evt : ModflowEvt object
            ModflowEvt object.

        Examples
        --------

        >>> import flopy
        >>> m = flopy.modflow.Modflow()
        >>> evt = flopy.modflow.mfevt.load('test.evt', m)

        """
        if type(f) is not file:
            filename = f
            f = open(filename, 'r')
        #dataset 0 -- header
        while True:
            line = f.readline()
            if line[0] != '#':
                break
        if "parameter" in line.lower():
            raw = line.strip().split()
            npar = int(raw[1])
            if npar > 0:
                print 'Parameters detected. Number of parameters = ', npar
            line = f.readline()
        #dataset 2
        t = line.strip().split()
        nevtop = int(t[0])
        ievtcb = 0
        try:
            if int(t[1]) != 0:
                ievtcb = 53
        except:
            pass

        #--dataset 3 and 4 - parameters data
        if npar > 0:
            pak_parms = mfparbc.loadarray(f, npar)


        if nper is None:
            nrow, ncol, nlay, nper = model.get_nrow_ncol_nlay_nper()
        #read data for every stress period
        surf = {}
        evtr = {}
        exdp = {}
        ievt = {}
        current_surf = []
        current_evtr = []
        current_exdp = []
        current_ievt = []
        for iper in xrange(nper):
            line = f.readline()
            t = line.strip().split()
            insurf = int(t[0])
            inevtr = int(t[1])
            inexdp = int(t[2])
            if (nevtop == 2):
              inievt = int(t[3])
            if insurf >= 0:
                print \
                    '   loading surf stress period {0:3d}...'.format(iper+1)
                t = util_2d.load(f, model, (nrow,ncol), np.float32, 'surf',
                                 ext_unit_dict)
                current_surf = t
            surf[iper] = current_surf

            if inevtr >= 0:
                if npar == 0:
                    print \
                        '   loading evtr stress period {0:3d}...'.format(iper+1)
                    t = util_2d.load(f, model, (nrow,ncol), np.float32, 'evtr',
                                     ext_unit_dict)
                else:
                    parm_dict = {}
                    for ipar in xrange(npar):
                        line = f.readline()
                        t = line.strip().split()
                        pname = t[0].lower()
                        try:
                            iname = t[1].lower()
                        except:
                            iname = 'static'
                        parm_dict[pname] = iname
                    t = mfparbc.parameter_bcfill(model, (nrow, ncol), 'rech', parm_dict, pak_parms)

                current_evtr = t
            evtr[iper] = current_evtr
            if inexdp >= 0:
                print \
                    '   loading exdp stress period {0:3d}...'.format(iper+1)
                t = util_2d.load(f, model, (nrow,ncol), np.float32, 'exdp',
                                 ext_unit_dict)
                current_exdp = t
            exdp[iper] = current_exdp
            if nevtop == 2:
                if inievt >= 0:
                    print '   loading ievt stress period {0:3d}...'.format(
                        iper+1)
                    t = util_2d.load(f, model, (nrow,ncol), np.int32, 'ievt',
                                     ext_unit_dict)
                    current_ievt = t
                ievt[iper] = current_ievt
        #--create evt object
        evt = ModflowEvt(model, nevtop=nevtop, ievtcb=ievtcb, 
                         surf=surf, evtr=evtr, exdp=exdp, ievt=ievt)
        return evt

#from numpy import ones, empty
import numpy as np
from flopy.mbase import Package
from flopy.utils import util_2d

class ModflowEvt(Package):
    'Evapotranspiration class'
    def __init__(self, model, nevtop=3, ievtcb=0, surf=0., evtr=1e-3, exdp=1., ievt=1, 
                 extension ='evt', unitnumber=22, external=True):
        '''
        external flag is used to control writing external arrays of constant value
        since this package has the potential to create a lot of external arrays
        '''
        Package.__init__(self, model, extension, 'EVT', unitnumber) # Call ancestor's init to set self.parent, extension, name and unit number
        nrow, ncol, nlay, nper = self.parent.nrow_ncol_nlay_nper
        self.heading = '# EVT for MODFLOW, generated by Flopy.'
        self.url = 'evt.htm'
        self.nevtop = nevtop
        self.ievtcb = ievtcb
        self.surf = []
        self.evtr = []
        self.exdp = []
        self.ievt = []        
        self.external = external
        if self.external is False:
            load = True
        else:
            load = model.load            
        if (not isinstance(surf, list)):
            surf = [surf]
        for i,a in enumerate(surf):
            s = util_2d(model,(nrow,ncol),np.float32,a,name='surf_'+str(i+1))
            self.surf = self.surf + [s]
        if (not isinstance(evtr, list)):
            evtr = [evtr]
        for i,a in enumerate(evtr):
            e = util_2d(model,(nrow,ncol),np.float32,a,name='etvr_'+str(i+1))
            self.evtr = self.evtr + [e]
        if (not isinstance(exdp, list)):
            exdp = [exdp]
        for i,a in enumerate(exdp):
            e = util_2d(model,(nrow,ncol),np.float32,a,name='exdp_'+str(i+1))
            self.exdp = self.exdp + [e]
        if (not isinstance(ievt, list)):
            ievt = [ievt]
        for i,a in enumerate(ievt):
            iv = util_2d(model,(nrow,ncol),np.int,a,name='ievt_'+str(i+1))
            self.ievt = self.ievt + [iv]
        self.np = 0
        self.parent.add_package(self)
    def __repr__( self ):
        return 'Evapotranspiration class'
    def ncells( self):
        # Returns the  maximum number of cells that have 
        # evapotranspiration (developped for MT3DMS SSM package)
        nrow, ncol, nlay, nper = self.parent.nrow_ncol_nlay_nper
        return (nrow * ncol)
    def write_file(self):
        nrow, ncol, nlay, nper = self.parent.nrow_ncol_nlay_nper
        f_evt = open(self.fn_path, 'w')        
        f_evt.write('{0:s}\n'.format(self.heading))        
        f_evt.write('{0:10d}{1:10d}\n'.format(self.nevtop,self.ievtcb))
        for n in range(nper):
            #comment = 'Evapotranspiration array for stress period ' + str(n + 1)
            if (n < len(self.surf)):
                insurf = 1
            else:
                insurf = -1
            if (n < len(self.evtr)):
                inevtr = 1
            else:
                inevtr = -1
            if (n < len(self.exdp)):
                inexdp = 1
            else:
                inexdp = -1
            if (n < len(self.ievt)):
                inievt = 1
            else:
                inievt = -1
            comment = 'Evapotranspiration surface array for stress period ' + str(n + 1)
            f_evt.write('{0:10d}{1:10d}{2:10d}{3:10d} # {4:s}\n'\
                .format(insurf,inevtr,inexdp,inievt,comment))
            
            if (n < len(self.surf)):
                f_evt.write(self.surf[n].get_file_entry())
            comment = 'Evapotranspiration rate array for stress period ' + str(n + 1)
            if (n < len(self.evtr)):
                f_evt.write(self.evtr[n].get_file_entry())
            comment = 'Evapotranspiration extinction depth array for stress period ' + str(n + 1)
            if (n < len(self.exdp)):
                f_evt.write(self.exdp[n].get_file_entry())
            comment = 'Evapotranspiration layers for stress period ' + str(n + 1)
            if ((n < len(self.ievt)) and (self.nevtop == 2)):
                f_evt.write(self.ievt[n].get_file_entry())
        f_evt.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
        -------
        evt : ModflowEvt object
            ModflowEvt object.

        Examples
        --------

        >>> import flopy
        >>> m = flopy.modflow.Modflow()
        >>> evt = flopy.modflow.mfevt.load('test.evt', m)

        """
        if type(f) is not file:
            filename = f
            f = open(filename, 'r')
        #dataset 0 -- header
        while True:
            line = f.readline()
            if line[0] != '#':
                break
        if "parameter" in line.lower():
            raw = line.strip().split()
            npar = int(raw[1])
            if npar > 0:
                print 'Parameters detected. Number of parameters = ', npar
            line = f.readline()
        #dataset 2
        t = line.strip().split()
        nevtop = int(t[0])
        ievtcb = 0
        try:
            if int(t[1]) != 0:
                ievtcb = 53
        except:
            pass

        #--dataset 3 and 4 - parameters data
        if npar > 0:
            pak_parms = mfparbc.loadarray(f, npar)


        if nper is None:
            nrow, ncol, nlay, nper = model.get_nrow_ncol_nlay_nper()
        #read data for every stress period
        surf = {}
        evtr = {}
        exdp = {}
        ievt = {}
        current_surf = []
        current_evtr = []
        current_exdp = []
        current_ievt = []
        for iper in xrange(nper):
            line = f.readline()
            t = line.strip().split()
            insurf = int(t[0])
            inevtr = int(t[1])
            inexdp = int(t[2])
            if (nevtop == 2):
              inievt = int(t[3])
            if insurf >= 0:
                print \
                    '   loading surf stress period {0:3d}...'.format(iper+1)
                t = util_2d.load(f, model, (nrow,ncol), np.float32, 'surf',
                                 ext_unit_dict)
                current_surf = t
            surf[iper] = current_surf

            if inevtr >= 0:
                if npar == 0:
                    print \
                        '   loading evtr stress period {0:3d}...'.format(iper+1)
                    t = util_2d.load(f, model, (nrow,ncol), np.float32, 'evtr',
                                     ext_unit_dict)
                else:
                    parm_dict = {}
                    for ipar in xrange(npar):
                        line = f.readline()
                        t = line.strip().split()
                        pname = t[0].lower()
                        try:
                            iname = t[1].lower()
                        except:
                            iname = 'static'
                        parm_dict[pname] = iname
                    t = mfparbc.parameter_bcfill(model, (nrow, ncol), 'rech', parm_dict, pak_parms)

                current_evtr = t
            evtr[iper] = current_evtr
            if inexdp >= 0:
                print \
                    '   loading exdp stress period {0:3d}...'.format(iper+1)
                t = util_2d.load(f, model, (nrow,ncol), np.float32, 'exdp',
                                 ext_unit_dict)
                current_exdp = t
            exdp[iper] = current_exdp
            if nevtop == 2:
                if inievt >= 0:
                    print '   loading ievt stress period {0:3d}...'.format(
                        iper+1)
                    t = util_2d.load(f, model, (nrow,ncol), np.int32, 'ievt',
                                     ext_unit_dict)
                    current_ievt = t
                ievt[iper] = current_ievt
        #--create evt object
        evt = ModflowEvt(model, nevtop=nevtop, ievtcb=ievtcb, 
                         surf=surf, evtr=evtr, exdp=exdp, ievt=ievt)
        return evt
#from numpy import ones, empty external flag is used to control writing external arrays of constant value
         since this package has the potential to create a lot of external arrays # Call ancestor's init to set self.parent, extension, name and unit number # Returns the  maximum number of cells that have # evapotranspiration (developped for MT3DMS SSM package) #comment = 'Evapotranspiration array for stress period ' + str(n + 1) # {4:s}\n'\ 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
         -------
         evt : ModflowEvt object
             ModflowEvt object.
 
         Examples
         --------
 
         >>> import flopy
         >>> m = flopy.modflow.Modflow()
         >>> evt = flopy.modflow.mfevt.load('test.evt', m) #dataset 0 -- header #dataset 2 #--dataset 3 and 4 - parameters data #read data for every stress period #--create evt object