#!/usr/bin/env python

'''
Written by Madeleine Ekblom, 2023-04-06
updates 2023-11-27
- Option of moving legend into plot when plotting on a global map(lower left corner)

Based on Jouni Räisänen's script for plotting p values on the map of Finland

Plots p values on a map restricted to an area limited by lons (19, 32) and lats (59, 71)

Example:
$ python3 plot_p_values_map.py p_values_as_text_file
$ python3 plot_p_values_map.py p-values

p_values_as_text_file is a text file containing stationID, longitude, latitude, and p value with one header row
'''

import numpy as np
import matplotlib.pyplot as plt
import sys
import cartopy.crs as ccrs
from matplotlib.lines import Line2D

def read_data(p_values_file):
    '''
    Description: Reads in file containing p values at different stations
    Input: file containing p values with stationID (sid), longitude (lon), latitude (lat), and p value (p). Header = 1 row.
    Output: structure numpy array with sid, lon, lat, p
    '''

    p_values_data = np.loadtxt(p_values_file, skiprows=1, dtype={'names': ('sid', 'lon', 'lat', 'p'), 'formats':('i4', 'f4', 'f4', 'f4')})

    return p_values_data


def plot_p_values(p_values):
    '''
    Description: plot p values on a map restricted to an area limited by lons (19, 32) and lats (59, 71)
    Input: numpy array  
    Output: saves a figure in the running directory
    '''
    # add more variables and units when needed
    variable_shortnames={'39': '2t'}
    variable_units={'39': '$^\circ$C'}
    
    # variables to add, shortnames and units need to be checked
    # total amount of clouds, 91
    # sea level pressure, 108 
    # 1-hour precipitation, 80
    # relative humidity, 58
    # 10-minute precipitation intensity, 999
    # snow depth, 71 
    # 2m temp, 39 --> 2t, degrees Celcius
    # dew point temperature, 40 
    # visibility, 6
    # wind direction, 111 
    # max wind gust in last 10 mins,261
    # surface wind speed, 221
    
    variable='39' # change to be an input value (as string)
    
    # if data only from SYNOP stations in Finland:
    only_Finland_data=True
    
    # Define area of map:
    if only_Finland_data:
        # Finland:
        lon_min=19
        lon_max=32
        lat_min=59
        lat_max=71
    else:
        # Global:
        lon_min=0
        lon_max=359
        lat_min=-90
        lat_max=90

    # define colours and limits for the colours:
    colors = ['darkviolet', 'blue', 'skyblue', 'lime', 'yellow', 'orange', 'red', 'grey']
    limits = [0,0.01, 0.05, 0.1, 0.2, 0.3, 0.4, 0.5,1]

    # create figure and add a subplot with a map:
    fig = plt.figure(figsize=(8,12))
    ax = fig.add_subplot(1,1,1,projection=ccrs.PlateCarree())
    ax.coastlines()
    ax.set_extent([lon_min, lon_max, lat_min, lat_max], crs=ccrs.PlateCarree())
    ax.gridlines(draw_labels=True)
    
    # Plotting details:
    fs1=20 # fontsize for title
    fs2=14 # fontsize for legend
    fs3=10 # fontisze for text box explanation
    
    #loop over limits:
    legend_elements = []
    for n in range(8):
        # find indices 
        p_ind = ((p_values['p'] > limits[n]) & (p_values['p'] <= limits[n+1]))
        # scatter plot
        ax.scatter(p_values['lon'][p_ind], p_values['lat'][p_ind], c=colors[n])
        # add to legend_elements, color and values
        legend_elements.append(Line2D([0], [0], marker='o', color='w',markerfacecolor=colors[n], label=str(limits[n])+'-'+str(limits[n+1])))

    # explanation for figure:
    text_expl ='p-value: the probability that the MSD statistics for the quantile rank\nhistogram would be exceeded by pure chance\nwhen comparing a perfect model simulation with observations.'

    #  add legend table with customized values
    if only_Finland_data:
        # If Finland map, put legend outside map
        ax.legend(handles=legend_elements, loc='center left', fontsize=fs3) #, bbox_to_anchor=(1.1, 0.5), fontsize=fs2)
        # add text box with explanation to the bottom of the figure: 0.1, -0.25
        ax.text(0.1, -0.25, text_expl, bbox={'pad':10, 'facecolor':'white'}, fontsize=fs3, transform=ax.transAxes)
    else:
        # If global map, put Legend in the lower left corner
        ax.legend(handles=legend_elements, loc='lower left', fontsize=fs3)
        # add text box with explanation to the bottom of the figure, 0.1, -0.45
        ax.text(0.1, -0.45, text_expl, bbox={'pad':10, 'facecolor':'white'}, fontsize=fs3, transform=ax.transAxes)	
    	
    # set title:
    fig.suptitle('SYNOP stations: p-values for ' + variable_shortnames[variable] +'\nquantiles in ' + timestring, fontsize=fs1)

    # save figure
    plt.savefig('p_values.png', dpi=300)
    

def main(p_values_file):
    '''
    Main file: (1) reads in p values from file, (2) plot p values on a map
    '''
    
    p_values_data = read_data(p_values_file)

    plot_p_values(p_values_data)
    

if __name__=='__main__':
    p_values_file = sys.argv[1]
    timestring = sys.argv[2]
    
    main(p_values_file)
