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

# Preamble
import ee
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import numpy as np
import pandas as pd
from datetime import datetime
import pickle
import pdb
import pytz
from pylab import ginput
import scipy.io as sio
import scipy.interpolate
import os

# image processing modules
import skimage.filters as filters 
import skimage.exposure as exposure
import skimage.transform as transform
import sklearn.decomposition as decomposition
import skimage.morphology as morphology
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
ee.Initialize()

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

#%%
# load SDS shorelines
with open('data\data_gt_l8.pkl', 'rb') as f:
    data = pickle.load(f)
    
# load quadbike dates and convert from datenum to datetime
suffix = '.mat'
dir_name = os.getcwd()
file_name = 'data\quadbike_dates'
file_path = os.path.join(dir_name, file_name + suffix)
quad_dates = sio.loadmat(file_path)['dates']
dt_quad = []
for i in range(quad_dates.shape[0]): 
    dt_quad.append(datetime(quad_dates[i,0], quad_dates[i,1], quad_dates[i,2], tzinfo=au_tz))

# remove overlapping images, keep the one with lowest cloud_cover 
n = len(data['cloud_cover'])
idx_worst = []
for i in range(n):
    date_im = data['date_acquired'][i]
    idx_double = np.isin(data['date_acquired'], date_im)
    if sum(idx_double.astype(int)) > 1:
        idx_worst.append(np.where(idx_double)[0][np.argmax(np.array(data['cloud_cover'])[idx_double])])
dt_sat = []
new_meta = {'contours':[],
            'cloud_cover':[],
            'geom_rmse_model':[],
            'gcp_model':[],
            'quality':[],
            'sun_azimuth':[],
            'sun_elevation':[]}
for i in range(n):
    if not np.isin(i,idx_worst):
        dt_sat.append(data['dt'][i].astimezone(au_tz))
        new_meta['contours'].append(data['contours'][i])
        new_meta['cloud_cover'].append(data['cloud_cover'][i])
        new_meta['geom_rmse_model'].append(data['geom_rmse_model'][i])
        new_meta['gcp_model'].append(data['gcp_model'][i])
        new_meta['quality'].append(data['quality'][i])
        new_meta['sun_azimuth'].append(data['sun_azimuth'][i])
        new_meta['sun_elevation'].append(data['sun_elevation'][i])
        
# calculate difference between days
diff_days = [ [(x - _).days for _ in dt_quad] for x in dt_sat]
day_thresh = 15
idx_close = [utils.find_indices(_, lambda e: abs(e) < day_thresh) for _ in diff_days]

# put everything in a dictionnary and save it
wl_comp = []
for i in range(len(dt_sat)):
    wl_comp.append({'sat dt': dt_sat[i],
                    'quad dt': [dt_quad[_] for _ in idx_close[i]],
                    'days diff': [diff_days[i][_] for _ in idx_close[i]],
                    'contours':  new_meta['contours'][i],
                    'cloud_cover':  new_meta['cloud_cover'][i],
                    'geom_rmse_model':  new_meta['geom_rmse_model'][i],
                    'gcp_model':  new_meta['gcp_model'][i],
                    'quality':  new_meta['quality'][i],
                    'sun_azimuth':  new_meta['sun_azimuth'][i],
                    'sun_elevation':  new_meta['sun_elevation'][i]})

with open('wl_l8_comparison.pkl', 'wb') as f:
    pickle.dump(wl_comp, f)
    
#%%
with open('data\wl_l8_comparison.pkl', 'rb') as f:
    wl = pickle.load(f)
# load quadbike dates and convert from datenum to datetime
suffix = '.mat'
dir_name = os.getcwd()
subfolder_name = 'data\quadbike_surveys'
file_path = os.path.join(dir_name, subfolder_name)
file_names = os.listdir(file_path)

for i in range(len(file_names)):
    fn_mat = os.path.join(file_path, file_names[i])
    years = int(file_names[i][6:10])
    months = int(file_names[i][11:13])
    days = int(file_names[i][14:16])
    for j in range(len(wl)):
        if wl[j]['quad dt'][0] == datetime(years, months, days, tzinfo=au_tz):
            quad_mat = sio.loadmat(fn_mat)
            wl[j].update({'quad_data':{'x':quad_mat['x'],
                         'y':quad_mat['y'],
                         'z':quad_mat['z'],
                         'dt': datetime(years, months, days, tzinfo=au_tz)}})
with open('data\wl_final.pkl', 'wb') as f:
    pickle.dump(wl, f) 
#%%
with open('data\wl_final.pkl', 'rb') as f:
    wl = pickle.load(f) 
    
i = 0        
x = wl[i]['quad_data']['x']
y = wl[i]['quad_data']['y']
z = wl[i]['quad_data']['z']
x = x.reshape(x.shape[0] * x.shape[1])
y = y.reshape(y.shape[0] * y.shape[1])
z = z.reshape(z.shape[0] * z.shape[1])
idx_nan = np.isnan(z)

x_nan = x[idx_nan]
y_nan = y[idx_nan]
z_nan = z[idx_nan]

x_nonan = x[~idx_nan]
y_nonan = y[~idx_nan]
z_nonan = z[~idx_nan]

xs = x_nonan[::10]
ys = y_nonan[::10]
zs = z_nonan[::10]

xq = wl[i]['contours'][:,0]
yq = wl[i]['contours'][:,1]

# cut xq around xs
np.min(xs)
np.max(xs)
np.min(ys)
np.max(ys)

idx_x = np.logical_and(xq < np.max(xs), xq > np.min(xs))
idx_y = np.logical_and(yq < np.max(ys), yq > np.min(ys))
idx_in = np.logical_and(idx_x, idx_y)

xq = xq[idx_in]
yq = yq[idx_in]

for i in range(len(xq)):
    idx_x = np.logical_and(xs < xq[i] + 10, xs > xq[i] - 10)
    idx_y = np.logical_and(ys < yq[i] + 10, ys > yq[i] - 10)
    xint = xs[idx_x]
    yint = ys[idx_y]

f = interpolate.interp2d(xs, ys, zs, kind='linear')
zq = f(xq,yq)

plt.figure()
plt.grid()
plt.scatter(xs, ys, s=10, c=zs, marker='o', cmap=cm.get_cmap('jet'),
            label='quad data')
plt.plot(xq,yq,'r-o', markersize=5, label='SDS')
plt.axis('equal')
plt.legend()
plt.colorbar(label='mAHD')
plt.xlabel('Eastings [m]')
plt.ylabel('Northings [m]')
plt.show()

plt.figure()
plt.plot(zq[:,0])
plt.show()


plt.figure()
plt.grid()
plt.scatter(x_nonan, y_nonan, s=10, c=z_nonan, marker='o', cmap=cm.get_cmap('jet'),
            label='quad data')
#plt.plot(x_nan, y_nan, 'k.', label='nans')
plt.plot(xq,yq,'r-o', markersize=5, label='SDS')
plt.axis('equal')
plt.legend()
plt.colorbar(label='mAHD')
plt.xlabel('Eastings [m]')
plt.ylabel('Northings [m]')
plt.show()


z2 = scipy.interpolate.griddata([x, y], z, [xq, yq], method='linear')

f_interp = scipy.interpolate.interp2d(x1,y1,z1, kind='linear')
    
    
sio.savemat('shoreline1.mat', {'x':xq, 'y':yq})

from scipy import interpolate
x = np.arange(-5.01, 5.01, 0.01)
y = np.arange(-5.01, 5.01, 0.01)
xx, yy = np.meshgrid(x, y)
z = np.sin(xx**2+yy**2)
f = interpolate.interp2d(x, y, z, kind='cubic')

xnew = np.arange(-5.01, 5.01, 1e-2)
ynew = np.arange(-5.01, 5.01, 1e-2)
znew = f(xnew, ynew)
plt.plot(x, z[:, 0], 'ro-', xnew, znew[:, 0], 'b-')
plt.show()