# -*- coding: utf-8 -*-

#==========================================================#
# Compare Narrabeen SDS with 3D quadbike surveys
#==========================================================#

# Initial settings

import os
import numpy as np
import matplotlib.pyplot as plt
import pdb
import ee

import matplotlib.dates as mdates
import matplotlib.cm as cm
from datetime import datetime, timedelta
import pickle
import pytz
import scipy.io as sio
import scipy.interpolate as interpolate
import statsmodels.api as sm
import skimage.measure as measure

# my functions
import functions.utils as utils
import functions.sds as sds

# some settings
np.seterr(all='ignore') # raise/ignore divisions by 0 and nans
plt.rcParams['axes.grid'] = True
plt.rcParams['figure.max_open_warning'] = 100

au_tz = pytz.timezone('Australia/Sydney')

# load quadbike dates and convert from datenum to datetime
filename = 'data\quadbike\survey_dates.mat'
filepath = os.path.join(os.getcwd(), filename)
dates_quad = sio.loadmat(filepath)['dates'] # matrix containing year, month, day
dates_quad = [datetime(dates_quad[i,0], dates_quad[i,1], dates_quad[i,2],
                       tzinfo=au_tz) for i in range(dates_quad.shape[0])]

# load timestamps from satellite images
satname = 'L8'
sitename = 'NARRA'
filepath = os.path.join(os.getcwd(), 'data', satname, sitename)
with open(os.path.join(filepath, sitename + '_output2' + '.pkl'), 'rb') as f:
    output = pickle.load(f)
dates_l8 = output['t']
# convert to AEST
dates_l8 = [_.astimezone(au_tz) for _ in dates_l8]

# load wave data (already AEST)
filename = 'data\wave\SydneyProcessed.mat'
filepath = os.path.join(os.getcwd(), filename)
wave_data = sio.loadmat(filepath)
idx = utils.find_indices(wave_data['dates'][:,0], lambda e: e >= dates_l8[0].year and e <= dates_l8[-1].year)
hsig = np.array([wave_data['Hsig'][i][0] for i in idx])
wdir = np.array([wave_data['Wdir'][i][0] for i in idx])
dates_wave = [datetime(wave_data['dates'][i,0], wave_data['dates'][i,1],
                       wave_data['dates'][i,2], wave_data['dates'][i,3],
                       wave_data['dates'][i,4], wave_data['dates'][i,5],
                       tzinfo=au_tz) for i in idx]
    
# load tide data (already AEST)
filename = 'SydTideData.mat'
filepath = os.path.join(os.getcwd(), 'data', 'tide', filename)
tide_data = sio.loadmat(filepath)
idx = utils.find_indices(tide_data['dates'][:,0], lambda e: e >= dates_l8[0].year and e <= dates_l8[-1].year)
tide = np.array([tide_data['tide'][i][0] for i in idx])
dates_tide = [datetime(tide_data['dates'][i,0], tide_data['dates'][i,1],
                       tide_data['dates'][i,2], tide_data['dates'][i,3],
                       tide_data['dates'][i,4], tide_data['dates'][i,5],
                       tzinfo=au_tz) for i in idx]

#%% make a plot of all the dates with wave data
orange = [255/255,140/255,0]
blue = [0,191/255,255/255]
f = plt.figure()
months = mdates.MonthLocator()
month_fmt = mdates.DateFormatter('%b %Y')
days = mdates.DayLocator()
years = [2013,2014,2015,2016]
for k in range(len(years)):
    sel_year = years[k]
    ax = plt.subplot(4,1,k+1)
    idx_year = utils.find_indices(dates_wave, lambda e : e.year >= sel_year and e.year <= sel_year)
    plt.plot([dates_wave[i] for i in idx_year], [hsig[i] for i in idx_year], 'k-', linewidth=0.5)
    hsigmax = np.nanmax([hsig[i] for i in idx_year])
    cbool = True
    for j in range(len(dates_quad)):
        if dates_quad[j].year == sel_year:
            if cbool:
                plt.plot([dates_quad[j], dates_quad[j]], [0, hsigmax], color=orange, label='survey')
                cbool = False
            else:
                plt.plot([dates_quad[j], dates_quad[j]], [0, hsigmax], color=orange)
    cbool = True
    for j in range(len(dates_l8)):
        if dates_l8[j].year == sel_year:
            if cbool:
                plt.plot([dates_l8[j], dates_l8[j]], [0, hsigmax], color=blue, label='landsat8')
                cbool = False
            else:
                plt.plot([dates_l8[j], dates_l8[j]], [0, hsigmax], color=blue)
    if k == 3:
        plt.legend()
    plt.xlim((datetime(sel_year,1,1), datetime(sel_year,12,31, tzinfo=au_tz)))
    plt.ylim((0, hsigmax))
    plt.ylabel('Hs [m]')
    ax.xaxis.set_major_locator = months
    ax.xaxis.set_major_formatter(month_fmt)
f.subplots_adjust(hspace=0.2)
plt.draw()

#%% calculate difference between dates (quad and sat)
diff_days = [ [(x - _).days for _ in dates_quad] for x in dates_l8]
max_diff = 10
idx_closest = [utils.find_indices(_, lambda e: abs(e) <= max_diff) for _ in diff_days]
# store in dates_diff dictionnary
dates_diff = []
cloud_cover = []
for i in range(len(idx_closest)):
    if not idx_closest[i]:
        continue
    elif len(idx_closest[i]) > 1:
        idx_best = np.argmin(np.abs([diff_days[i][_] for _ in idx_closest[i]]))
        dates_temp = [dates_quad[_] for _ in idx_closest[i]]
        days_temp = [diff_days[i][_] for _ in idx_closest[i]]
        dates_diff.append({"date sat": dates_l8[i],
                           "date quad": dates_temp[idx_best],
                           "days diff": days_temp[idx_best]})
    else:
        dates_diff.append({"date sat": dates_l8[i],
                           "date quad": dates_quad[idx_closest[i][0]],
                           "days diff": diff_days[i][idx_closest[i][0]]
                           }) 
    # store cloud data
    cloud_cover.append(output['cloud_cover'][i])

# store wave data
wave_hsig = []
for i in range(len(dates_diff)):
    wave_hsig.append(hsig[np.argmin(np.abs(np.array([(dates_diff[i]['date sat'] - _).total_seconds() for _ in dates_wave])))])

# make a plot
plt.figure()
counter = 0
for i in range(len(dates_diff)):
    counter = counter + 1
    if dates_diff[i]['date quad'] > dates_diff[i]['date sat']:
        date_min = dates_diff[i]['date sat']
        date_max = dates_diff[i]['date quad']
        color1 = orange
        color2 = blue
    else:
        date_min = dates_diff[i]['date quad']
        date_max = dates_diff[i]['date sat']
        color1 = blue
        color2 = orange
    idx_t = utils.find_indices(dates_wave, lambda e : e >= date_min and e <= date_max)
    hsigmax = np.nanmax([hsig[i] for i in idx_t])
    hsigmin = np.nanmin([hsig[i] for i in idx_t])
    if counter > 9:
        counter = 1
        plt.figure()
    ax = plt.subplot(3,3,counter)
    plt.plot([dates_wave[i] for i in idx_t], [hsig[i] for i in idx_t], 'k-', linewidth=1.5)
    plt.plot([date_min, date_min], [0, 4.5], color=color2, label='survey')
    plt.plot([date_max, date_max], [0, 4.5], color=color1, label='landsat8')
    plt.ylabel('Hs [m]')
    ax.xaxis.set_major_locator(mdates.DayLocator(tz=au_tz))
    ax.xaxis.set_minor_locator(mdates.HourLocator(tz=au_tz))
    ax.xaxis.set_major_formatter(mdates.DateFormatter('%d'))
    ax.xaxis.set_minor_locator(months)
    plt.title(dates_diff[i]['date sat'].strftime('%b %Y') + '   (' + str(abs(dates_diff[i]['days diff'])) + ' days)')
    plt.draw()
    plt.gcf().subplots_adjust(hspace=0.5)   

# mean day difference
np.mean([ np.abs(_['days diff']) for _ in dates_diff])

#%% Compare shorelines in elevation

dist_thresh = 200   # maximum distance between an sds point and a narrabeen point
frac_smooth = 1./10 # fraction of the data used for smoothing (the bigger the smoother)
dist_buffer = 50    # buffer of points selected for interpolation

# load quadbike .mat files 
foldername = 'data\quadbike\surveys3D'
folderpath = os.path.join(os.getcwd(), foldername)
filenames = os.listdir(folderpath)

# get the satellite shorelines
sl = output['shorelines']

# load narrabeen beach points (manually digitized)
with open(os.path.join(os.getcwd(), 'olddata', 'narra_beach' + '.pkl'), 'rb') as f:
    narrabeach = pickle.load(f)

# get dates from filenames
dates_quad = [datetime(int(_[6:10]), int(_[11:13]), int(_[14:16]), tzinfo= au_tz) for _ in filenames]

zav = []
ztide = []
sl_gt = []
sl_sds = []
for i in range(len(dates_diff)):
    
    # select closest 3D survey and load .mat file
    idx_closest = np.argmin(np.abs(np.array([(dates_diff[i]['date sat'] - _).days for _ in dates_quad])))
    survey3d = sio.loadmat(os.path.join(folderpath, filenames[idx_closest]))
    # reshape to a vector
    xs = survey3d['x'].reshape(survey3d['x'].shape[0] * survey3d['x'].shape[1])
    ys = survey3d['y'].reshape(survey3d['y'].shape[0] * survey3d['y'].shape[1])
    zs = survey3d['z'].reshape(survey3d['z'].shape[0] * survey3d['z'].shape[1])
    # remove nan values
    idx_nan = np.isnan(zs)
    xs = xs[~idx_nan]
    ys = ys[~idx_nan]
    zs = zs[~idx_nan]
    
    # select point of sds that are close to the manually digitized points
    idx_beach = [np.min(np.linalg.norm(sl[i][k,:] - narrabeach, axis=1)) < dist_thresh for k in range(sl[i].shape[0])]
    # smooth (LOWESS) satellite shoreline    
    sl_smooth = sm.nonparametric.lowess(sl[i][idx_beach,0],sl[i][idx_beach,1], frac=frac_smooth, it = 6)
    sl_smooth = sl_smooth[:,[1,0]]
    sl_sds.append(sl_smooth)
    # find water level at the time the image was acquired
    idx_closest = np.argmin(np.abs(np.array([(dates_diff[i]['date sat'] - _).total_seconds() for _ in dates_tide])))
    tide_level = tide[idx_closest]
    ztide.append(tide_level)
    
    # find contour corresponding to the water level on 3D surface (if below minimum, add 0.05m increments)
    if tide_level < np.nanmin(survey3d['z']):
        tide_level = np.nanmin(survey3d['z'])
        sl_tide = measure.find_contours(survey3d['z'], tide_level)
        sl_tide = sl_tide[np.argmax(np.array([len(_) for _ in sl_tide]))]
        count = 0
        while len(sl_tide) < 900:
            count = count + 1
            tide_level = tide_level + 0.05*count
            sl_tide = measure.find_contours(survey3d['z'], tide_level)
            sl_tide = sl_tide[np.argmax(np.array([len(_) for _ in sl_tide]))]
        print('added ' + str(0.05*count) + ' cm   -  contour with ' + str(len(sl_tide)) + ' points')    
    else:
        sl_tide = measure.find_contours(survey3d['z'], tide_level)
        sl_tide = sl_tide[np.argmax(np.array([len(_) for _ in sl_tide]))]
    # remove nans
    if np.any(np.isnan(sl_tide)):
        index_nan = np.where(np.isnan(sl_tide))[0]
        sl_tide = np.delete(sl_tide, index_nan, axis=0) 
    # get x,y coordinates
    xtide = [survey3d['x'][int(np.round(sl_tide[m,0])), int(np.round(sl_tide[m,1]))] for m in range(sl_tide.shape[0])]
    ytide = [survey3d['y'][int(np.round(sl_tide[m,0])), int(np.round(sl_tide[m,1]))] for m in range(sl_tide.shape[0])]
    sl_gt.append(np.transpose(np.array([np.array(xtide), np.array(ytide)])))
    # interpolate SDS on 3D surface to get elevation (point by point)
    zq = np.zeros((sl_smooth.shape[0], 1))
    for j in range(sl_smooth.shape[0]):
        xq = sl_smooth[j,0]
        yq = sl_smooth[j,1]
        dist_q = np.linalg.norm(np.transpose(np.array([[xq - _ for _ in xs],[yq - _ for _ in ys]])), axis=1)
        idx_buffer = dist_q <= dist_buffer       
        tck = interpolate.bisplrep(xs[idx_buffer], ys[idx_buffer], zs[idx_buffer])
        zq[j] = interpolate.bisplev(xq, yq, tck)
#        plt.figure()
#        plt.axis('equal')
#        plt.scatter(xs, ys, s=10, c=zs, marker='o', cmap=cm.get_cmap('jet'),
#                    label='quad data')
#        plt.plot(xs[idx_buffer], ys[idx_buffer], 'ko')
#        plt.plot(xq,yq,'ro')
#        plt.draw()
        
    # store the alongshore median elevation
    zav.append(np.median(utils.reject_outliers(zq, m=2)))
    
    # make plot
    red = [255/255, 0, 0]
    gray = [0.75, 0.75, 0.75]
    plt.figure()
    plt.subplot(121)
    plt.axis('equal')
    plt.scatter(xs, ys, s=10, c=zs, marker='o', cmap=cm.get_cmap('jet'),
            label='3D survey')
    plt.plot(xtide, ytide, '--', color=gray, linewidth=2.5, label='tide level contour')
    plt.plot(sl_smooth[:,0], sl_smooth[:,1], '-', color=red, linewidth=2.5, label='SDS')
#    plt.plot(sl[i][idx_beach,0], sl[i][idx_beach,1], 'w-', linewidth=2)
    plt.xlabel('Eastings [m]')
    plt.ylabel('Northings [m]')
    plt.title('Shoreline comparison')
    plt.colorbar(label='mAHD')
    plt.legend()
    plt.ylim((6266100, 6267000))
    plt.subplot(122)
    plt.plot(sl_smooth[:,1], zq, 'ko-', markersize=5)
    plt.plot([sl_smooth[0,1], sl_smooth[-1,1]], [zav[i], zav[i]], 'r--', label='median')
    plt.plot([sl_smooth[0,1], sl_smooth[-1,1]], [ztide[i], ztide[i]], 'g--', label = 'measured tide')
    plt.xlabel('Northings [m]')
    plt.ylabel('Elevation [mAHD]')
    plt.title('Alongshore SDS elevation')
    plt.legend()
    mng = plt.get_current_fig_manager()                                         
    mng.window.showMaximized()
    plt.tight_layout()   
    plt.draw()
    
    print(i)

#%% Calculate some error statistics
zav = np.array(zav)
ztide = np.array(ztide)


f = plt.figure()
plt.subplot(3,1,1)
plt.bar(np.linspace(1,len(zav),len(zav)), zav-ztide)
plt.ylabel('Error in z [m]')
plt.title('Elevation error')
plt.xticks([])
plt.draw()

plt.subplot(3,1,2)
plt.bar(np.linspace(1,len(zav),len(zav)), wave_hsig, color=orange)
plt.ylabel('Hsig [m]')
plt.xticks([])
plt.draw()

plt.subplot(3,1,3)
plt.bar(np.linspace(1,len(zav),len(zav)), np.array(cloud_cover)*100, color='g')
plt.ylabel('Cloud cover %')
plt.xlabel('comparison #')
plt.grid(False)
plt.grid(axis='y')
f.subplots_adjust(hspace=0)
plt.draw()

np.sqrt(np.mean((zav - ztide)**2))
#%% plot to show LOWESS smoothing
#i = 0
#idx_beach = [np.min(np.linalg.norm(sl[i][k,:] - narrabeach, axis=1)) < dist_thresh for k in range(sl[i].shape[0])]
#x = sl[i][idx_beach,0]
#y = sl[i][idx_beach,1]
#sl_smooth = lowess(x,y, frac=1./10, it = 10)
#
#plt.figure()
#plt.axis('equal')
#plt.scatter
#plt.plot(x,y,'bo', linewidth=2, label='original SDS')
#plt.plot(sl_smooth[:,1], sl_smooth[:,0], 'ro', linewidth=2, label='smoothed SDS')
#plt.legend()
#plt.xlabel('Eastings [m]')
#plt.ylabel('Northings [m]')
#plt.title('Local weighted scatterplot smoothing (LOWESS)')
#plt.draw()