import sys
import os
sys.path.append(os.environ.get('NOTEBOOK_ROOT'))

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

def bathymetry_index(df, m0 = 1, m1 = 0):
    return m0*(np.log(df.blue)/np.log(df.green))+m1

"""
Explanation: <a id="Shallow_Water_Bathymetry_top"></a>
Shallow Water Bathymetry
Visualizing Differences in Depth With Spectral Analysis
<hr>

Notebook Summary

Import data from LANDSAT 8
A bathymetry index is calculated
Contrast is adjusted to make a more interpretable visualization.
Citation: Stumpf, Richard P., Kristine Holderied, and Mark Sinclair. "Determination of water depth with high‐resolution satellite imagery over variable bottom types." Limnology and Oceanography 48.1part2 (2003): 547-556.



<hr>

Index

Import Dependencies and Connect to the Data Cube
Choose the Platform and Product
Define the Extents of the Analysis
Retrieve the Data
Calculate the Bathymetry and NDWI Indices
Export Unmasked GeoTIFF
Mask the Dataset Using the Quality Column and NDWI
Create a Visualization Function
Visualize the Bathymetry
Visualize the Bathymetry With Adjusted Contrast

<hr>

How It Works
Bathymetry is the measurement of depth in bodies of water(Oceans, Seas or Lakes).  This notebook illustrates a technique for deriving depth of shallow water areas using purely optical features from Landsat Collection 1 imagery and draws heavily from the publication Determination of water depth with high-resolution satelite imagery over variable bottom types.
<br>
Bathymetry Index  
This bathymetry index uses optical green and blue values on a logarithmic scale with two tunable coefficients m0 and m1.
$$ BATH =  m_0*\frac{ln(blue)}{ln(green)} -m_1$$  
Where: 
- m0 is a tunable scaling factor to tune the ratio to depth <br>
- m1 is the offset for a depth of 0 meters.
<br>
<div class="alert-info"><br>
<b>Note: </b> that for our purposes, $m_0$ and $m_1$ are equal to <b>1</b> and <b>0</b> respectively, since we cannot determine the baseline nor the offset from spectral reflectance alone. This effectively simplifies the formula to: $$\frac{ln(blue)}{ln(green)}$$
<br>

</div>

Bathymetry Index Function
End of explanation
"""


from datacube.utils.aws import configure_s3_access
configure_s3_access(requester_pays=True)

import datacube
dc = datacube.Datacube()

"""
Explanation: <span id="Shallow_Water_Bathymetry_import">Import Dependencies and Connect to the Data Cube &#9652;</span>
End of explanation
"""


#List the products available on this server/device
dc.list_products()

#create a list of the desired platforms
platform = 'LANDSAT_8'
product = 'ls8_level1_usgs'

"""
Explanation: <span id="Shallow_Water_Bathymetry_plat_prod">Choose the Platform and Product &#9652;</span>
End of explanation
"""


# East Coast of Australia
lat_subsect = (-31.7, -32.2)
lon_subsect = (152.4, 152.9)

print('''
Latitude:\t{0}\t\tRange:\t{2} degrees
Longitude:\t{1}\t\tRange:\t{3} degrees
'''.format(lat_subsect,
           lon_subsect,
           max(lat_subsect)-min(lat_subsect),
           max(lon_subsect)-min(lon_subsect)))

"""
Explanation: <span id="Shallow_Water_Bathymetry_define_extents">Define the Extents of the Analysis &#9652;</span>
Region bounds
End of explanation
"""


from utils.data_cube_utilities.dc_display_map import display_map      
display_map(latitude = lat_subsect,longitude = lon_subsect)

"""
Explanation: Display
End of explanation
"""


%%time
ds = dc.load(lat = lat_subsect,
             lon = lon_subsect,
             platform = platform,
             product = product,
             output_crs = "EPSG:32756",
             measurements = ["red","blue","green","nir","quality"],
             resolution = (-30,30))

ds

"""
Explanation: <span id="Shallow_Water_Bathymetry_retrieve_data">Retrieve the Data &#9652;</span>
Load and integrate datasets
End of explanation
"""


from utils.data_cube_utilities.dc_rgb import rgb
rgb(ds.isel(time=6), x_coord='x', y_coord='y')
plt.show()

"""
Explanation: Preview the Data
End of explanation
"""


# Create Bathemtry Index column
ds["bathymetry"] = bathymetry_index(ds)

from utils.data_cube_utilities.dc_water_classifier import NDWI
# (green - nir) / (green + nir)
ds["ndwi"] = NDWI(ds, band_pair=1)

"""
Explanation: <span id="Shallow_Water_Bathymetry_bathymetry">Calculate the Bathymetry and NDWI Indices &#9652;</span>


The bathymetry function is located at the top of this notebook.
End of explanation
"""


ds

"""
Explanation: <hr>

Preview Combined Dataset
End of explanation
"""


import os
from utils.data_cube_utilities.import_export import export_xarray_to_multiple_geotiffs

unmasked_dir = "geotiffs/landsat8/unmasked"
if not os.path.exists(unmasked_dir):
    os.makedirs(unmasked_dir)
export_xarray_to_multiple_geotiffs(ds, unmasked_dir + "/unmasked.tif",
                                   x_coord='x', y_coord='y')

"""
Explanation: <span id="Shallow_Water_Bathymetry_export_unmasked">Export Unmasked GeoTIFF &#9652;</span>
End of explanation
"""


# preview values
np.unique(ds["quality"])

"""
Explanation: <span id="Shallow_Water_Bathymetry_mask">Mask the Dataset Using the Quality Column and NDWI &#9652;</span>
End of explanation
"""


# Tunable threshold for masking the land out
threshold = .05

water = (ds.ndwi>threshold).values

#preview one time slice to determine the effectiveness of the NDWI masking
rgb(ds.where(water).isel(time=6), x_coord='x', y_coord='y')
plt.show()

from utils.data_cube_utilities.dc_mosaic import ls8_oli_unpack_qa
clear_xarray  = ls8_oli_unpack_qa(ds.quality, "clear")

full_mask = np.logical_and(clear_xarray, water)

ds = ds.where(full_mask)

"""
Explanation: Use NDWI to Mask Out Land

The threshold can be tuned if need be to better fit the RGB image above. <br>
Unfortunately our existing WOFS algorithm is designed to work with Surface Reflectance (SR) and does not work with this data yet but with a few modifications it could be made to do so.  We will approximate the WOFs mask with NDWI for now.
End of explanation
"""


plt.figure(figsize=[15,5])

#Visualize the distribution of the remaining data
sns.boxplot(ds['bathymetry'])
plt.show()

"""
Explanation: <span id="Shallow_Water_Bathymetry_vis_func">Create a Visualization Function &#9652;</span>
Visualize the distribution of the bathymetry index for the water pixels
End of explanation
"""


#set the quantile range in either direction from the median value
def get_quantile_range(col, quantile_range = .25):
    low = ds[col].quantile(.5 - quantile_range,["time","y","x"]).values
    high = ds[col].quantile(.5 + quantile_range,["time","y","x"]).values
    return low,high

#Custom function for a color mapping object
from matplotlib.colors import LinearSegmentedColormap

def custom_color_mapper(name = "custom", val_range = (1.96,1.96), colors = "RdGnBu"):
    custom_cmap = LinearSegmentedColormap.from_list(name,colors=colors)
    
    min, max = val_range
    step = max/10.0
    Z = [min,0],[0,max]
    levels = np.arange(min,max+step,step)
    cust_map = plt.contourf(Z, 100, cmap=custom_cmap)
    plt.clf()
    return cust_map.cmap

def mean_value_visual(ds, col, figsize = [15,15], cmap = "GnBu", low=None, high=None):
    if low is None: low = np.min(ds[col]).values
    if high is None: high = np.max(ds[col]).values
    ds.reduce(np.nanmean,dim=["time"])[col].plot.imshow(figsize = figsize, cmap=cmap,  
                                                        vmin=low, vmax=high)

"""
Explanation: <b>Interpretation: </b> We can see that most of the values fall within a very short range.  We can scale our plot's cmap limits to fit the specific quantile ranges for the bathymetry index so we can achieve better contrast from our plots.
End of explanation
"""


mean_value_visual(ds, "bathymetry", cmap="GnBu")

"""
Explanation: <span id="Shallow_Water_Bathymetry_bath_vis">Visualize the Bathymetry &#9652;</span>
End of explanation
"""


# create range using the 10th and 90th quantile
low, high = get_quantile_range("bathymetry", .40)


custom = custom_color_mapper(val_range=(low,high),
                             colors=["darkred","red","orange","yellow","green",
                                     "blue","darkblue","black"])

mean_value_visual(ds, "bathymetry", cmap=custom, low=low, high=high)

# create range using the 5th and 95th quantile
low, high = get_quantile_range("bathymetry", .45)


custom = custom_color_mapper(val_range=(low,high),
                             colors=["darkred","red","orange","yellow","green",
                                     "blue","darkblue","black"])

mean_value_visual(ds, "bathymetry", cmap = custom, low=low, high = high)

# create range using the 2nd and 98th quantile
low, high = get_quantile_range("bathymetry", .48)


custom = custom_color_mapper(val_range=(low,high),
                             colors=["darkred","red","orange","yellow","green",
                                     "blue","darkblue","black"])

mean_value_visual(ds, "bathymetry", cmap=custom, low=low, high=high)

# create range using the 1st and 99th quantile
low, high = get_quantile_range("bathymetry", .49)


custom = custom_color_mapper(val_range=(low,high),
                             colors=["darkred","red","orange","yellow","green",
                                     "blue","darkblue","black"])

mean_value_visual(ds, "bathymetry", cmap=custom, low=low, high=high)

"""
Explanation: <span id="Shallow_Water_Bathymetry_bath_vis_better">Visualize the Bathymetry With Adjusted Contrast &#9652;</span>

If we clamp the range of the plot using different quantile ranges we can see relative differences in higher contrast.
End of explanation
"""


