%matplotlib inline
from __future__ import print_function
import os
import sys
import platform
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import pandas as pd
import flopy

print(sys.version)
print('numpy version: {}'.format(np.__version__))
print('matplotlib version: {}'.format(mpl.__version__))
print('pandas version: {}'.format(pd.__version__))
print('flopy version: {}'.format(flopy.__version__))

# Set path to example datafiles
loadpth = os.path.join('..', 'data', 'zonbud_examples')
cbc_f = os.path.join(loadpth, 'freyberg_mlt', 'freyberg.gitcbc')

"""
Explanation: FloPy
ZoneBudget Example
This notebook demonstrates how to use the ZoneBudget class to extract budget information from the cell by cell budget file using an array of zones.
First set the path and import the required packages. The flopy path doesn't have to be set if you install flopy from a binary installer. If you want to run this notebook, you have to set the path to your own flopy path.
End of explanation
"""


from flopy.utils import read_zbarray

zone_file = os.path.join(loadpth, 'zonef_mlt')
zon = read_zbarray(zone_file)
nlay, nrow, ncol = zon.shape

fig = plt.figure(figsize=(10, 4))

for lay in range(nlay):
    ax = fig.add_subplot(1, nlay, lay+1)
    im = ax.pcolormesh(zon[lay, :, :])
    cbar = plt.colorbar(im)
    plt.gca().set_aspect('equal')
    
plt.show()
np.unique(zon)

"""
Explanation: Read File Containing Zones
Using the read_zbarray utility, we can import zonebudget-style array files.
End of explanation
"""


# Create a ZoneBudget object and get the budget record array
zb = flopy.utils.ZoneBudget(cbc_f, zon, kstpkper=(0, 1096))
zb.get_budget()

# Get a list of the unique budget record names
zb.get_record_names()

# Look at a subset of fluxes
names = ['RECHARGE_IN', 'ZONE_1_IN', 'ZONE_3_IN']
zb.get_budget(names=names)

# Look at fluxes in from zone 2
names = ['RECHARGE_IN', 'ZONE_1_IN', 'ZONE_3_IN']
zones = ['ZONE_2']
zb.get_budget(names=names, zones=zones)

# Look at all of the mass-balance records
names = ['TOTAL_IN', 'TOTAL_OUT', 'IN-OUT', 'PERCENT_DISCREPANCY']
zb.get_budget(names=names)

"""
Explanation: Extract Budget Information from ZoneBudget Object
At the core of the ZoneBudget object is a numpy structured array. The class provides some wrapper functions to help us interogate the array and save it to disk.
End of explanation
"""


cmd = flopy.utils.ZoneBudget(cbc_f, zon, kstpkper=(0, 0))
cfd = cmd / 35.3147
inyr = (cfd / (250 * 250)) * 365 * 12

cmdbud = cmd.get_budget()
cfdbud = cfd.get_budget()
inyrbud = inyr.get_budget()

names = ['RECHARGE_IN']
rowidx = np.in1d(cmdbud['name'], names)
colidx = 'ZONE_1'

print('{:,.1f} cubic meters/day'.format(cmdbud[rowidx][colidx][0]))
print('{:,.1f} cubic feet/day'.format(cfdbud[rowidx][colidx][0]))
print('{:,.1f} inches/year'.format(inyrbud[rowidx][colidx][0]))

cmd is cfd

"""
Explanation: Convert Units
The ZoneBudget class supports the use of mathematical operators and returns a new copy of the object.
End of explanation
"""


aliases = {1: 'SURF', 2:'CONF', 3: 'UFA'}
zb = flopy.utils.ZoneBudget(cbc_f, zon, totim=[1097.], aliases=aliases)
zb.get_budget()

"""
Explanation: Alias Names
A dictionary of {zone: "alias"} pairs can be passed to replace the typical "ZONE_X" fieldnames of the ZoneBudget structured array with more descriptive names.
End of explanation
"""


zon = np.ones((nlay, nrow, ncol), np.int)
zon[1, :, :] = 2
zon[2, :, :] = 3

aliases = {1: 'SURF', 2:'CONF', 3: 'UFA'}
zb = flopy.utils.ZoneBudget(cbc_f, zon, kstpkper=None, totim=None, aliases=aliases)
df = zb.get_dataframes()
print(df.head())
print(df.tail())

"""
Explanation: Return the Budgets as a Pandas DataFrame
Set kstpkper and totim keyword args to None (or omit) to return all times.
The get_dataframes() method will return a DataFrame multi-indexed on totim and name.
End of explanation
"""


dateidx1 = 1092.
dateidx2 = 1097.
names = ['RECHARGE_IN', 'WELLS_OUT']
zones = ['SURF', 'CONF']
df.loc[(slice(dateidx1, dateidx2), names), :][zones]

"""
Explanation: Slice the multi-index dataframe to retrieve a subset of the budget
End of explanation
"""


dateidx1 = 1092.
dateidx2 = 1097.
zones = ['SURF']

# Pull out the individual records of interest
rech = df.loc[(slice(dateidx1, dateidx2), ['RECHARGE_IN']), :][zones]
pump = df.loc[(slice(dateidx1, dateidx2), ['WELLS_OUT']), :][zones]

# Remove the "record" field from the index so we can 
# take the difference of the two DataFrames
rech = rech.reset_index()
rech = rech.set_index(['totim'])
rech = rech[zones]
pump = pump.reset_index()
pump = pump.set_index(['totim'])
pump = pump[zones] * -1

# Compute pumping as a percentage of recharge
pump_as_pct = (pump / rech) * 100.
pump_as_pct

# Use "slice(None)" to return all records
df.loc[(slice(dateidx1, dateidx2), slice(None)), :][zones]

# Or all times
df.loc[(slice(None), names), :][zones]

"""
Explanation: Look at pumpage (WELLS_OUT) as a percentage of recharge (RECHARGE_IN)
End of explanation
"""


df = zb.get_dataframes(start_datetime='1970-01-01', timeunit='D')
dateidx1 = pd.Timestamp('1972-12-01')
dateidx2 = pd.Timestamp('1972-12-06')
names = ['RECHARGE_IN', 'WELLS_OUT']
zones = ['SURF', 'CONF']
df.loc[(slice(dateidx1, dateidx2), names), :][zones]

"""
Explanation: Pass start_datetime and timeunit keyword arguments to return a dataframe with a datetime multi-index
End of explanation
"""


df = zb.get_dataframes(index_key='kstpkper')
df.head()

"""
Explanation: Pass index_key to indicate which fields to use in the multi-index (defualt is "totim"; valid keys are "totim" and "kstpkper")
End of explanation
"""


zb = flopy.utils.ZoneBudget(cbc_f, zon, kstpkper=[(0, 0), (0, 1096)])
zb.to_csv(os.path.join(loadpth, 'zonbud.csv'))

# Read the file in to see the contents
fname = os.path.join(loadpth, 'zonbud.csv')
try:
    import pandas as pd
    print(pd.read_csv(fname).to_string(index=False))
except:
    with open(fname, 'r') as f:
        for line in f.readlines():
            print('\t'.join(line.split(',')))

"""
Explanation: Write Budget Output to CSV
We can write the resulting recarray to a csv file with the .to_csv() method of the ZoneBudget object.
End of explanation
"""


