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

import xarray as xr  
import numpy as np
import datacube
from utils.data_cube_utilities.data_access_api import DataAccessApi

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

api = DataAccessApi()
dc = api.dc

"""
Explanation: <a id="top"></a>
Export Notebook
<hr>

Notebook Summary
The code in this notebook subsets a data cube, selects a specific set of variables, generates some additional data from those and then outputs that data into a GeoTIFF file. The goal is to be able to do external analyses of this data using other data analysis tools or GIS tools. The files would be reasonable in size, since we would restrict the region and parameters in the output.
<hr>

Index

Import Dependencies and Connect to the Data Cube
Choose Platforms and Products
Get the Extents of the Cube
Define the Extents of the Analysis
Load Data from the Datacube
Derive Products
Combine Data
Export Data
Export to GeoTIFF
Export to NetCDF



<hr>

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


list_of_products = dc.list_products()
list_of_products

"""
Explanation: <span id="plat_prod">Choose Platforms and Products &#9652;</span>
List available products for each platform
End of explanation
"""


platform = 'LANDSAT_7'
product = 'ls7_usgs_sr_scene'
collection = 'c1'
level = 'l2'

"""
Explanation: Choose product
End of explanation
"""


from utils.data_cube_utilities.dc_load import get_product_extents
from utils.data_cube_utilities.dc_time import dt_to_str

full_lat, full_lon, min_max_dates = get_product_extents(api, platform, product)

# Print the extents of the combined data.
print("Latitude Extents:", full_lat)
print("Longitude Extents:", full_lon)
print("Time Extents:", list(map(dt_to_str, min_max_dates)))

## 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(full_lat, full_lon)

"""
Explanation: <span id="extents">Get the Extents of the Cube &#9652;</span>
End of explanation
"""


######### Ghana - Pambros Salt Ponds ################## 
lon = (-0.3013, -0.2671)
lat = (5.5155, 5.5617)

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

from utils.data_cube_utilities.dc_display_map import display_map
display_map(lat, lon)

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


landsat_dataset = dc.load(latitude = lat,
                          longitude = lon,
                          platform = platform,
                          time = time_extents,
                          product = product,
                          measurements = ['red', 'green', 'blue', 'nir', 'swir1', 'swir2', 'pixel_qa'],
                          dask_chunks={'time':1, 'latitude':1000, 'longitude':1000})

"""
Explanation: <span id="retrieve_data">Load Data from the Data Cube &#9652;</span>
End of explanation
"""


from utils.data_cube_utilities.clean_mask import landsat_qa_clean_mask

plt_col_lvl_params = dict(platform=platform, collection=collection, level=level)
clear_xarray = landsat_qa_clean_mask(landsat_dataset, cover_types=['clear'], **plt_col_lvl_params)
water_xarray = landsat_qa_clean_mask(landsat_dataset, cover_types=['water'], **plt_col_lvl_params)
shadow_xarray = landsat_qa_clean_mask(landsat_dataset, cover_types=['cld_shd'], **plt_col_lvl_params)

# clean_xarray = xr.ufuncs.logical_or(clear_xarray, water_xarray).rename("clean_mask")
clean_xarray = (clear_xarray | water_xarray).rename("clean_mask")

# landsat_qa_clean_mask(cover_types=[])

"""
Explanation: <span id="derive_products">Derive Products &#9652;</span>

Masks
End of explanation
"""


from utils.data_cube_utilities.dc_water_classifier import wofs_classify

water_classification = wofs_classify(landsat_dataset,
                                     clean_mask = clean_xarray, 
                                     mosaic = False) 

wofs_xarray = water_classification.wofs

"""
Explanation: Water Classification
End of explanation
"""


def NDVI(dataset):
    return ((dataset.nir - dataset.red)/(dataset.nir + dataset.red)).rename("NDVI")

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

def NDBI(dataset):
        return ((dataset.swir2 - dataset.nir)/(dataset.swir2 + dataset.nir)).rename("NDBI")

ndbi_xarray = NDBI(landsat_dataset)  # Urbanization - Reds
ndvi_xarray = NDVI(landsat_dataset)  # Dense Vegetation - Greens
ndwi_xarray = NDWI(landsat_dataset)  # High Concentrations of Water - Blues  

"""
Explanation: Normalized Indices
End of explanation
"""


from utils.data_cube_utilities.dc_water_quality import tsm

tsm_xarray = tsm(landsat_dataset, clean_mask = wofs_xarray.values.astype(bool) ).tsm

"""
Explanation: TSM
End of explanation
"""


def EVI(dataset, c1 = None, c2 = None, L = None):
        return ((dataset.nir - dataset.red)/((dataset.nir  + (c1 * dataset.red) - (c2 *dataset.blue) + L))).rename("EVI")

evi_xarray = EVI(landsat_dataset, c1 = 6, c2 = 7.5, L = 1 )

"""
Explanation: EVI
End of explanation
"""


combined_dataset = xr.merge([landsat_dataset,
          ## <span id="combine_data">Combine Data [&#9652;](#top)</span>  clean_xarray,
          clear_xarray,
          water_xarray,
          shadow_xarray,
          evi_xarray,
          ndbi_xarray,
          ndvi_xarray,
          ndwi_xarray,
          wofs_xarray,
          tsm_xarray])

# Copy original crs to merged dataset 
combined_dataset = combined_dataset.assign_attrs(landsat_dataset.attrs)

combined_dataset

"""
Explanation: <span id="combine_data">Combine Data &#9652;</span>
End of explanation
"""


from utils.data_cube_utilities.import_export import export_xarray_to_multiple_geotiffs

# Ensure the output directory exists before writing to it.
if platform == 'LANDSAT_7':
    !mkdir -p output/geotiffs/landsat7
else:
    !mkdir -p output/geotiffs/landsat8

output_path = "output/geotiffs/landsat{0}/landsat{0}".format(7 if platform=='LANDSAT_7' else 8)

export_xarray_to_multiple_geotiffs(combined_dataset, output_path)

"""
Explanation: <span id="export">Export Data &#9652;</span>
<span id="export_geotiff">Export to GeoTIFF &#9652;</span>
Export each acquisition as a GeoTIFF.
End of explanation
"""


if platform == 'LANDSAT_7':
    !ls -lah output/geotiffs/landsat7/*.tif
else:
    !ls -lah output/geotiffs/landsat8/*.tif

"""
Explanation: Check to see what files were exported. The size of these files is also shown.
End of explanation
"""


if platform == 'LANDSAT_7':
    !gdalinfo output/geotiffs/landsat7/landsat7_2015_01_09_03_06_13.tif
else:
    !gdalinfo output/geotiffs/landsat8/landsat8_2015_01_01_03_07_41.tif

"""
Explanation: Sanity check using gdalinfo to make sure that all of our bands exist    .
End of explanation
"""


if platform == 'LANDSAT_7':
    !tar -cvzf output/geotiffs/landsat7/landsat_7.tar.gz output/geotiffs/landsat7/*.tif
else:
    !tar -cvzf output/geotiffs/landsat8/landsat_8.tar.gz output/geotiffs/landsat8/*.tif

"""
Explanation: Zip all GeoTIFFs.
End of explanation
"""


combined_dataset

import os
import pathlib
from utils.data_cube_utilities.import_export import export_xarray_to_netcdf

# Ensure the output directory exists before writing to it.
ls_num = 7 if platform=='LANDSAT_7' else 8
output_dir = f"output/netcdfs/landsat{ls_num}"
pathlib.Path(output_dir).mkdir(parents=True, exist_ok=True)

output_file_path = output_dir + f"/ls{ls_num}_netcdf_example.nc"
export_xarray_to_netcdf(combined_dataset.red, output_file_path)

"""
Explanation: <span id="export_netcdf">Export to NetCDF &#9652;</span>
Export all acquisitions together as a single NetCDF.
End of explanation
"""


