import pandas as pd
import matplotlib as mpl
import matplotlib.pyplot as plt
from math import sin, cos, sqrt, atan2, radians
import numpy as np
import geopandas as gp
from matplotlib.tri import Triangulation

#read in data
df = pd.read_csv('NC_CurrentMeter_All_Phase1_all_mobile_2013_05_02.csv')
#convert data to CRS in delft model
gdf = gp.GeoDataFrame(df, geometry = gp.points_from_xy(df.Longitude,df.Latitude))
gdf.crs = {'init':'epsg:4326'}
gdf.geometry = gdf.geometry.to_crs({'init':'epsg:32118'})
gdf.geometry.crs = (4326)
gdf.geometry = gdf.geometry.to_crs(32118)
gdf['UTM_X'] = gdf.geometry.x
gdf['UTM_Y'] = gdf.geometry.y
gdf.to_csv('NC_CurrentMeter_All_Phase1_all_mobile_2013_05_02_EPSG_32118_V2.csv',index=False)
#Define variables and columns for calculations
df['latR'] = ''
df['lonR'] = ''
df['dlat']= ''
df['dlon']= ''
df['a']= ''
df['c']= ''
df['d']= ''
R = 6373 #Radius of Earth
#loop through all transects
for transect_id in df['transect'].unique():
    T1 = df.loc[(df['transect']==transect_id)]
    for i in range(1,len(T1)):#loop to calculate distance between points
            T1.iloc[0,T1.columns.get_loc('latR')] = radians(T1.iloc[0,9])
            T1.iloc[0,T1.columns.get_loc('lonR')] = radians(T1.iloc[0,10])
            T1.iloc[i,T1.columns.get_loc('latR')] = radians(T1.iloc[i,9])#convert lat to radians
            T1.iloc[i,T1.columns.get_loc('lonR')] = radians(T1.iloc[i,10])#convert lon to radians
            T1.iloc[i,T1.columns.get_loc('dlat')] = T1.iloc[i,T1.columns.get_loc('latR')]-T1.iloc[i-1,T1.columns.get_loc('latR')]#get dif between lat
            T1.iloc[i,T1.columns.get_loc('dlon')] = T1.iloc[i,T1.columns.get_loc('lonR')]-T1.iloc[i-1,T1.columns.get_loc('lonR')]#get dif between lon
            T1.iloc[i,T1.columns.get_loc('a')] = (sin(T1.iloc[i,T1.columns.get_loc('dlat')]/2))**2 + cos(T1.iloc[i,T1.columns.get_loc('latR')]) * cos(T1.iloc[i-1,T1.columns.get_loc('latR')]) * (sin(T1.iloc[i,T1.columns.get_loc('dlon')]/2))**2 #calc a
            T1.iloc[i,T1.columns.get_loc('c')] = 2 * atan2(sqrt(T1.iloc[i,T1.columns.get_loc('a')]), sqrt(1 - T1.iloc[i,T1.columns.get_loc('a')])) #calc c 
            T1.iloc[i,T1.columns.get_loc('d')] = 1000*R*T1.iloc[i,T1.columns.get_loc('c')] #calc dist between points
    T1['d']=pd.to_numeric(T1['d']) 
    T1['dist']=T1['d'].cumsum() #get cumulative sum of distances
    #Column names to distances data for plotting
    depths1 = T1.columns[11:43] #ADCP sensor depth reading column names to list
    depths = [float(d.replace('m','')) for d in depths1]
    V = T1.values[:,11:43].astype(float) #get velocity data all rows columns 11-43
    T1['date'] = T1['year'].astype(str) + '-' + T1['mon'].astype(str).str.zfill(2) + '-' + T1['day'].astype(str).str.zfill(2)
    # V = V.flatten() #flatten matrix into list
    # V = pd.to_numeric(V)
    # dist = np.repeat(T1['dist'],len(depths1))
    #Plotting 
    # print(f'Plotting Transect {transect_id} Location {T1.iloc[i,0]}')
    # plt.scatter(dist,depths,c=V)
    # plt.title(f"Location {T1.iloc[i,0]} Transect {transect_id}")
    # plt.xlabel("Distance Along Transect (m)")
    # plt.ylabel("Height Above Bed (m)")
    # plt.colorbar()
    # plt.savefig(f"Location {T1.iloc[i,0]} Transect {transect_id}")
# #__________________________________________________________________________________________________________________________
    #quad contourf
    from math import ceil

    #x = np.reshape(dist.values, (len(T1), len(depths1)))
    #y = np.reshape(depths, (len(T1), len(depths1)))
    #c = np.reshape(V, (len(T1), len(depths1)))

    fig, ax = plt.subplots(figsize=(16,4))
    ax.set_xlabel('Distance along Transect (m)')
    ax.set_ylabel('Depth below WSL (m)')
    ax.grid(linestyle=':')
    colormap = 'jet'
    vmin=0
    vmax=0.5
    #vmax=ceil(np.nanmax(V))
    cnt = ax.imshow(V.T, extent=[T1['dist'].min(), T1['dist'].max(), min(depths), max(depths)], 
                    aspect='auto', origin='lower', cmap=colormap, vmin=vmin, vmax=vmax)
    # cnt = ax.contourf(x,y,c, cmap='jet', levels=100, vmin=0, vmax=ceil(np.nanmax(c)))
    ax.set_xlim(0, ceil(np.nanmax(T1['dist'])/10)*10)
    ax.set_ylim(ceil(np.nanmax(depths)),0)
    cm = plt.cm.ScalarMappable(cmap=colormap)
    cm.set_array(V.T)
    cm.set_clim(vmin, vmax)
    cb = fig.colorbar(cm)
    cb.set_label('Velocity Magnitude (m/s)')
    cb.ax.tick_params(labelsize=8)
    cb.set_ticks(np.arange(0, ceil(np.nanmax(V))+0.1, 0.1))
    ax.set_title(f"Transect {transect_id} : {T1['date'].iloc[0]}")
    fig.savefig(f"transect_{transect_id}.png", dpi=400, bbox_to_inches='tight', pad_inches=0)
    plt.close(fig)