#!/usr/bin/env python3
#  -*- coding: utf-8 -*-

"""
Plot shipping density in 2019 in the Arabian Gulf

Author: Thomas Anselain, Earth and Life Institute, UCLouvain, Belgium
Last modified: 12 April 2022
"""
#%% Import packages and set up
import matplotlib.pyplot as plt
import geopandas as gpd
#from cartopy import crs as ccrs
import pathlib
import numpy as np
import matplotlib
from shapely.geometry import Polygon
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib import colors
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
import matplotlib.patches as patches
from matplotlib.colors import LinearSegmentedColormap, BoundaryNorm

basedir = "/export/miro/students/tanselain/OpenOil/"


#%% Define colormap
def truncate_colormap(cmap, minval=0.0, maxval=1.0, n=100):
    new_cmap = colors.LinearSegmentedColormap.from_list(
        'trunc({n},{a:.2f},{b:.2f})'.format(n=cmap.name, a=minval, b=maxval),
        cmap(np.linspace(minval, maxval, n)))
    return new_cmap


#%% Define main function
def plot_intersect_ships(vmin, vmax,outfile) :
    
    # Call files
    path = basedir + f'Indicators/Oil_shipping_exposure_map/Density_shipping_2019/shp_data/Persian_shipping_density_2019'
    df2 = gpd.read_file(path)
    df2 = df2.to_crs("EPSG:4326")
    total = df2['density'].sum() #sum of density to determine number of ships in along month

        
#%%Plot
    # Main axis
    fig = plt.figure(figsize = (12,12))     
    ax = fig.add_subplot(111)
    
    # Second axis to draw cbbox
    fc = colors.to_rgba('white')
    ec = colors.to_rgba('gray')
    fc = fc[:-1] + (0,)
    cbbox = inset_axes(ax, '35%', '15%',loc=3, bbox_to_anchor=(0.02, 0.01, 1, 1), bbox_transform=ax.transAxes, borderpad=0)  #30% et 15%
    cbbox.tick_params(left='off', top='off', right='off', bottom='off', labelleft='off', labeltop='off', labelright='off', labelbottom='off')
    cbbox.spines['bottom'].set_visible(False)
    cbbox.spines['top'].set_visible(False) 
    cbbox.spines['right'].set_visible(False)
    cbbox.spines['left'].set_visible(False)
    cbbox.set_facecolor(fc)
    
    # Subaxis to add colorbar to cbbox
    subax = inset_axes(cbbox,
                    width="100%", # width = 30% of parent_bbox   de base 100%
                    height="36%", # height : 1 inch     de base 36%
                    loc = 10,
                    bbox_to_anchor=(0,0,1,0.7),
                    bbox_transform=cbbox.transAxes,
                    borderpad=0,
                    )
    
    # Fig to draw zoom on Ras Laffan
    rs_ax = inset_axes(ax, '30%', '30%',loc=1, bbox_to_anchor=(0, 0, 1.4, 1), bbox_transform=ax.transAxes, borderpad=0) #35 et 35, 1.6 1.07
    ec2 = colors.to_rgba('darkblue')
    for axis in ['top','bottom','left','right']:
        rs_ax.spines[axis].set_linewidth(2)
        rs_ax.spines[axis].set_color(ec2)
    rs_ax.tick_params(axis='both', left='off', top='off', right='off', bottom='off', labelleft='off', labeltop='off', labelright='off', labelbottom='off')
    
    # Fig to draw zoom on Abu Fontas
    af_ax = inset_axes(ax, '30%', '46%',loc=1, bbox_to_anchor=(0, 0, 1.4, 0.65), bbox_transform=ax.transAxes, borderpad=0) # 30 et 41, 1.5 0.71       44% 0.68
    ec3 = colors.to_rgba('red')
    for axis in ['top','bottom','left','right']:
        af_ax.spines[axis].set_linewidth(2)
        af_ax.spines[axis].set_color(ec3)
    af_ax.tick_params(axis='both', left='off', top='off', right='off', bottom='off', labelleft='off', labeltop='off', labelright='off', labelbottom='off')
    
    # Fig to draw zoom on Umm al Houl
    uh_ax = inset_axes(ax, '30%', '100%',loc=1, bbox_to_anchor=(0, 0, 1.4, 0.3), bbox_transform=ax.transAxes, borderpad=0) # 30 et 87, 1.6 0.38
    ec4 = colors.to_rgba('green')
    for axis in ['top','bottom','left','right']:
        uh_ax.spines[axis].set_linewidth(2)
        uh_ax.spines[axis].set_color(ec4)
    uh_ax.tick_params(axis='both', left='off', top='off', right='off', bottom='off', labelleft='off', labeltop='off', labelright='off', labelbottom='off')
    
    # Add total ships 
    #ax.text(47.85, 24.25, f'Sum of ships: {total}', fontsize=18, horizontalalignment = 'left', verticalalignment = 'center', zorder=6)
    
    # Put land, border and eez in background
    bgdir = basedir + "Indicators/Background_map/"
    gpd_landp  = gpd.read_file(bgdir+"PG_layer_poly_new_2/PG_layer_poly_new_2.shp")
    gpd_bordl  = gpd.read_file(bgdir+"natural_earth/ne_10m_admin_0_boundary_lines_land.shp")
    gpd_eezl  = gpd.read_file(bgdir+"qatar_eez/qatar_eez_line.shp")
    
    gpd_landp.plot(ax=ax, linewidth=0.5, edgecolors="k", color='lightgray',zorder=4)
    gpd_bordl.plot(ax=ax, linewidth=0.5, linestyle='-', color="k", zorder=4)
    gpd_eezl.plot(ax=ax, linewidth=2, color="darkgreen",zorder=4)
    
    gpd_landp.plot(ax=rs_ax, linewidth=0.5, edgecolors="k", color='lightgray',zorder=0.9)
    gpd_bordl.plot(ax=rs_ax, linewidth=0.5, linestyle='-', color="k", zorder=0.9)
    gpd_eezl.plot(ax=rs_ax, linewidth=2, color="darkgreen",zorder=0.9)
    rs_ax.set_xlim(51.5,51.7)
    rs_ax.set_ylim(25.83,26.03)
    rect1 = patches.Rectangle((51.5, 25.83), (51.7 - 51.5),(26.03 - 25.83), linewidth=2, edgecolor=ec2, facecolor = (1,0,0,0), zorder= 5) # Rectangle on main figure
    ax.add_patch(rect1)
    
    gpd_landp.plot(ax=af_ax, linewidth=0.5, edgecolors="k", color='lightgray',zorder=0.9)
    gpd_bordl.plot(ax=af_ax, linewidth=0.5, linestyle='-', color="k", zorder=0.9)
    gpd_eezl.plot(ax=af_ax, linewidth=2, color="darkgreen",zorder=0.9)
    af_ax.set_xlim(51.61,51.645)
    af_ax.set_ylim(25.19,25.225)
    rect2 = patches.Rectangle((51.61, 25.19), (51.645 - 51.61), (25.225 - 25.19), linewidth=2, edgecolor=ec3, facecolor = (1,0,0,0), zorder= 5) # Rectangle on main figure
    ax.add_patch(rect2)
    
    gpd_landp.plot(ax=uh_ax, linewidth=0.5, edgecolors="k", color='lightgray',zorder=0.9)
    gpd_bordl.plot(ax=uh_ax, linewidth=0.5, linestyle='-', color="k", zorder=0.9)
    gpd_eezl.plot(ax=uh_ax, linewidth=2, color="darkgreen",zorder=0.9)
    uh_ax.set_xlim(51.6,51.65)
    uh_ax.set_ylim(25.085,25.135)
    rect3 = patches.Rectangle((51.6, 25.085), (51.65 - 51.6), (25.135 - 25.085), linewidth=2, edgecolor=ec4, facecolor = (1,0,0,0), zorder= 5) # Rectangle on main figure
    ax.add_patch(rect3)
    
    # Plot density mesh
    cmap = plt.get_cmap('YlOrRd')
    new_cmap = truncate_colormap(cmap, 0, 1) #set colormap
    
    df2I = df2.loc[df2['density'] >= vmin]
    column = df2I['density'].values.astype(float)
    vmax2 = max(column)
    print(vmax2)
    
    bounds = np.array([0, 10 , 100, 1000, 10000, 1000000])
    stretched_bounds = np.interp(np.linspace(0, 1, 257), np.linspace(0, 1, len(bounds)), bounds)
    norm = matplotlib.colors.BoundaryNorm(stretched_bounds, ncolors=256)
    
    pcm = df2I.plot(column='density', ax=ax, cmap= new_cmap, norm=norm, vmin = vmin, vmax = vmax2)
    
    # Set axis
    #ax.set_xlim(47.7, 56.4)
    #ax.set_ylim(23.8, 30.3)
    ax.set_xlim(50, 53.2)
    ax.set_ylim(24, 27.5)
    #yticks = np.arange(24, 30.5, .5)
    yticks = np.arange(24, 28, .5)
    ax.set_yticks(yticks)
    ax.set_yticklabels([f"{y}°N" for y in yticks])
    ax.tick_params(labelleft=True, labelbottom = True) 
    #xticks = np.arange(48,56.5,0.5)
    xticks = np.arange(50,53.5,0.5)
    ax.set_xticks(xticks)
    ax.set_xticklabels([f"{x}°E" for x in xticks])
    
    # Set colorbar   
    cticks = np.array([0, 10 , 100, 1000, 10000, 1000000]) 
    cbar = matplotlib.colorbar.ColorbarBase(subax, orientation = 'horizontal', cmap=new_cmap, norm=norm, ticks=[cticks]) 
    cbar.ax.set_xticklabels(['0', '10','10$^2$', '10$^3$', '10$^4$', '$\geq$10$^6$'])
    cbar.set_label(label = 'Shipping density', labelpad=-70, y=1, rotation=0, fontsize = 15)
    cbbox.text(0.5, 0.55, '(routes/km$^2$ year)', horizontalalignment = 'center', verticalalignment = 'center', rotation = 0, fontsize = 10)
    
    
    # Add number for each dangerous shipping routes
    point1 = ax.text(50.9, 26.49, '1', color = 'royalblue', fontsize = 17, fontweight='bold')
    point2 = ax.text(51.24, 26.34, '2', color = 'royalblue', fontsize = 17, fontweight='bold')
    point3 = ax.text(51.71, 26.15, '3', color = 'royalblue', fontsize = 17, fontweight='bold')
    point4 = ax.text(52, 26.06, '4', color = 'royalblue', fontsize = 17, fontweight='bold')
    point5 = ax.text(51.75, 25.12, '5', color = 'royalblue', fontsize = 17, fontweight='bold')
    point6 = ax.text(50.57, 26.88, '6', color = 'royalblue', fontsize = 17, fontweight='bold')
    point7 = ax.text(52.15, 26.37, '7', color = 'royalblue', fontsize = 17, fontweight='bold')
    point8 = ax.text(52.62, 25.5, '8', color = 'royalblue', fontsize = 17, fontweight='bold')
    
    # Add point and text on subplot + Legend
    p1 = rs_ax.scatter(51.55, 25.9464, color ='#1f77b4')
    p2 = rs_ax.scatter(51.6682, 25.9204, color ='limegreen', marker = '*', s = 80)
    rs_ax.text((51.5 + (0.5*(51.7-51.5))), (25.83 + (0.9*(26.03-25.83))), 'Ras Laffan', fontsize = 14, fontweight='bold', zorder= 5)
    af_ax.scatter(51.62558, 25.2040, color = '#1f77b4')
    af_ax.text((51.61 + (0.5*(51.645-51.61))), (25.19 + (0.9*(25.225-25.19))), 'Abu Fontas', fontsize = 14, fontweight='bold',zorder=5)
    uh_ax.scatter(51.633, 25.1135, color = '#1f77b4')
    uh_ax.text((51.6 + (0.45*(51.65-51.6))), (25.085 + (0.9*(25.135-25.085))), 'Umm al Houl', fontsize = 14, fontweight='bold', zorder= 5)
    
    leg = rs_ax.legend([p1, p2], ['Desalination plant intake','LNG terminal entry'], bbox_to_anchor=(-0.18, -2.65), markerscale = 2, loc= 'lower left', frameon = False, fontsize = 14)
    #leg.legendHandles[0]._legmarker.set_markersize(6)
    #leg.legendHandles[1]._legmarker.set_markersize(6)
    
    # Zoom letter
    ax.text(51.46, 25.83, 'A', color = 'steelblue', fontsize = 16, fontweight='bold', horizontalalignment = 'center', verticalalignment = 'center', zorder = 5)
    ax.text(51.57,25.19, 'B', color = 'steelblue', fontsize = 16, fontweight='bold', horizontalalignment = 'center', verticalalignment = 'center', zorder = 5)
    ax.text(51.56, 25.085,'C', color = 'steelblue', fontsize = 16, fontweight='bold', horizontalalignment = 'center', verticalalignment = 'center', zorder = 5)
    rs_ax.text(51.52, 26.01, 'A', color = 'steelblue', fontsize = 16, fontweight='bold', horizontalalignment = 'center', verticalalignment = 'center', zorder = 5)
    af_ax.text(51.6135, 25.2215, 'B', color = 'steelblue', fontsize = 16, fontweight='bold', horizontalalignment = 'center', verticalalignment = 'center', zorder = 5)
    uh_ax.text(51.605, 25.13, 'C', color = 'steelblue', fontsize = 16, fontweight='bold', horizontalalignment = 'center', verticalalignment = 'center', zorder = 5)
   
    # Add cities
    ax.scatter(51.5310-0.01, 25.2854, marker = '*', color = 'black', s=80, zorder=5)#, fontsize = 14, horizontalalignment = 'center', verticalalignment = 'center', zorder = 5)
    ax.scatter(51.5493+0.01, 24.9909, marker = '*', color = 'black', s=80, zorder=5)#, fontsize = 14, horizontalalignment = 'center', verticalalignment = 'center', zorder = 5)
    ax.text(51.5310-0.040, 25.2854, 'Doha', color = 'black', fontsize = 14, horizontalalignment = 'right', verticalalignment = 'center', zorder = 5)
    ax.text(51.5493-0.020, 24.9909, 'Mesaaied', color = 'black', fontsize = 14, horizontalalignment = 'right', verticalalignment = 'center', zorder = 5)

    # Draw scales
    ax.hlines(24.03,52.84,53.05,lw=5,color='k', zorder=5) 
    ax.text(52.945, 24.085, '21 km', color = 'k', fontsize = 12, horizontalalignment = 'center', verticalalignment = 'center', zorder = 5)
    
    rs_ax.hlines(25.84,51.65,51.69,lw=5,color='k')
    rs_ax.text(51.67, 25.85, '4 km', color = 'k', fontsize = 12, horizontalalignment = 'center', verticalalignment = 'center')
    
    af_ax.hlines(25.19175, 51.63625, 51.64325,lw=5,color='k')
    af_ax.text(51.63975, 25.1935, '700 m', color = 'k', fontsize = 12, horizontalalignment = 'center', verticalalignment = 'center')
    
    uh_ax.hlines(25.0875,51.6375,51.6475,lw=5,color='k')
    uh_ax.text(51.6425, 25.09, '1 km', color = 'k', fontsize = 12, horizontalalignment = 'center', verticalalignment = 'center')
    
    
    rs_ax.set_aspect("equal")
    af_ax.set_aspect("equal")
    uh_ax.set_aspect("equal")
    ax.set_aspect("equal")

    rs_ax.axes.xaxis.set_ticks([]); rs_ax.axes.yaxis.set_ticks([]);
    af_ax.axes.xaxis.set_ticks([]); af_ax.axes.yaxis.set_ticks([]);
    uh_ax.axes.xaxis.set_ticks([]); uh_ax.axes.yaxis.set_ticks([]);
    cbbox.axes.xaxis.set_ticks([]); cbbox.axes.yaxis.set_ticks([]);

#%% Final save
    save = outfile + f"all_inputs_qatar.png"
    plt.savefig(save, dpi=300, bbox_inches="tight")
    plt.close(fig)
   

#%% Beginning of loops    
# Give the list of month assessed      
list_mois = ['january', 'february', 'march', 'april', 'may', 'june', 'july', 'august', 'september', 'october', 'november', 'december']

# Start loops for main function
fname = basedir + f"Indicators/Oil_shipping_exposure_map/Density_shipping_2019/plot_csv_data/all_inputs/"
plot_intersect_ships(0, 2000, fname)