import datacube
import utils.data_cube_utilities.data_access_api as dc_api  

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

api = dc_api.DataAccessApi()
dc = datacube.Datacube(app = 'ardc_task_b')
api.dc = dc

"""
Explanation: ARDC Training: Python Notebooks
Task-B: Water Extent (WOFS) and Water Quality (TSM)

Import the Datacube Configuration
End of explanation
"""


list_of_products = dc.list_products()
netCDF_products = list_of_products[list_of_products['format'] == 'NetCDF']
netCDF_products

"""
Explanation: Browse the available Data Cubes
End of explanation
"""


# Change the data platform and data cube here

platform = 'LANDSAT_7'
product = 'ls7_usgs_sr_scene'

"""
Explanation: Pick a product
Use the platform and product names from the previous block to select a Data Cube.
End of explanation
"""


from utils.data_cube_utilities.dc_time import _n64_to_datetime, dt_to_str

extents = api.get_full_dataset_extent(platform = platform, product = product, measurements=[])

latitude_extents = (min(extents['latitude'].values),max(extents['latitude'].values))
longitude_extents = (min(extents['longitude'].values),max(extents['longitude'].values))
time_extents = (min(extents['time'].values),max(extents['time'].values))

print("Latitude Extents:", latitude_extents)
print("Longitude Extents:", longitude_extents)
print("Time Extents:", list(map(dt_to_str, map(_n64_to_datetime, time_extents))))

"""
Explanation: Display Latitude-Longitude and Time Bounds of the Data Cube
End of explanation
"""


## The code below renders a map that can be used to orient yourself with the region.
from utils.data_cube_utilities.dc_display_map import display_map
display_map(latitude = latitude_extents, longitude = longitude_extents)

"""
Explanation: Visualize Data Cube Region
End of explanation
"""


## Vietnam - Central Lam Dong Province ##
longitude_extents = (107.0, 107.2)
latitude_extents  = (11.7, 12.0)

## Tanzania - Lake Sulunga
longitude_extents = (35.00, 35.37)
latitude_extents = (-6.28,-5.87)

time_extents = ('2015-01-01', '2015-12-31')

display_map(latitude = latitude_extents, longitude = longitude_extents)

"""
Explanation: Pick a smaller analysis region and display that region
Try to keep your region to less than 0.2-deg x 0.2-deg for rapid processing. You can click on the map above to find the Lat-Lon coordinates of any location. You will want to identify a region with an inland water body. Pick a time window of a few months so we can pick out some clear pixels and plot the water.
End of explanation
"""


landsat_dataset = dc.load(latitude = latitude_extents,
                          longitude = longitude_extents,
                          platform = platform,
                          time = time_extents,
                          product = product,
                          measurements = ['red', 'green', 'blue', 'nir', 'swir1', 'swir2', 'pixel_qa']) 

landsat_dataset
#view the dimensions and sample content from the cube

"""
Explanation: Load the dataset and the required spectral bands or other parameters
After loading, you will view the Xarray dataset. Notice the dimensions represent the number of pixels in your latitude and longitude dimension as well as the number of time slices (time) in your time series.
End of explanation
"""


acquisition_number = 2
# select an acquisition number from 1 to "time" using the array limits above

%matplotlib inline
#landsat_dataset.red.isel(time = acquisition_number).plot(cmap = "Greys")
landsat_dataset.green.isel(time = acquisition_number).plot(cmap = "Greys")
#landsat_dataset.blue.isel(time = acquisition_number).plot(cmap = "Greys")
#landsat_dataset.nir.isel(time = acquisition_number).plot(cmap = "Greys")
#landsat_dataset.swir1.isel(time = acquisition_number).plot(cmap = "Greys")
#landsat_dataset.swir2.isel(time = acquisition_number).plot(cmap = "Greys")

"""
Explanation: Display Example Images

Single band visualization
For a quick inspection, let's look at one image. The code will allow the selection of any band (red, blue, green, nir, swir1, swir2) to produce a grey-scale image. Select the desired acquisition (time slice) in the block below. You can select from 1 to #, where the max value is the number of time slices noted in the block above. Change the comment statements below to select the bands for the first image.
End of explanation
"""


import numpy as np  

def generate_cloud_mask(dataset, include_shadows = False):
    #Create boolean Masks for clear and water pixels
    clear_pixels = dataset.pixel_qa.values == 2  + 64
    water_pixels = dataset.pixel_qa.values == 4  + 64
    shadow_pixels= dataset.pixel_qa.values == 8 + 64
    
    a_clean_mask = np.logical_or(clear_pixels, water_pixels)
    
    if include_shadows:
        a_clean_mask = np.logical_or(a_clean_mask, shadow_pixels)
        
    return np.invert(a_clean_mask)

def remove_clouds(dataset, include_shadows = False):
    #Create boolean Masks for clear and water pixels
    clear_pixels = dataset.pixel_qa.values == 2  + 64
    water_pixels = dataset.pixel_qa.values == 4  + 64
    shadow_pixels= dataset.pixel_qa.values == 8 + 64
    
    a_clean_mask = np.logical_or(clear_pixels, water_pixels)
    
    if include_shadows:
        a_clean_mask = np.logical_or(a_clean_mask, shadow_pixels)
        
    return dataset.where(a_clean_mask)

cloud_mask = generate_cloud_mask(landsat_dataset)
cloudless = remove_clouds(landsat_dataset) #landsat_dataset.where(image_is_clean)

"""
Explanation: Define Cloud Masking Function
Removes clouds and cloud shadows based on the Landsat pixel QA information
This is only for reference ... nothing to modify here
End of explanation
"""


from utils.data_cube_utilities.dc_rgb import rgb

"""
Explanation: Set up plotting function (to be used later)
Nothing to modify here
End of explanation
"""


from utils.data_cube_utilities.dc_mosaic import create_mosaic

def mrf_mosaic(dataset):
    # The mask here is based on pixel_qa products. It comes bundled in with most Landsat Products.
    cloud_free_boolean_mask = np.invert(generate_cloud_mask(dataset))
    return create_mosaic(dataset, clean_mask = cloud_free_boolean_mask)

recent_composite = mrf_mosaic(landsat_dataset)

recent_composite.nir.plot(cmap = "Greys")

rgb(recent_composite, width = 20)

"""
Explanation: Most Recent Pixel Mosaic  
Masks clouds from imagery and uses the most recent cloud-free pixels.
End of explanation
"""


from utils.data_cube_utilities.dc_water_classifier import wofs_classify

water_classification = wofs_classify(recent_composite, clean_mask = np.ones(recent_composite.pixel_qa.shape).astype(np.bool),  mosaic = True) 

water_classification.wofs.plot(cmap='Blues')

"""
Explanation: Plot WOFS water detection results
This example uses the Australian Water Detection from Space (WOFS) algorithm for water detection. The base image will use a most-recent pixel composite (from above). When reviewing the results, 1=water, 0=no water.
End of explanation
"""


def NDWI(dataset):
    return (dataset.green - dataset.nir)/(dataset.green + dataset.nir)

ndwi = NDWI(recent_composite)  # High Concentrations of Water - Blues  

(ndwi).plot(cmap = "Blues")

"""
Explanation: Plot NDWI water detection results
This example uses the Normalized Difference Water Index (NDWI) which is a spectral "index" that correlates well with the existance of water. 
<br>
$$ NDWI =  \frac{GREEN - NIR}{GREEN + NIR}$$
End of explanation
"""


from utils.data_cube_utilities.dc_water_quality import tsm

mask_that_only_includes_water_pixels = water_classification.wofs == 1  
tsm_dataset = tsm(recent_composite, clean_mask = mask_that_only_includes_water_pixels )

tsm_dataset.tsm.plot(cmap = "jet")

"""
Explanation: Plot TSM water quality results
This example uses the Australian Total Suspended Matter (TSM) algorithm. The TSM value is the mean over the entire time range. This parameter is a measure of the particulate matter in water and is often a proxy for water quality.
End of explanation
"""


ts_water_classification = wofs_classify(landsat_dataset,clean_mask = np.invert(cloud_mask)) 

# Apply nan to no_data values
ts_water_classification = ts_water_classification.where(ts_water_classification != -9999)

##Time series aggregation that ignores nan values.    
water_classification_percentages = (ts_water_classification.mean(dim = ['time']) * 100).wofs.rename('water_classification_percentages')

## import color-scheme and set nans to black
from matplotlib.cm import jet_r as jet_r
jet_r.set_bad('black',1)

## apply nan to percentage values that aren't greater than 0, then plot
water_classification_percentages\
    .where(water_classification_percentages > 0)\
    .plot(cmap = jet_r)

"""
Explanation: Time Series Water Detection Analysis
Time series output of the Australian Water Detection from Space (WOFS) results. The results show the percent of time that a pixel is classified as water over the entire time series. BLUE = frequent water, RED = infrequent water.
End of explanation
"""


pixel_lat = 11.84
pixel_lon = 107.09

pixel = ts_water_classification.sel( latitude  = pixel_lat,
                                     longitude = pixel_lon,
                                     method = 'nearest') # nearest neighbor selection  

import matplotlib.pyplot as plt 
plt.figure(figsize = (20,5)) 
plt.scatter(pixel.time.values, pixel.wofs.values)

"""
Explanation: Create a WOFS plot for a single pixel
First select the Lat-Lon position. Then the code will find the closest pixel in the dataset using a "nearest neighbor" selection.
End of explanation
"""


