#!/usr/bin/env python

'''
Written by Jouni Räisänen, 2023-10-07
Updates by Madeleine Ekblom, 2023-11-29
- move legend into lower left corner of plot
- figure title added, subplot titles changed

Plots p values from quantile bin rank histogram analysis to global maps (scale 0-1, only low values indicate significance). 
P-values are plotted separately for 00 UTC and 12 UTC, to take into account varying data availability.

Example:
$ python3 plot_p_values_map_rank_hist_temp.py p_values_as_text_file_00 p_values_as_text_file_12 timestring figname

p_values_as_text_file_XX are text files containing stationID, longitude, latitude, and p value with one header row
timestring indicates the simulation period used. It is only used in the map headers.
'''

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':('S3', 'f4', 'f4', 'f4')})

    return p_values_data


def plot_p_values_00_12(p_values_00, p_values_12):
    '''
    Description: plot p values at 00 and 12 UTC on global maps
    Input: two numpy arrays  
    Output: saves a figure in the running directory
    '''

    # define lat and lon for global map
    lon_min=-180
    lon_max=180
    lat_min=-90
    lat_max=90

    # define colour and limits for the dots
    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]

    # set up figure
    fig = plt.figure(figsize=(12,11))
    
    # plotting details:
    fs1=20 # fontsize for title
    fs2=12 # fontsize for labels, subplot titles
    fs3=10 # fontsize for figure explanation
    
    ##  00 UTC ##
    # add subplot with map, for 00 UTC data
    ax = fig.add_subplot(2,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)
    
    #loop over limits:
    legend_elements = []
    for n in range(8):
        # find indices for p-values
        p_mod = 0.5+1.0001*(p_values_00['p']-0.5)
        p_ind = ((p_mod > limits[n]) & (p_mod <= limits[n+1]))
        # scatter plot with dots coloured based on p-values
        ax.scatter(p_values_00['lon'][p_ind], p_values_00['lat'][p_ind], c=colors[n])
        # add to legend_elements: colour and limit
        legend_elements.append(Line2D([0], [0], marker='o', color='w',markerfacecolor=colors[n], label=str(limits[n])+'-'+str(limits[n+1])))

    # add legend to lower left corner of map
    ax.legend(handles=legend_elements, loc='lower left')

    # add subplot title
    ax.set_title('00 UTC', fontsize=fs2, fontweight='bold')

    ##   12 UTC  ##
    # add subplot with map, for 12 UTC data
    ax = fig.add_subplot(2,1,2,projection=ccrs.PlateCarree())
    ax.coastlines()
    ax.set_extent([lon_min, lon_max, lat_min, lat_max], crs=ccrs.PlateCarree())
    ax.gridlines(draw_labels=True)
    
    #loop over limits:
    legend_elements = []
    for n in range(8):
        # find indices for p-values
        p_mod = 0.5+1.0001*(p_values_12['p']-0.5)
        p_ind = ((p_mod > limits[n]) & (p_mod <= limits[n+1]))
        # scatter plot with dots coloured based on p-values
        ax.scatter(p_values_12['lon'][p_ind], p_values_12['lat'][p_ind], c=colors[n])
        # add to legend elements: colour and limit
        legend_elements.append(Line2D([0], [0], marker='o', color='w',markerfacecolor=colors[n], label=str(limits[n])+'-'+str(limits[n+1])))
        
    # add legend to lower left corner of map
    ax.legend(handles=legend_elements, loc='lower left')
    
    # add subplot title
    ax.set_title('12 UTC', fontsize=fs2, fontweight='bold')
    
    # explanation for figure:
    text_expl ='p-value: the probability that the MSD statistics for the quantile rank \nhistogram would be exceeded by pure chance when comparing \na perfect model simulation with observations.'
    
    # add text box with explanation to the bottom of the figure
    ax.text(0.2, -0.25, text_expl, bbox={'pad':10, 'facecolor':'white'}, fontsize=fs3, transform=ax.transAxes)
    
    # add figure title:
    fig.suptitle('RADSOUND stations: p-values for $\Delta$T = 2t - t850\nquantiles in ' + timestring, fontsize=fs1) #, fontweight='bold')
    
    # save figure
    plt.savefig(figname, dpi=300)
    

def main(p_values_file_00, p_values_file_12):
    '''
    Main: (1) reads in p values at 00 UTC and 12 UTC (2) plots p-values on a global map 
    '''
    
    p_values_data_00 = read_data(p_values_file_00)
    p_values_data_12 = read_data(p_values_file_12)

    plot_p_values_00_12(p_values_data_00, p_values_data_12)    

if __name__=='__main__':
    p_values_file_00 = sys.argv[1]
    p_values_file_12 = sys.argv[2]
    timestring = sys.argv[3]
    figname = sys.argv[4]    
    
    main(p_values_file_00, p_values_file_12)
