import collections
import os
import shutil
import tempfile
import unittest

import numpy
import shapely.geometry
from osgeo import gdal
from osgeo import ogr
from osgeo import osr


class WatershedDelineationTests(unittest.TestCase):
    def setUp(self):
        self.workspace_dir = tempfile.mkdtemp()

    def tearDown(self):
        shutil.rmtree(self.workspace_dir)

    def test_watersheds_trivial(self):
        """PGP watersheds: test trivial delineation."""
        import pygeoprocessing.testing
        import pygeoprocessing.routing

        flow_dir_array = numpy.array([
            [6, 6, 6, 6, 6, 6, 6, 6, 6, 6],
            [6, 6, 6, 6, 6, 6, 6, 6, 6, 6],
            [6, 6, 6, 6, 6, 6, 6, 6, 6, 6],
            [6, 6, 6, 6, 6, 6, 6, 6, 6, 255],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [2, 2, 2, 2, 2, 2, 2, 2, 2, 255],
            [2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
            [2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
            [2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
            [2, 2, 2, 2, 2, 2, 2, 2, 2, 2]],
            dtype=numpy.int8)

        srs = osr.SpatialReference()
        srs.ImportFromEPSG(32731)  # WGS84 / UTM zone 31s
        srs_wkt = srs.ExportToWkt()

        flow_dir_path = os.path.join(self.workspace_dir, 'flow_dir.tif')
        driver = gdal.GetDriverByName('GTiff')
        flow_dir_raster = driver.Create(
            flow_dir_path, flow_dir_array.shape[1], flow_dir_array.shape[0],
            1, gdal.GDT_Byte, options=(
                'TILED=YES', 'BIGTIFF=YES', 'COMPRESS=LZW',
                'BLOCKXSIZE=256', 'BLOCKYSIZE=256'))
        flow_dir_raster.SetProjection(srs_wkt)
        flow_dir_band = flow_dir_raster.GetRasterBand(1)
        flow_dir_band.WriteArray(flow_dir_array)
        flow_dir_band.SetNoDataValue(255)
        flow_dir_geotransform = [2, 2, 0, -2, 0, -2]
        flow_dir_raster.SetGeoTransform(flow_dir_geotransform)
        flow_dir_raster = None

        # These geometries test:
        #  * Delineation works with varying geometry types
        #  * That we exclude seed pixels that are over nodata
        #  * That we exclude seed pixels off the bounds of the raster
        horizontal_line = shapely.geometry.LineString([(19, -11), (25, -11)])
        vertical_line = shapely.geometry.LineString([(21, -9), (21, -13)])
        square = shapely.geometry.box(17, -13, 21, -9)
        point = shapely.geometry.Point(21, -11)

        outflow_vector_path = os.path.join(self.workspace_dir, 'outflow.gpkg')
        pygeoprocessing.testing.create_vector_on_disk(
            [horizontal_line, vertical_line, square, point],
            projection=srs_wkt,
            fields={'polygon_id': 'int',
                    'field_string': 'string',
                    'other': 'real'},
            attributes=[
                {'polygon_id': 1, 'field_string': 'hello world', 'other': 1.111},
                {'polygon_id': 2, 'field_string': 'hello foo', 'other': 2.222},
                {'polygon_id': 3, 'field_string': 'hello bar', 'other': 3.333},
                {'polygon_id': 4, 'field_string': 'hello baz', 'other': 4.444}
            ],
            vector_format='GPKG', filename=outflow_vector_path)

        target_watersheds_path = os.path.join(self.workspace_dir, 'watersheds.gpkg')

        pygeoprocessing.routing.delineate_watersheds_d8(
            (flow_dir_path, 1), outflow_vector_path, target_watersheds_path)

        watersheds_vector = gdal.OpenEx(target_watersheds_path, gdal.OF_VECTOR)
        watersheds_layer = watersheds_vector.GetLayer('watersheds')
        self.assertEqual(watersheds_layer.GetFeatureCount(), 4)

        # All features should have the same watersheds, both in area and
        # geometry.
        flow_dir_bbox = pygeoprocessing.get_raster_info(flow_dir_path)['bounding_box']
        expected_watershed_geometry = shapely.geometry.box(*flow_dir_bbox)
        expected_watershed_geometry = expected_watershed_geometry.difference(
            shapely.geometry.box(20, -2, 22, -10))
        expected_watershed_geometry = expected_watershed_geometry.difference(
            shapely.geometry.box(20, -12, 22, -22))
        pygeoprocessing.testing.create_vector_on_disk(
            [expected_watershed_geometry],
            srs_wkt, vector_format='GPKG', filename=os.path.join(self.workspace_dir, 'foo.gpkg'))

        id_to_fields = {}
        for feature in watersheds_layer:
            geometry = feature.GetGeometryRef()
            shapely_geom = shapely.wkb.loads(geometry.ExportToWkb())
            self.assertEqual(shapely_geom.area, expected_watershed_geometry.area)
            self.assertEqual(
                shapely_geom.intersection(
                    expected_watershed_geometry).area,
                expected_watershed_geometry.area)
            self.assertEqual(
                shapely_geom.difference(
                    expected_watershed_geometry).area, 0)

            field_values = feature.items()
            id_to_fields[field_values['polygon_id']] = field_values

        outflow_vector = gdal.OpenEx(outflow_vector_path, gdal.OF_VECTOR)
        outflow_layer = outflow_vector.GetLayer()
        try:
            for feature in outflow_layer:
                self.assertEqual(
                    id_to_fields[feature.GetField('polygon_id')],
                    feature.items())
        finally:
            outflow_layer = None
            outflow_vector = None

    def test_split_geometry_into_seeds(self):
        """PGP watersheds: Test geometry-to-seed extraction."""
        import pygeoprocessing
        from pygeoprocessing.routing import watershed
        nodata = 255
        flow_dir_array= numpy.array([
            [0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0],
            [nodata, nodata, nodata, nodata, nodata, nodata, nodata],
            [0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 2, 2]], dtype=numpy.uint8)
        srs = osr.SpatialReference()
        srs.ImportFromEPSG(32731)  # WGS84 / UTM zone 31s
        srs_wkt = srs.ExportToWkt()

        flow_dir_path = os.path.join(self.workspace_dir, 'flow_dir.tif')
        driver = gdal.GetDriverByName('GTiff')
        flow_dir_raster = driver.Create(
            flow_dir_path, flow_dir_array.shape[1], flow_dir_array.shape[0],
            1, gdal.GDT_Byte, options=(
                'TILED=YES', 'BIGTIFF=YES', 'COMPRESS=LZW',
                'BLOCKXSIZE=256', 'BLOCKYSIZE=256'))
        flow_dir_raster.SetProjection(srs_wkt)
        flow_dir_band = flow_dir_raster.GetRasterBand(1)
        flow_dir_band.WriteArray(flow_dir_array)
        flow_dir_band.SetNoDataValue(255)
        flow_dir_geotransform = [2, 2, 0, -2, 0, -2]
        flow_dir_raster.SetGeoTransform(flow_dir_geotransform)
        flow_dir_raster = None
        flow_dir_info = pygeoprocessing.get_raster_info(flow_dir_path)

        point = shapely.geometry.Point(2.5, -2.5)
        linestring = shapely.geometry.LineString([(10, -2), (10, -9.9)])
        box = shapely.geometry.box(4.1, -7.9, 7.9, -4.1)

        for index, (geometry, expected_seeds) in enumerate((
                (point, set([(0, 0)])),
                (linestring, set([(4, 0), (4, 1), (4, 2), (4, 3)])),  # includes nodata pixel
                (box, set([(1, 1), (2, 1), (1, 2), (2, 2)])))):  # includes nodata pixels

            raster_path = os.path.join(self.workspace_dir, '%s.tif' % index)
            diagnostic_path = os.path.join(self.workspace_dir, '%s.gpkg' % index)
            result_seeds = watershed._split_geometry_into_seeds(
                geometry.wkb, flow_dir_info['geotransform'], srs,
                flow_dir_array.shape[1], flow_dir_array.shape[0],
                raster_path, diagnostic_path)

            self.assertEqual(result_seeds, expected_seeds)

    def test_split_geometry_into_seeds_willamette(self):
        """PGP watersheds: Test geometry-to-seed extraction in Willamette."""
        import pygeoprocessing
        from pygeoprocessing.routing import watershed
        nodata = 255
        flow_dir_array= numpy.array([
            [0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0]], dtype=numpy.uint8)
        srs = osr.SpatialReference()
        srs.ImportFromEPSG(3157)  # UTM zone 17N
        srs_wkt = srs.ExportToWkt()

        flow_dir_path = os.path.join(self.workspace_dir, 'flow_dir.tif')
        driver = gdal.GetDriverByName('GTiff')
        flow_dir_raster = driver.Create(
            flow_dir_path, flow_dir_array.shape[1], flow_dir_array.shape[0],
            1, gdal.GDT_Byte, options=(
                'TILED=YES', 'BIGTIFF=YES', 'COMPRESS=LZW',
                'BLOCKXSIZE=256', 'BLOCKYSIZE=256'))
        flow_dir_raster.SetProjection(srs_wkt)
        flow_dir_band = flow_dir_raster.GetRasterBand(1)
        flow_dir_band.WriteArray(flow_dir_array)
        flow_dir_band.SetNoDataValue(255)
        pixel_xsize = 30
        pixel_ysize = -30
        flow_dir_geotransform = [
            443723, pixel_xsize, 0,
            4956546, 0, pixel_ysize]
        flow_dir_raster.SetGeoTransform(flow_dir_geotransform)
        flow_dir_raster = None
        flow_dir_info = pygeoprocessing.get_raster_info(flow_dir_path)

        pixel_indexes_array = numpy.arange(
            flow_dir_array.size).reshape(flow_dir_array.shape)
        pixel_indexes_path = os.path.join(self.workspace_dir, 'pixel_indexes.tif')
        driver = gdal.GetDriverByName('GTiff')
        pixel_indexes_raster = driver.Create(
            pixel_indexes_path, pixel_indexes_array.shape[1], pixel_indexes_array.shape[0],
            1, gdal.GDT_Byte, options=(
                'TILED=YES', 'BIGTIFF=YES', 'COMPRESS=LZW',
                'BLOCKXSIZE=256', 'BLOCKYSIZE=256'))
        pixel_indexes_raster.SetProjection(srs_wkt)
        pixel_indexes_band = pixel_indexes_raster.GetRasterBand(1)
        pixel_indexes_band.WriteArray(pixel_indexes_array)
        pixel_indexes_band.SetNoDataValue(255)
        pixel_xsize = 30
        pixel_ysize = -30
        pixel_indexes_geotransform = [
            443723, pixel_xsize, 0,
            4956546, 0, pixel_ysize]
        pixel_indexes_raster.SetGeoTransform(pixel_indexes_geotransform)
        pixel_indexes_raster = None

        point = shapely.geometry.Point(
            flow_dir_geotransform[0] + pixel_xsize / 2.,
            flow_dir_geotransform[3] + pixel_ysize / 2.)
        linestring = shapely.geometry.LineString([
            (flow_dir_geotransform[0] + pixel_xsize * 4,
             flow_dir_geotransform[3] - pixel_ysize * 2),  # extend beyond y boundary
            (flow_dir_geotransform[0] + pixel_xsize * 4,
             flow_dir_geotransform[3] + pixel_ysize * 5)])
        box = shapely.geometry.box(
            flow_dir_geotransform[0] + pixel_xsize * 2.1,
            flow_dir_geotransform[3] + pixel_ysize * 4.1,
            flow_dir_geotransform[0] + pixel_xsize * 3.9,
            flow_dir_geotransform[3] + pixel_ysize * 2.1)

        for index, (geometry, expected_seeds) in enumerate((
                (point, set([(0, 0)])),
                (linestring, set([(4, 0), (4, 1), (4, 2), (4, 3)])),
                (box, set([(2, 2), (2, 3), (3, 2), (3, 3)])))):

            result_seeds = watershed._split_geometry_into_seeds(
                geometry.wkb, flow_dir_info['geotransform'], srs,
                flow_dir_array.shape[1], flow_dir_array.shape[0],
                os.path.join(self.workspace_dir, '%s.tif' % index),
                os.path.join(self.workspace_dir, '%s.gpkg' % index))

            self.assertEqual(result_seeds, expected_seeds)
