'''
adapted from Andrew T Leaf's ipython notebook
https://github.com/aleaf/Notebooks/blob/master/rasterize_features_simple.ipynb
create an IBOUND array
'''
import matplotlib.pyplot as plt
from shapely.geometry import mapping
#import rasterio
from rasterio import features
from rasterio import Affine
#from flopy.utils import SpatialReference
from GISio import shp2df, get_proj4

# model grid info
xul, yul = 522000., 8420000.
rotation = 0

dx = 250. # cell spacing in meters
dy = 250.
nrow, ncol = 40, 50 # number of rows and columns

# finer grid
#dx = 100. # cell spacing in meters
#dy = 100.
#nrow, ncol = 100, 125 # number of rows and columns

fname = './Unit_boundaries.shp'

df = shp2df(fname)
print(df.geometry[0].bounds) # domain extent

# convert feature to GeoJSON
feature_gj = mapping(df.geometry[0])

# create a list of (feature, number) tuples
# the number for each feature will be assigned to the intersecting raster cells
shapes = [(feature_gj, 1)]

# create a rasterio.Affine reference for the grid
trans = Affine(dx, 0, xul,
               0, -dy, yul)

# "rasterize" the features to a numpy array
result = features.rasterize(shapes, out_shape=(nrow, ncol), transform=trans)

print(result) # this is your IBOUND for mf

# plot IBOUND array
plt.figure()
plt.imshow(result) # image show
plt.title('IBOUND')
plt.show()

# get the proj4 string for the shape file
proj4 = get_proj4(fname)

# export to geotiff
#meta = {'count': 1,
#        'dtype': result.dtype,
#        'driver': 'GTiff',
#        'height': result.shape[0],
#        'width': result.shape[1],
#        'crs': proj4,
#        'transform': trans}

#with rasterio.open('./ibound.tif', 'w', **meta) as dest:
#    dest.write(result, 1)