# -*- coding: utf-8 -*-
"""
Created on Tue Mar 20 16:15:51 2018

@author: z5030440
"""

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

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
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
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
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 days difference
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]
dates_diff = []
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]]
                           }) 

# 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)   

np.mean([ np.abs(_['days diff']) for _ in dates_diff])
#%% compare shorelines

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)

# load 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)
    
dates_quad = [datetime(int(_[6:10]), int(_[11:13]), int(_[14:16]), tzinfo= au_tz) for _ in filenames]
zav = []
ztide = []
for i in range(len(dates_diff)):
    # select closest 3D survey
    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]))
    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])
    idx_nan = np.isnan(zs)
    xs = xs[~idx_nan]
    ys = ys[~idx_nan]
    zs = zs[~idx_nan]
    # smooth (LOWESS) satellite shoreline
    idx_beach = [np.min(np.linalg.norm(sl[i][k,:] - narrabeach, axis=1)) < dist_thresh for k in range(sl[i].shape[0])]
    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]]
    # 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
    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(str(0.05*count) + ' - ' + str(len(sl_tide)))    
    else:
        sl_tide = measure.find_contours(survey3d['z'], tide_level)
        sl_tide = sl_tide[np.argmax(np.array([len(_) for _ in sl_tide]))]
    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) 
    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])]
    # interpolate SDS on 3D surface to get elevation
    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()

    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], 'go', markersize=3)
    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)

plt.figure()
plt.plot(zav - ztide)
plt.draw()

zav - ztide
#%% 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()