# -*- coding: utf-8 -*-
"""
Created on Thu Mar  1 14:32:08 2018

@author: z5030440

Main code to extract shorelines from Landsat imagery
"""
# Preamble

import ee
from IPython import display
import math
import matplotlib.pyplot as plt
import numpy as np
import pdb

# image processing modules
import skimage.filters as filters 
import skimage.exposure as exposure
import skimage.transform as transform
import sklearn.decomposition as decomposition
import skimage.morphology as morphology
import skimage.measure as measure

from shapely.geometry import Polygon

from osgeo import gdal
from osgeo import osr
import tempfile
import urllib
from urllib.request import urlretrieve
import zipfile

# my modules
from utils import *
# from sds import *

np.seterr(all='ignore') # raise/ignore divisions by 0 and nans
ee.Initialize()
plot_bool = True # if you want the plots

def download_tif(image, bandsId):
    """downloads tif image (region and bands) from the ee server and stores it in a temp file"""
    url = ee.data.makeDownloadUrl(ee.data.getDownloadId({
        'image': image.serialize(),
        'bands': bandsId,
        'filePerBand': 'false',
        'name': 'data',
        }))
    local_zip, headers = urlretrieve(url)
    with zipfile.ZipFile(local_zip) as local_zipfile:
        return local_zipfile.extract('data.tif', tempfile.mkdtemp())

def load_image(image, bandsId): 
    """loads an ee.Image() as a np.array. e.Image() is retrieved from the EE database."""   
    local_tif_filename = download_tif(image, bandsId)
    dataset = gdal.Open(local_tif_filename, gdal.GA_ReadOnly)
    bands = [dataset.GetRasterBand(i + 1).ReadAsArray() for i in range(dataset.RasterCount)]
    return np.stack(bands, 2), dataset



im = ee.Image('LANDSAT/LC08/C01/T1_RT_TOA/LC08_089083_20130411')

lon = [151.2820816040039, 151.3425064086914]
lat = [-33.68206818063878, -33.74775138989556]
polygon = [[lon[0], lat[0]], [lon[1], lat[0]], [lon[1], lat[1]], [lon[0], lat[1]]];
             
# get image metadata into dictionnary
im_dic = im.getInfo()
im_bands = im_dic.get('bands')
# delete dimensions key from dictionnary, otherwise the entire image is extracted
#for i in range(len(im_bands)): del im_bands[i]['dimensions']
pan_band = [im_bands[7]]
ms_bands = [im_bands[1], im_bands[2], im_bands[3]]
im_full, dataset_full = load_image(im, ms_bands)

plt.figure()
plt.imshow(np.clip(im_full[:,:,[2,1,0]] * 3, 0, 1))
plt.show()
#%%

def download_tif(image, polygon, bandsId):
    """downloads tif image (region and bands) from the ee server and stores it in a temp file"""
    url = ee.data.makeDownloadUrl(ee.data.getDownloadId({
        'image': image.serialize(),
        'region': polygon,
        'bands': bandsId,
        'filePerBand': 'false',
        'name': 'data',
        }))
    local_zip, headers = urlretrieve(url)
    with zipfile.ZipFile(local_zip) as local_zipfile:
        return local_zipfile.extract('data.tif', tempfile.mkdtemp())

def load_image(image, polygon, bandsId): 
    """
    Loads an ee.Image() as a np.array. e.Image() is retrieved from the EE database.
    The geographic area and bands to select can be specified
    
    KV WRL 2018
    
    Arguments:
    -----------
        image: ee.Image()
            image objec from the EE database
        polygon: list
            coordinates of the points creating a polygon. Each point is a list with 2 values
        bandsId: list
            bands to select, each band is a dictionnary in the list containing the following keys:
            crs, crs_transform, data_type and id. NOTE: you have to remove the key dimensions, otherwise
            the entire image is retrieved.
            
    Returns:
    -----------
        image_array : np.ndarray
            An array containing the image (2D if one band, otherwise 3D)
    """
    
    local_tif_filename = download_tif(image, polygon, bandsId)
    dataset = gdal.Open(local_tif_filename, gdal.GA_ReadOnly)
    bands = [dataset.GetRasterBand(i + 1).ReadAsArray() for i in range(dataset.RasterCount)]
    return np.stack(bands, 2), dataset


for i in range(len(im_bands)): del im_bands[i]['dimensions']
ms_bands = [im_bands[1], im_bands[2], im_bands[3]]

im_cropped, dataset_cropped = load_image(im, polygon, ms_bands)

plt.figure()
plt.imshow(np.clip(im_cropped[:,:,[2,1,0]] * 3, 0, 1))
plt.show()

#%%
crs_full  = dataset_full.GetGeoTransform()
crs_cropped  = dataset_cropped.GetGeoTransform()
scale = crs_full[1]
ul_full = np.array([crs_full[0], crs_full[3]])
ul_cropped = np.array([crs_cropped[0], crs_cropped[3]])

delta = np.abs(ul_full - ul_cropped)/scale

u0 = delta[0].astype('int')
v0 = delta[1].astype('int')

im_full[v0,u0,:]
im_cropped[0,0,:]

lrx = ul_cropped[0] + (dataset_cropped.RasterXSize * scale)
lry = ul_cropped[1] + (dataset_cropped.RasterYSize * (-scale))

lr_cropped = np.array([lrx, lry])

delta = np.abs(ul_full - lr_cropped)/scale
u1 = delta[0].astype('int')
v1 = delta[1].astype('int')

im_cropped2 = im_full[v0:v1,u0:u1,:]

#%%
crs_full  = dataset_full.GetGeoTransform()
source = osr.SpatialReference()
source.ImportFromWkt(dataset_full.GetProjection())

target = osr.SpatialReference()
target.ImportFromEPSG(4326)

transform = osr.CoordinateTransformation(source, target)

transform.TransformPoint(ulx, uly)

#%%
crs_cropped  = dataset_cropped.GetGeoTransform()
ulx = crs_cropped[0]
uly = crs_cropped[3]
source = osr.SpatialReference()
source.ImportFromWkt(dataset_cropped.GetProjection())

target = osr.SpatialReference()
target.ImportFromEPSG(4326)

transform = osr.CoordinateTransformation(source, target)

transform.TransformPoint(lrx, lry)



#%%
source = osr.SpatialReference()
source.ImportFromEPSG(4326)

target = osr.SpatialReference()
target.ImportFromEPSG(32656)

coords = transform.TransformPoint(151.2820816040039, -33.68206818063878)
coords[0] - ulx
coords[1] - uly
#%%
x_ul_full = ms_bands[0]['crs_transform'][2]
y_ul_full = ms_bands[0]['crs_transform'][5]
scale = ms_bands[0]['crs_transform'][0]

x_ul_cropped = np.array([340756.105840223, 346357.851288875, 346474.839525944, 340877.362938763])
y_ul_cropped = np.array([-3728229.45372866, -3728137.91775723, -3735421.58347927,  -3735513.20696522])

dx = abs(x_ul_full - x_ul_cropped)
dy = abs(y_ul_full - y_ul_cropped)

u_coord = np.round(dx/scale).astype('int')
v_coord = np.round(dy/scale).astype('int')

im_cropped2 = im_full[np.min(v_coord):np.max(v_coord), np.min(u_coord):np.max(u_coord),:]

plt.figure()
plt.imshow(np.clip(im_cropped2[:,:,[2,1,0]] * 3, 0, 1), cmap='gray')
plt.show()

sum(sum(sum(np.equal(im_cropped,im_cropped2).astype('int')-1)))

