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

class ModflowRch(Package):
    'Recharge class'
    def __init__(self, model, nrchop=3, irchcb=0, rech=1e-3, irch=1, extension ='rch', unitnumber=19,bin=False):
        '''
       
        '''
        Package.__init__(self, model, extension, 'RCH', 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 = '# RCH for MODFLOW, generated by Flopy.'
        self.url = 'rch.htm'
        self.nrchop = nrchop
        self.irchcb = irchcb
        self.rech = []
        self.irch = []        
        if (not isinstance(rech, list)):
            rech = [rech]
        for i,a in enumerate(rech):
            r = util_2d(model,(nrow,ncol),np.float32,a,name='rech_'+str(i+1),bin=bin)
            self.rech = self.rech + [r]
        if (not isinstance(irch, list)):
            irch = [irch]
        for i,a in enumerate(irch):
            ir = util_2d(model,(nrow,ncol),np.int,a,name='irech_'+str(i+1))
            self.irch = self.irch + [ir]
        self.np = 0
        self.parent.add_package(self)
    def __repr__( self ):
        return 'Recharge class'
    def ncells( self):
        # Returns the  maximum number of cells that have recharge (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
        # Open file for writing
        f_rch = open(self.fn_path, 'w')
        f_rch.write('{0:s}\n'.format(self.heading))
        f_rch.write('{0:10d}{1:10d}\n'.format(self.nrchop,self.irchcb))
        for n in range(nper):
            if (n < len(self.rech)):
                inrech = 1
            else:
                inrech = -1
            if (n < len(self.irch)):
                inirch = 1
            else:
                inirch = -1
            comment = 'Recharge array for stress period ' + str(n + 1)
            f_rch.write('{0:10d}{1:10d} #{2:s}\n'.format(inrech,inirch,comment))
            if (n < len(self.rech)):
                f_rch.write(self.rech[n].get_file_entry())
            if ((n < len(self.irch)) and (self.nrchop == 2)):
                f_rch.write(self.irch[n].get_file_entry())
        f_rch.close()

