# Import Python Libraries #
import pandas as pd
import os, glob, wget
import matplotlib.pyplot as plt
from osgeo import gdal
import rasterio
import numpy as np
import math
import re

f = 'landsat_8_aws_scene_list_20200808'
df = pd.read_csv(f)
# Check Scene List #
print(df.shape)
display(df.tail())

# Location of the TA Observatory in the Landsat 8 page #
CLF_path = 38
CLF_row = 33

clf_df = df[df['path'] == CLF_path]
clf_df = clf_df[clf_df['row'] == CLF_row]

## Turn acquisition time into datetime ##
clf_df['acquisitionDate'] = pd.to_datetime(clf_df['acquisitionDate'])

## Cut on cloud cover to less than 5% ##
clf_df = clf_df[clf_df['cloudCover'] < 5 ]

## Find a Spring Month (For aesthetics) ##
clf_df = clf_df[clf_df['acquisitionDate'].dt.month == 4]

## Sort by Most Recent Pass ##
clf_df = clf_df.sort_values('acquisitionDate',ascending=False)
display(clf_df.head())

## Download Directory ##
dir = '/GDF/TAResearch/TAmap/Landsat_imgs'

# Select Scene to download #
map_df = clf_df.head(1)
display(map_df)

# Scene Download URL
download_url = str(map_df['download_url'].values[0])
# print download_url
productId = str(map_df['productId'].values[0])

# Format list of Images to download #
download_urls = download_url.replace('/index.html','/{0}'.format(productId))
download_urls = [ '{0}_{1}.TIF'.format(download_urls, band) for band in bands]
# print download_urls[0]

def centerPOI(max_lat,
    min_lat,
    max_lon,
    min_lon,
    raster,
    poi_lat = ta_df[ta_df['Site'] == 'CLF']['Latitude'].values[0],
    poi_lon = ta_df[ta_df['Site'] == 'CLF']['Longitude'].values[0],
    pixels = 2000):
    '''
    Create a Window for the Rasters to Center on a Point of Interest (POI) within the Raster with given latitude and longitude of site.
    poi_lat = POI latitude in Degrees Decimal
    poi_lon = POI longitude in Degrees Decimal
    p = width of window in pixels. Each Pixel is 30m in width

    '''
    # Raster Longitude and Latitude Width #
    lon_range = max_lat - min_lat
    lat_range = max_lon - min_lon
#     print(max_lat, min_lat, lon_range)
#     print(max_lon, min_lon, lat_range)

    # Relative Position of POI #
    poi_x = (poi_lon - clf_df['min_lon'].values[0]) / lon_range
    poi_y = 1 - (poi_lat - clf_df['min_lat'].values[0]) / lat_range # Flipped y axis #

    # Open Raster and Find Size #
    raster =  gdal.Open(raster)
    dx, dy = raster.RasterXSize, raster.RasterYSize

    # Find Position of POI in Raster #
    poi_x = poi_x * dx
    poi_y = poi_y * dy

    # Calculate Window for Focus on POI #
    # Window Corner #
    x_o, y_o = poi_x - .5 * pixels, poi_y - .5 * pixels
    # Window Width #
    d_x, d_y = pixels, pixels
    window = rasterio.windows.Window(x_o, y_o, d_x, d_y)

    return window

# Create Window #
window = centerPOI(poi_lat = ta_df[ta_df['Site'] == 'CLF']['Latitude'].values[0],
              poi_lon = ta_df[ta_df['Site'] == 'CLF']['Longitude'].values[0],
              max_lat = map_df['max_lon'].values[0],
              min_lat = map_df['min_lon'].values[0],
              max_lon = map_df['max_lat'].values[0],
              min_lon = map_df['min_lat'].values[0],
              raster = image_files[3],
              pixels = 2000)

# print(window)

# Load RGB Bands #
r = rasterio.open(image_files[3]).read(1, window=window)
g = rasterio.open(image_files[2]).read(1, window=window)
b = rasterio.open(image_files[1]).read(1, window=window)
pan = rasterio.open(image_files[8]).read(1, window=window)

# Normalize RGB Bands #
b2 = norm(b.astype(np.float))
b3 = norm(g.astype(np.float))
b4 = norm(r.astype(np.float))
b8 = norm(pan.astype(np.float))
print(np.shape(b4))

# Create RGB Composite by stacking layers #
rgb = np.dstack((b4,b3,b2))
print(np.shape(rgb))

## Create An Image of Each Site Zoomed in ##
fig = plt.figure(figsize=(16,16))
for (i, site, lat, lon) in zip(range(1,len(ta_df)+1),ta_df['Site'].values,ta_df['Latitude'].values,ta_df['Longitude'].values):
    plt.subplot(2,2,i)
    print(i, site, lat, lon)

    ## Finding Offset from recorded CLF lat and Long values to Center CLF ##
    lat_offset, lon_offset = -.015, -.017

    # Create Window #
    window = centerPOI(poi_lat = lat + lat_offset,
              poi_lon = lon +  lon_offset,
              max_lat = map_df['max_lon'].values[0],
              min_lat = map_df['min_lon'].values[0],
              max_lon = map_df['max_lat'].values[0],
              min_lon = map_df['min_lat'].values[0],
              raster = image_files[3],
              pixels = 200)

    # Load Bands #
    r = rasterio.open(image_files[3]).read(1, window=window)
    g = rasterio.open(image_files[2]).read(1, window=window)
    b = rasterio.open(image_files[1]).read(1, window=window)

    # Normalize Bands #
    b2 = norm(b.astype(np.float))
    b3 = norm(g.astype(np.float))
    b4 = norm(r.astype(np.float))

    # Stack RGB for Composite #
    rgb = np.dstack((b4,b3,b2))

    # Plot Site #
    plt.imshow(rgb)
    plt.title('{0} Site ({1:.6f}, {2:.6f})'.format(site,lat,lon))

fig.suptitle('TA Sites with LandSat 8 Composite Images (30m Resolution)', fontsize=16)
plt.savefig('ta_sites.svg')
plt.savefig('ta_sites.png', dpi=300)

## Finding Offset from recorded CLF lat and Long values to Center CLF ##
lat_offset, lon_offset = -.019, -.02

# Create Window #
window = centerPOI(poi_lat = ta_df[ta_df['Site'] == 'CLF']['Latitude'].values[0] + lat_offset,
              poi_lon = ta_df[ta_df['Site'] == 'CLF']['Longitude'].values[0] + lon_offset,
              max_lat = map_df['max_lon'].values[0],
              min_lat = map_df['min_lon'].values[0],
              max_lon = map_df['max_lat'].values[0],
              min_lon = map_df['min_lat'].values[0],
              raster = image_files[3],
              pixels = 2000)

# Load Bands #
r = rasterio.open(image_files[3]).read(1, window=window)
g = rasterio.open(image_files[2]).read(1, window=window)
b = rasterio.open(image_files[1]).read(1, window=window)

# Normalize Bands #
b2 = norm(b.astype(np.float))
b3 = norm(g.astype(np.float))
b4 = norm(r.astype(np.float))

# Create RGB Composite #
rgb = np.dstack((b4,b3,b2))

# FOV Rings #
FOV_x, FOV_y = [], []
FOV_c = []   
for (ang_i, ang_f, xo, yo, color) in zip(ta_df['FOVi'].values,ta_df['FOVf'].values,ta_df['X'].values,ta_df['Y'].values,ta_df['Color'].values):
    r = 30
    x, y = FOV_arc(r, ang_i, ang_f, xo, yo)
    FOV_x.append(x)
    FOV_y.append(y)
    FOV_c.append(color)

## Composite Array Plot ##
fig = plt.figure(figsize=(16,16))

# Satellite Background #
ax=fig.add_subplot(111)
ax.imshow(rgb)
ax.tick_params(size=0)
ax.set_aspect('equal')
ax.set_yticklabels([])
ax.set_xticklabels([])

# FD Site Back Grounds #
clr_clf = 'red'
clr_fd = '#065398'
clr_sd = 'white'

ax2 = fig.add_subplot(111, frame_on=False)

# Plot FOV #
for (x,y,c) in zip(FOV_x, FOV_y, FOV_c):
    if x[0] != 0 :
        alpha = .2
        ax2.fill(x,y,color=clr_fd,alpha=alpha)

# FD Sites #
# for (lat,lon,color,label) in zip(ta_df['Latitude DD'],ta_df['Longitude DD'],ta_df['Color'],ta_df['Site']):
#     x, y = CLF_coord(lon,lat)
for (x,y,label) in zip(ta_df['X'],ta_df['Y'],ta_df['Site']):
    if label == 'CLF' : ax2.scatter(x,y,marker='h', s=90, color=clr_clf)
    else : ax2.scatter(x,y,marker='^', s=90, color=clr_fd)
# SD Sites #
for (x,y) in zip(sd_df['X'],sd_df['Y']):
    ax2.scatter(x,y,marker='.', s=90, color=clr_sd)
# Label Sites #
for (x,y,text) in zip(ta_df['X'],ta_df['Y'],ta_df['Site']):
    #print (x,y,color,label)
    if text != 'CLF':
        if text == 'MD':
            ax2.text(x,y + 1.25,text, verticalalignment='center', horizontalalignment='center', fontsize=16, color='white')
        else :
            ax2.text(x,y - 1.25,text, verticalalignment='center', horizontalalignment='center', fontsize=16, color='white')
ax2.set_xlim([-30,30])
ax2.set_ylim([-30,30])
ax2.tick_params(size=0)
ax2.set_yticklabels([])
ax2.set_xticklabels([])

# Stuff for Legend #
ax2.scatter(40,40,marker='^',color=clr_fd, label='FD Site')
ax2.scatter(40,40,marker='h',color=clr_clf, label='CLF Site')
ax2.scatter(40,40,marker='.',color=clr_sd, label='SD Site')

# Scale #
ax2.plot([-20,-10],[27,27],color='white')
ax2.text(-15, 28, '10 km', verticalalignment='center', horizontalalignment='center', color='white', fontsize=14)

# North Compass #
ax2.plot([-25,-25],[22,27],color='white')
ax2.text(-25, 28, 'North', verticalalignment='center', horizontalalignment='center', color='white', fontsize=14)
ax2.fill([-25,-25.5,-24.5],[27,26,26],color='white')

## Attribution ##
ax.text(0.98, 0.02, 'Greg Furlich\n USGS Landsat 8\n Telescope Array Cosmic Ray Observatory Satellite Map',
        verticalalignment='bottom', horizontalalignment='right',
        transform=ax.transAxes,
        weight='medium',
        color='white', fontsize=15)

# Plot Info #
# plt.title('Telescope Array Cosmic Ray Observatory')
# plt.xlabel('East [km]')
# plt.ylabel('North [km]')

leg =plt.legend(loc='upper right', bbox_to_anchor=(0.98, 0.98), prop={'size':16}, markerscale=2, framealpha = .2)
frame = leg.get_frame()
frame.set_color('white')
frame.set_linewidth(0)
for text in leg.get_texts():
    plt.setp(text, color = 'w')

ax2.set_aspect('equal')

# plt.show()
plt.savefig('ta_map.svg', bbox_inches='tight', pad_inches=-.1)
plt.savefig('ta_map.png', bbox_inches='tight', pad_inches=-.1, dpi=300)