import numpy as np
import flopy
reload(flopy)
from flopy.modflow import *
from flopy.mt3d import *

nlay, nrow, ncol = 3, 10, 10
ml = Modflow("test")
dis = ModflowDis(ml,nlay=nlay, nrow=nrow, ncol=ncol)


"""
Explanation: Demo for new kwargs-based ncomp
End of explanation
"""


mt = Mt3dms(modflowmodel=ml)
btn = Mt3dBtn(mt, sconc=1.0,ncomp=2)

"""
Explanation: for btn, sconc is passed as a normal util_3d-compatible argument, no need for the list anymore
End of explanation
"""


mt = Mt3dms(modflowmodel=ml)
btn = Mt3dBtn(mt,sconc=1.0,ncomp=2,sconc2=1.5)
print btn.sconc[1].name_base

"""
Explanation: or we pass in the sconc2 kwarg explicitly
End of explanation
"""


mt = Mt3dms(modflowmodel=ml)
btn = Mt3dBtn(mt,sconc=1.0,ncomp=2,sconc3=1.5)

"""
Explanation: error raised for unrecognized kwargs
End of explanation
"""


mt = Mt3dms(modflowmodel=ml)
btn = Mt3dBtn(mt,sconc=1.0,ncomp=2, mcomp=2, sconc2=1.5)
ssm = Mt3dSsm(mt,crch=1.0,crch2=1.0,cevt=0.5,cevt2=1.5)
print ssm.crch,ssm.cevt

"""
Explanation: same behavior for SSM
End of explanation
"""


dsp = Mt3dDsp(mt,dmcoef2=1.0)

"""
Explanation: DSP does not need multDiff flag for multi species.  But in this case, arrays are by layer, not by layer, row, column.
End of explanation
"""


dsp = Mt3dDsp(mt,dmcoef2=1.0,multiDiff=True)
print dsp.dmcoef

"""
Explanation: If we set multiDiff to True, then there will be mcomp (not ncomp, according to mt3d supplemental guide) number of 3d dmcoef arrays.
End of explanation
"""


