# -*- coding: utf-8 -*-
"""
Created on Wed Feb 21 18:05:01 2018

@author: z5030440
"""

#%% Initial settings

# import packages
import ee
from IPython import display
import math
import matplotlib.pyplot as plt
import numpy as np
from osgeo import gdal
import tempfile
import tensorflow as tf
import urllib
from urllib.request import urlretrieve
import zipfile
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 scripts
from GEEImageFunctions import *

np.seterr(all='ignore') # raise divisions by 0 and nans
ee.Initialize()

# Load image collection and filter it based on location (Narrabeen)

input_col = ee.ImageCollection('LANDSAT/LC08/C01/T1_RT_TOA')
#n_img = input_col.size().getInfo()
#print('Number of images in collection:', n_img)

# filter based on location (Narrabeen-Collaroy)
rect_narra = [[[151.3473129272461,-33.69035274454718],
              [151.2820816040039,-33.68206818063878], 
              [151.27281188964844,-33.74775138989556],
              [151.3425064086914,-33.75231878701767],
              [151.3473129272461,-33.69035274454718]]]; 
              
flt_col = input_col.filterBounds(ee.Geometry.Polygon(rect_narra))
n_img = flt_col.size().getInfo()
print('Number of images covering Narrabeen:', n_img)

# Select the most recent image and download it

im = ee.Image(flt_col.sort('SENSING_TIME',False).first())
im_dic = im.getInfo()
image_prop = im_dic.get('properties')
im_bands = im_dic.get('bands')
for i in range(len(im_bands)): del im_bands[i]['dimensions'] # delete the dimensions key

# download the panchromatic band (B8)
pan_band = [im_bands[7]]
im_pan = load_image(im, rect_narra, pan_band)
im_pan = im_pan[:,:,0]
size_pan = im_pan.shape
vec_pan = im_pan.reshape(size_pan[0] * size_pan[1])
# download the QA band (BQA)
qa_band = [im_bands[11]]
im_qa = load_image(im, rect_narra, qa_band)
im_qa = im_qa[:,:,0]

# convert QA bits
cloud_values = [2800, 2804, 2808, 2812, 6896, 6900, 6904, 6908]
cloud_shadow_values = [2976, 2980, 2984, 2988, 3008, 3012, 3016, 3020]

# Create cloud mask (resized to be applied to the Pan band)
im_cloud = np.isin(im_qa, cloud_values)
im_cloud_shadow = np.isin(im_qa, cloud_shadow_values)
im_cloud_res = transform.resize(im_cloud,(im_pan.shape[0], im_pan.shape[1]), order=0, preserve_range=True).astype('bool_')
vec_cloud = im_cloud.reshape(im_cloud.shape[0] * im_cloud.shape[1])
vec_cloud_res = im_cloud_res.reshape(size_pan[0] * size_pan[1])


# download the other bands (B2,B3,B4,B5,B6) = (blue,green,red,nir,swir1)
ms_bands = [im_bands[1], im_bands[2], im_bands[3], im_bands[4], im_bands[5]]
im_ms = load_image(im, rect_narra, ms_bands)
size_ms = im_ms.shape
vec_ms = im_ms.reshape(size_ms[0] * size_ms[1], size_ms[2])

# Plot the RGB image and cloud masks
plt.figure()
ax1 = plt.subplot(121)
plt.imshow(im_ms[:,:,[2,1,0]])
plt.title('RGB')
ax2 = plt.subplot(122)
plt.imshow(im_cloud, cmap='gray')
plt.title('Cloud mask')
#ax3 = plt.subplot(133, sharex=ax1, sharey=ax1)
#plt.imshow(im_cloud_shadow)
#plt.title('Cloud mask shadow')
plt.show()

# Resize multispectral bands (30m) to the size of the pan band (15m) using bilinear interpolation
im_ms_res = transform.resize(im_ms,(size_pan[0], size_pan[1]), order=1, preserve_range=True, mode='constant')
vec_ms_res = im_ms_res.reshape(size_pan[0] * size_pan[1], size_ms[2])

# Adjust intensities (set cloud pixels to 0 intensity)
cloud_value = np.nan

prc_low = 0 # lower percentile
prob_high = 99.9 # upper percentile probability to clip

# Rescale intensities between 0 and 1
vec_ms_adj = np.ones((len(vec_cloud_res),size_ms[2])) * np.nan
for i in range(im_ms.shape[2]):
    prc_high = np.percentile(vec_ms_res[~vec_cloud_res,i], prob_high)
    vec_rescaled = exposure.rescale_intensity(vec_ms_res[~vec_cloud_res,i], in_range=(prc_low,prc_high))
    plt.figure()
    plt.hist(vec_rescaled, bins = 300)
    plt.show()
    vec_ms_adj[~vec_cloud_res,i] = vec_rescaled

    
im_ms_adj = vec_ms_adj.reshape(size_pan[0], size_pan[1], size_ms[2])

# same for the pan band
vec_pan_adj = np.ones(len(vec_cloud_res)) * np.nan
prc_high = np.percentile(vec_pan[~vec_cloud_res],prob_high)
vec_rescaled = exposure.rescale_intensity(vec_pan[~vec_cloud_res], in_range=(prc_low,prc_high))
plt.figure()
plt.hist(vec_rescaled, bins = 300)
plt.show()
vec_pan_adj[~vec_cloud_res] = vec_rescaled
im_pan_adj = vec_pan_adj.reshape(size_pan[0], size_pan[1])

# Plot adjusted images
plt.figure()
plt.subplot(131)
plt.imshow(im_pan_adj, cmap='gray')
plt.title('PANCHROMATIC (15 m pixel)')

plt.subplot(132)
plt.imshow(im_ms_adj[:,:,[2,1,0]])
plt.title('RGB (30 m pixel)')
plt.show()

plt.subplot(133)
plt.imshow(im_ms_adj[:,:,[3,1,0]])
plt.title('NIR-GB (30 m pixel)')
plt.show()


#%% Pansharpening (PCA)
# Run PCA on selected bands

sel_bands = [0,1,2]
temp = vec_ms_adj[:,sel_bands]
vec_ms_adj_nocloud = temp[~vec_cloud_res,:]
pca = decomposition.PCA()
vec_pcs = pca.fit_transform(vec_ms_adj_nocloud)
vec_pcs_all = np.ones((len(vec_cloud_res),len(sel_bands))) * np.nan
vec_pcs_all[~vec_cloud_res,:] = vec_pcs
im_pcs = vec_pcs_all.reshape(size_pan[0], size_pan[1], vec_pcs.shape[1])

plt.figure()
plt.subplot(221)
plt.imshow(im_pcs[:,:,0], cmap='gray')
plt.title('Component 1')
plt.subplot(222)
plt.imshow(im_pcs[:,:,1], cmap='gray')
plt.title('Component 2')
plt.subplot(223)
plt.imshow(im_pcs[:,:,2], cmap='gray')
plt.title('Component 3')
plt.show()

# Compare the Pan image with the 1st Principal component
compare_images(im_pan_adj,im_pcs[:,:,0])
intensity_histogram(im_pan_adj)
intensity_histogram(im_pcs[:,:,0])

# Match histogram of the pan image with the 1st principal component and replace the 1st component
vec_pcs[:,0] = hist_match(vec_pan_adj[~vec_cloud_res], vec_pcs[:,0])
vec_ms_ps = pca.inverse_transform(vec_pcs)

# normalise between 0 and 1
for i in range(vec_pcs.shape[1]):
    vec_ms_ps[:,i] = np.divide(vec_ms_ps[:,i] - np.min(vec_ms_ps[:,i]),
               np.max(vec_ms_ps[:,i]) - np.min(vec_ms_ps[:,i]))

vec_ms_ps_all = np.ones((len(vec_cloud_res),len(sel_bands))) * np.nan
vec_ms_ps_all[~vec_cloud_res,:] = vec_ms_ps
im_ms_ps = vec_ms_ps_all.reshape(size_pan[0], size_pan[1], len(sel_bands))
vec_ms_ps_all = np.append(vec_ms_ps_all, vec_ms_adj[:,[3,4]], axis=1)
im_ms_ps = np.append(im_ms_ps, im_ms_adj[:,:,[3,4]], axis=2)

# Plot adjusted images
plt.figure()
plt.subplot(121)
plt.imshow(im_ms_adj[:,:,[2,1,0]])
plt.title('Original RGB')
plt.show()

plt.subplot(122)
plt.imshow(im_ms_ps[:,:,[2,1,0]])
plt.title('Pansharpened RGB')
plt.show()

plt.figure()
plt.subplot(121)
plt.imshow(im_ms_adj[:,:,[3,1,0]])
plt.title('Original NIR-GB')
plt.show()

plt.subplot(122)
plt.imshow(im_ms_ps[:,:,[3,1,0]])
plt.title('Pansharpened NIR-GB')
plt.show()
#%% Compute Normalized Difference Water Index (NDWI)

# With NIR
vec_ndwi_nir = np.ones(len(vec_cloud_res)) * np.nan
temp = np.divide(vec_ms_ps_all[~vec_cloud_res,3] - vec_ms_ps_all[~vec_cloud_res,1],
                         vec_ms_ps_all[~vec_cloud_res,3] + vec_ms_ps_all[~vec_cloud_res,1])
vec_ndwi_nir[~vec_cloud_res] = temp
im_ndwi_nir = vec_ndwi_nir.reshape(size_pan[0], size_pan[1])

# With SWIR_1
vec_ndwi_swir = np.ones(len(vec_cloud_res)) * np.nan
temp = np.divide(vec_ms_ps_all[~vec_cloud_res,4] - vec_ms_ps_all[~vec_cloud_res,1],
                         vec_ms_ps_all[~vec_cloud_res,4] + vec_ms_ps_all[~vec_cloud_res,1])
vec_ndwi_swir[~vec_cloud_res] = temp
im_ndwi_swir = vec_ndwi_swir.reshape(size_pan[0], size_pan[1])

plt.figure()
ax1 = plt.subplot(211)
plt.hist(vec_ndwi_nir[~vec_cloud_res], bins=300, label='NIR')
plt.hist(vec_ndwi_swir[~vec_cloud_res], bins=300, label='SWIR', alpha=0.5)
plt.legend()
ax2 = plt.subplot(212, sharex=ax1)
plt.hist(vec_ndwi_nir[~vec_cloud_res], bins=300, cumulative=True, histtype='step', label='NIR')
plt.hist(vec_ndwi_swir[~vec_cloud_res], bins=300, cumulative=True, histtype='step', label='SWIR')
plt.legend()
plt.show()

compare_images(im_ndwi_nir,im_ndwi_swir)

plt.figure()
plt.imshow(im_ndwi_nir, cmap='seismic')
plt.title('Water Index')
plt.colorbar()
plt.show()

#%% Extract shorelines (NIR)

ndwi_nir = vec_ndwi_nir[~vec_cloud_res]

t_otsu = filters.threshold_otsu(ndwi_nir)
t_min = filters.threshold_minimum(ndwi_nir)
t_mean = filters.threshold_mean(ndwi_nir)
t_li = filters.threshold_li(ndwi_nir)
# try all thresholding algorithms

plt.figure()
plt.hist(ndwi_nir, bins=300)
plt.plot([t_otsu, t_otsu],[0, 15000], 'r-', label='Otsu threshold')
#plt.plot([t_min, t_min],[0, 15000], 'g-', label='min')
#plt.plot([t_mean, t_mean],[0, 15000], 'y-', label='mean')
#plt.plot([t_li, t_li],[0, 15000], 'm-', label='li')
plt.legend()
plt.show()

plt.figure()
plt.imshow(im_ndwi_nir > t_otsu, cmap='gray')
plt.title('Binary image')
plt.show()

im_bin = im_ndwi_nir > t_otsu
im_open = morphology.binary_opening(im_bin,morphology.disk(1))
im_close = morphology.binary_closing(im_open,morphology.disk(1))
im_bin_coast_in = im_close ^ morphology.erosion(im_close,morphology.disk(1))
im_bin_sl_in = morphology.remove_small_objects(im_bin_coast_in,100,8)

compare_images(im_close,im_bin_sl_in)

plt.figure()
plt.subplot(121)
plt.imshow(im_close, cmap='gray')
plt.title('morphological closing')
plt.subplot(122)
plt.imshow(im_bin_sl_in, cmap='gray')
plt.title('Water mark')
plt.show()


im_bin_coast_out = morphology.dilation(im_close,morphology.disk(1)) ^ im_close
im_bin_sl_out = morphology.remove_small_objects(im_bin_coast_out,100,8)


# Plot shorelines on top of RGB image
im_rgb_sl = np.copy(im_ms_ps[:,:,[2,1,0]])

im_rgb_sl[im_bin_sl_in,0] = 0
im_rgb_sl[im_bin_sl_in,1] = 1
im_rgb_sl[im_bin_sl_in,2] = 1

im_rgb_sl[im_bin_sl_out,0] = 1
im_rgb_sl[im_bin_sl_out,1] = 0
im_rgb_sl[im_bin_sl_out,2] = 1

plt.figure()
plt.imshow(im_rgb_sl)
plt.title('Pansharpened RGB')
plt.show()

#%% Extract shorelines SWIR

ndwi_swir = vec_ndwi_swir[~vec_cloud_res]
t_otsu = filters.threshold_otsu(ndwi_swir)

plt.figure()
plt.hist(ndwi_swir, bins=300)
plt.plot([t_otsu, t_otsu],[0, 15000], 'r-', label='Otsu threshold')
#plt.plot([t_min, t_min],[0, 15000], 'g-', label='min')
#plt.plot([t_mean, t_mean],[0, 15000], 'y-', label='mean')
#plt.plot([t_li, t_li],[0, 15000], 'm-', label='li')
plt.legend()
plt.show()

plt.figure()
plt.imshow(im_ndwi_swir > t_otsu, cmap='gray')
plt.title('Binary image')
plt.show()

im_bin = im_ndwi_swir > t_otsu
im_open = morphology.binary_opening(im_bin,morphology.disk(1))
im_close = morphology.binary_closing(im_open,morphology.disk(1))
im_bin_coast_in = im_close ^ morphology.erosion(im_close,morphology.disk(1))
im_bin_sl_in = morphology.remove_small_objects(im_bin_coast_in,100,8)

compare_images(im_close,im_bin_sl_in)


im_bin_coast_out = morphology.dilation(im_close,morphology.disk(1)) ^ im_close
im_bin_sl_out = morphology.remove_small_objects(im_bin_coast_out,100,8)


# Plot shorelines on top of RGB image
im_rgb_sl = np.copy(im_ms_ps[:,:,[2,1,0]])

im_rgb_sl[im_bin_sl_in,0] = 0
im_rgb_sl[im_bin_sl_in,1] = 1
im_rgb_sl[im_bin_sl_in,2] = 1

im_rgb_sl[im_bin_sl_out,0] = 1
im_rgb_sl[im_bin_sl_out,1] = 0
im_rgb_sl[im_bin_sl_out,2] = 1

plt.figure()
plt.imshow(im_rgb_sl)
plt.title('Pansharpened RGB')
plt.show()