import geopandas
import geoviews as gv
import xarray as xr
import pyproj
from pyproj import Proj
import requests

import boto3
import rasterio as rio
from rasterio.session import AWSSession
import os
from netrc import netrc
from subprocess import Popen
from platform import system
from getpass import getpass

import hvplot.xarray
import holoviews as hv
from bokeh.models import FixedTicker
hv.extension('bokeh')
gv.extension('bokeh', 'matplotlib')

from matplotlib.colors import ListedColormap
import numpy as np

import warnings
warnings.filterwarnings('ignore')

# User-specified input
data_dir ='https://archive.podaac.earthdata.nasa.gov/podaac-ops-cumulus-protected/OPERA_L3_DSWX-HLS_PROVISIONAL_V0/OPERA_L3_DSWx_HLS_T42RUQ_20220930T055641Z_20221028T020516Z_S2A_30_v0.0_'
bandlist = ['B01_WTR', 'B02_BWTR', 'B03_CONF']
bandpath = f"{data_dir}%s.tiff"

# Generates authentication token
# Asks for your Earthdata username and password for first time, if netrc does not exists.

urs = 'urs.earthdata.nasa.gov'    # Earthdata URL endpoint for authentication
prompts = ['Enter NASA Earthdata Login Username: ',
           'Enter NASA Earthdata Login Password: ']

# Determine the OS (Windows machines usually use an '_netrc' file)
netrc_name = "_netrc" if system()=="Windows" else ".netrc"

# Determine if netrc file exists, and if so, if it includes NASA Earthdata Login Credentials
try:
    netrcDir = os.path.expanduser(f"~/{netrc_name}")
    netrc(netrcDir).authenticators(urs)[0]

# Below, create a netrc file and prompt user for NASA Earthdata Login Username and Password
except FileNotFoundError:
    homeDir = os.path.expanduser("~")
    Popen('touch {0}{2} | echo machine {1} >> {0}{2}'.format(homeDir + os.sep, urs, netrc_name), shell=True)
    Popen('echo login {} >> {}{}'.format(getpass(prompt=prompts[0]), homeDir + os.sep, netrc_name), shell=True)
    Popen('echo \'password {} \'>> {}{}'.format(getpass(prompt=prompts[1]), homeDir + os.sep, netrc_name), shell=True)
    # Set restrictive permissions
    Popen('chmod 0600 {0}{1}'.format(homeDir + os.sep, netrc_name), shell=True)

    # Determine OS and edit netrc file if it exists but is not set up for NASA Earthdata Login
except TypeError:
    homeDir = os.path.expanduser("~")
    Popen('echo machine {1} >> {0}{2}'.format(homeDir + os.sep, urs, netrc_name), shell=True)
    Popen('echo login {} >> {}{}'.format(getpass(prompt=prompts[0]), homeDir + os.sep, netrc_name), shell=True)
    Popen('echo \'password {} \'>> {}{}'.format(getpass(prompt=prompts[1]), homeDir + os.sep, netrc_name), shell=True)

# Generates the temporary S3 credentials
s3_cred_endpoint = 'https://archive.podaac.earthdata.nasa.gov/s3credentials'
def get_temp_creds():
    temp_creds_url = s3_cred_endpoint
    return requests.get(temp_creds_url).json()
temp_creds_req = get_temp_creds()

session = boto3.Session(aws_access_key_id=temp_creds_req['accessKeyId'], 
                        aws_secret_access_key=temp_creds_req['secretAccessKey'],
                        aws_session_token=temp_creds_req['sessionToken'],
                        region_name='us-west-2')

rio_env = rio.Env(AWSSession(session),
                  GDAL_DISABLE_READDIR_ON_OPEN='EMPTY_DIR',
                  CPL_VSIL_CURL_ALLOWED_EXTENSIONS='TIF, TIFF',
                  GDAL_HTTP_COOKIEFILE=os.path.expanduser('~/cookies.txt'),
                  GDAL_HTTP_COOKIEJAR=os.path.expanduser('~/cookies.txt'))
rio_env.__enter__()

# Function to read each layers and stack them to create a geocube
def stack_bands(bandpath:str, bandlist:list): 
    '''
    Returns geocube with three bands stacked into one multi-dimensional array.
            Parameters:
                    bandpath (str): Path to bands that should be stacked
                    bandlist (list): Three bands that should be stacked
            Returns:
                    bandStack (xarray Dataset): Geocube with stacked bands
                    crs (int): Coordinate Reference System corresponding to bands
    '''
    bandStack = []; bandS = []; bandStack_ = [];
    for i,band in enumerate(bandlist):
        if i==0:
            bandStack_ = xr.open_rasterio(bandpath%band)
            crs = pyproj.CRS.to_epsg(pyproj.CRS.from_proj4(bandStack_.crs))
            bandStack_ = bandStack_ * bandStack_.scales[0]
            bandStack = bandStack_.squeeze(drop=True)
            bandStack = bandStack.to_dataset(name='z')
            bandStack.coords['band'] = i+1
            bandStack = bandStack.rename({'x':'longitude', 'y':'latitude', 'band':'band'})
            bandStack = bandStack.expand_dims(dim='band')  
        else:
            bandS = xr.open_rasterio(bandpath%band)
            bandS = bandS * bandS.scales[0]
            bandS = bandS.squeeze(drop=True)
            bandS = bandS.to_dataset(name='z')
            bandS.coords['band'] = i+1
            bandS = bandS.rename({'x':'longitude', 'y':'latitude', 'band':'band'})
            bandS = bandS.expand_dims(dim='band')
            bandStack = xr.concat([bandStack, bandS], dim='band')
    return bandStack, crs

# Creates geocube of stacked bands
da, crs = stack_bands(bandpath, bandlist)

# Creates basemap using geoviews
base = gv.tile_sources.EsriImagery.opts(width=1000, height=1000, padding=0.1)

# Mask nodata values (255)
da_masked = da.where(da['z'] != 255.) 
B01_WTR = da_masked.z.sel({'band':1})
B02_BWTR = da_masked.z.sel({'band':2}) 
B03_CONF = da_masked.z.sel({'band':3}) 

# Visualize B01 - WATER CLASSIFICATION LAYER
# Parameters for Colorbar
levels = [0, 0.9, 1.9, 2.9, 7.9, 8.9, 10]
color_key = {
    "Not Water": "#ffffff",
    "Open Water": "#0000ff",
    "Partial Surface Water": "#00ff00",
    "Reserved": "#000000",
    "Snow/Ice": "#00ffff",
    "Clouds/Cloud Shadow": "#7f7f7f"
}

ticks = [0.5, 1.5, 2.5, 5.5, 8.5, 9.5]
ticker = FixedTicker(ticks=ticks)
labels = dict(zip(ticks, color_key))

fig = B01_WTR.hvplot.image(x='longitude', 
                          y='latitude', 
                          crs=crs, 
                          rasterize=False, 
                          dynamic=False, 
                          aspect='equal', 
                          frame_width=500,
                          clim=(0,255),
                          frame_height=500, 
                          alpha=0.8).opts(title=f"B01 WTR", xlabel='Longitude', ylabel='Latitude', color_levels= levels, cmap=tuple(color_key.values()), 
                                        colorbar_opts={'ticker': ticker, 'major_label_overrides': labels}, clim=(0,10)) * base
#Uncomment to save
#hvplot.save(fig, 'B01_WTR.png')
fig

# Visualize B02 - BINARY WATER LAYER
# Parameters for Colorbar
levels = [0, 0.9, 1.9, 7.9, 8.9, 10]
color_key = {
    "Not Water": "#ffffff",
    "Water": "#0000ff",
    "Reserved": "#000000",
    "Snow/Ice": "#00ffff",
    "Clouds/Cloud Shadow": "#7f7f7f"
}

ticks = [0.5, 1.5, 5.5, 8.5, 9.5]
ticker = FixedTicker(ticks=ticks)
labels = dict(zip(ticks, color_key))

fig = B02_BWTR.hvplot.image(x='longitude', 
                          y='latitude', 
                          crs=crs, 
                          rasterize=False, 
                          dynamic=False, 
                          aspect='equal', 
                          frame_width=500,
                          clim=(0,255),
                          frame_height=500, 
                          alpha=0.8).opts(title=f"B02 BWTR", xlabel='Longitude', ylabel='Latitude', color_levels= levels, cmap=tuple(color_key.values()), 
                                        colorbar_opts={'ticker': ticker, 'major_label_overrides': labels}, clim=(0,10)) * base
#Uncomment to save
#hvplot.save(fig, 'B02_BWTR.png')
fig

# Visualize B03 - CONFIDENCE LAYER
# Parameters for Colorbar
with rio.open(f"{data_dir}B03_CONF.tiff") as ds:
    colormap = ds.colormap(1)

color_key = ListedColormap([np.array(colormap[key]) / 255 for key in range(256)])
ticks = [0, 50, 100, 175, 245, 255]
ticker = FixedTicker(ticks=ticks)
labels = dict(zip(ticks, ["0", "Confidence", "100", "Reserved", "Snow/Ice", "Clouds/Cloud Shadow"]))

fig = B03_CONF.hvplot.image(x='longitude', 
                          y='latitude', 
                          crs=crs, 
                          rasterize=False, 
                          dynamic=False, 
                          aspect='equal', 
                          frame_width=500,
                          clim=(0,255),
                          frame_height=500, 
                          alpha=0.8).opts(title=f"B03 CONF", xlabel='Longitude', ylabel='Latitude', cmap=color_key,
                                    colorbar_opts={'ticker': ticker, 'major_label_overrides': labels}) * base
#Uncomment to save
#hvplot.save(fig, 'B03_CONF.png')
fig

rio_env.__exit__()