"""Module to handle processing the raw GeoTIFF data from satellite imagery."""

import rasterio
import rasterio.warp
import itertools
import numpy as np
import itertools
from osm_feature_extraction import extract_water


def read_geotiff(file_name):
    """TODO: Docstring."""
    raster_dataset = rasterio.open(file_name)
    bands = read_bands(raster_dataset)
    return raster_dataset, bands

def read_bands(raster_dataset):
    bands = [raster_dataset.read(band_number)
             for band_number in raster_dataset.indexes]
    bands = np.dstack(bands)
    return bands


def get_bounds(raster_dataset):
    """TODO: Docstring."""
    # Coordinate reference system of the GeoTIFF.
    src_crs = raster_dataset.crs
    # Destination coordinate reference system. We want to get the longitude
    # and latitude so we use EPSG:4326 also known as WGS 84.
    dst_crs = "EPSG:4326"
    west, south, east, north = rasterio.warp.transform_bounds(
        src_crs, dst_crs, *raster_dataset.bounds)
    return {
        "west": west,
        "south": south,
        "east": east,
        "north": north
    }


def read_bitmap(file_name):
    """TODO: Docstring."""
    # TODO: Outputshape of bitmap?
    raster_dataset, bitmap = read_geotiff(file_name)
    bitmap[bitmap == 255] = 1
    return raster_dataset, bitmap


def create_tiles(bands_data, tile_size, path_to_geotiff):
    """From https://github.com/trailbehind/DeepOSM."""
    # TODO: Select bands.

    rows, cols, n_bands = bands_data.shape

    all_tiled_data = []

    tile_indexes = itertools.product(range(0, rows, tile_size), range(0, cols, tile_size))

    for (row, col) in tile_indexes:
        in_bounds = row + tile_size < rows and col + tile_size < cols
        if in_bounds:
            new_tile = bands_data[row:row + tile_size, col:col + tile_size, 0:n_bands]
            all_tiled_data.append((new_tile, (row, col), path_to_geotiff)) 

    """
    for row in range(0, cols, tile_size):
        for col in range(0, rows, tile_size):
            in_bounds = row + tile_size < rows and col + tile_size < cols
            if in_bounds:
                new_tile = bands_data[row:row + tile_size, col:col + tile_size, 0:n_bands]
                all_tiled_data.append((new_tile, (row, col), path_to_geotiff)) 
    """

    return all_tiled_data

def image_from_tiles(tiles, tile_size, image_shape):
    image = np.zeros(image_shape, dtype=np.uint8)

    for tile, (row, col), _ in tiles:
        image[row:row + tile_size, col:col + tile_size, :] = tile

    return image


def overlay_bitmap(bitmap, raster_dataset, out_path):
    # TODO: Choose color.
    red, green, blue = raster_dataset.read()
    red[bitmap == 1] = 0
    green[bitmap == 1] = 0
    blue[bitmap == 1] = 255
    profile = raster_dataset.profile
    with rasterio.open(out_path, 'w', **profile) as dst:
        dst.write(red, 1)
        dst.write(green, 2)
        dst.write(blue, 3)


def create_bitmap(raster_dataset):
    """TODO: Docstring"""
    bounds = get_bounds(raster_dataset)
    water_features = extract_water(bounds)
    bitmap = np.zeros(raster_dataset.shape, dtype=np.int)
    for feature in water_features:
        feature = [ lat_lon_to_pixel(point, raster_dataset) for point in feature ]
        bitmap = add_feature(bitmap, feature)
    return bitmap

@@ -389,16 +389,17 @@
  bands%0A%0A
+%0A
 def read
@@ -435,16 +435,25 @@
 ands = %5B
+%0A        
 raster_d
@@ -477,21 +477,16 @@
 number)%0A
-     
         
@@ -526,16 +526,21 @@
 .indexes
+%0A    
 %5D%0A    ba
@@ -1011,25 +1011,16 @@
 return %7B
-%0A        
 %22west%22: 
@@ -1024,24 +1024,16 @@
 %22: west,
-%0A       
  %22south%22
@@ -1040,24 +1040,16 @@
 : south,
-%0A       
  %22east%22:
@@ -1054,24 +1054,16 @@
 %22: east,
-%0A       
  %22north%22
@@ -1069,21 +1069,16 @@
 %22: north
-%0A    
 %7D%0A%0A%0Adef 
@@ -1528,16 +1528,25 @@
 product(
+%0A        
 range(0,
@@ -1791,35 +1791,24 @@
 + tile_size,
- 0:n_bands%5D
 %0A           
@@ -1812,348 +1812,29 @@
     
-all_tiled_data.append((new_tile, (row, col), path_to_geotiff)) %0A%0A    %22%22%22%0A    for row in range(0, cols, tile_size):%0A        for col in range(0, rows, tile_size):%0A            in_bounds = row + tile_size %3C rows and col + tile_size %3C cols%0A            if in_bounds:%0A                new_tile = bands_data%5Brow:row + tile_size, col:col + tile_size,
+                     
  0:n
@@ -1841,20 +1841,16 @@
 _bands%5D%0A
-    
         
@@ -1919,17 +1919,8 @@
 ff))
- %0A    %22%22%22
 %0A%0A  
@@ -1944,16 +1944,17 @@
 d_data%0A%0A
+%0A
 def imag
@@ -2811,16 +2811,28 @@
 ture = %5B
+%0A           
  lat_lon
@@ -2884,16 +2884,24 @@
  feature
+%0A       
  %5D%0A     

remove buggy code