import rasterio
from pkg_resources import resource_filename

import pandas as pd
import numpy as np

from .spatial_context import SpatialContext


class FloodPlainError(Exception):
    """"""


class FloodPlain(object):
    def __init__(self, depths=None, duration=None, frequency=None,
                 lnk_potential=None, potential=None):
        self._ct = dict()
        self._veg = dict()

        for i in ["duration", "frequency", "lnk_potential", "potential"]:
            if locals()[i] is None:
                ct = resource_filename(
                        "niche_vlaanderen",
                        "system_tables/floodplains/"+i+".csv")
            else:
                ct = locals()[i]
            self._ct[i] = pd.read_csv(ct)

    def _calculate(self, depth, frequency, duration, period):
        orig_shape = depth.shape
        depth = depth.flatten()

        for veg_code, subtable_veg in \
                self._ct["lnk_potential"].groupby(["veg_code"]):
            subtable_veg = subtable_veg.reset_index()
            self._veg[veg_code] = np.zeros(depth.shape, dtype="uint8")
            groupby_cols = ["period", "frequency", "duration"]
            for index, subtable in subtable_veg.groupby(groupby_cols):
                if (period, frequency, duration) == index:
                    subtable.reset_index()
                    for row in subtable.itertuples():
                        veg = self._veg[veg_code]
                        veg[row.depth == depth] = row.potential
                        self._veg[veg_code] = np.maximum(veg,
                                                         self._veg[veg_code])
            self._veg[veg_code] = self._veg[veg_code].reshape(orig_shape)

    def calculate(self, depth_file, frequency, period, duration):
        with rasterio.open(depth_file) as dst:
            depth = dst.read(1)
            self._context = SpatialContext(dst)
        self._calculate(depth, frequency, duration, period)
        self.options = (frequency, duration, period)

    def show(self, key):
        plt = rasterio.plot.get_plt()
        import matplotlib.patches as mpatches
        from matplotlib.colors import Normalize

        if key not in self._veg.keys():
            print("vegetation type not modeled")
            return

        fig, ax = plt.subplots()
        ((a, b), (c, d)) = self._context.extent
        mpl_extent = (a, c, d, b)

        im = plt.imshow(self._veg[key], extent=mpl_extent,
                        norm=Normalize(0, 3))
        options = list(self.options)
        options[2] = "< 14 days" if self.options[2] == 1 else "> 14 days"
        ax.set_title("{} ({})".format(key, options))

        labels = self._ct["potential"]["description"]
        values = self._ct["potential"]["code"]

        colors = [im.cmap(value/(len(values) - 1)) for value in values]
        patches = [mpatches.Patch(color=colors[i],
                                  label=labels[i]) for i in values]
        plt.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2,
                   borderaxespad=0.)

        plt.show()

    def combine(self, niche_result):
        # check niche model has been run
        if not hasattr(niche_result, "_vegetation"):
            raise FloodPlainError(
                "Niche model must be run prior to running this module.")

@@ -9,16 +9,26 @@
 asterio%0A
+import os%0A
 from pkg
@@ -167,20 +167,24 @@
 odPlainE
-rror
+xception
 (Excepti
@@ -1983,16 +1983,17 @@
 period)%0A
+%0A
         
@@ -3107,24 +3107,1033 @@
 plt.show()%0A%0A
+    def write(self, folder):%0A        if not hasattr(self, %22_veg%22):%0A            raise FloodPlainException(%0A                %22A valid run must be done before writing the output.%22)%0A%0A        if not os.path.exists(folder):%0A            os.makedirs(folder)%0A%0A        params = dict(%0A            driver='GTiff',%0A            height=self._context.height,%0A            width=self._context.width,%0A            crs=self._context.crs,%0A            transform=self._context.transform,%0A            count=1,%0A            dtype=%22uint8%22,%0A            nodata=255,%0A            compress=%22DEFLATE%22%0A        )%0A%0A        self._files_written = dict()%0A%0A        for vi in self._veg:%0A            filename = %22F%7B%7D-%7B%7D-P%7B%7D-%7B%7D.tif%22.format(vi, self.options%5B0%5D,%0A                                              self.options%5B1%5D, self.options%5B2%5D)%0A            path = folder + %22/%22 + filename%0A            with rasterio.open(path, 'w', **params) as dst:%0A                dst.write(self._veg%5Bvi%5D, 1)%0A                self._files_written%5Bvi%5D = os.path.normpath(path)%0A%0A%0A
     def comb
@@ -4284,12 +4284,16 @@
 ainE
-rror
+xception
 (%0A  

Add write function to floodplain module