#!/usr/bin/env python

'''
Written by Madeleine Ekblom, 2023-04-06

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
    '''

    lon_min=19
    lon_max=32
    lat_min=59
    lat_max=71

    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]

    fig = plt.figure(figsize=(12,8))
    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)
    
    # basic plot, not changing colors of p:       
    #ax.scatter(p_values['lon'], p_values['lat'], c=p_values['p'])

    #loop over limits:
    legend_elements = []
    for n in range(8):
        p_ind = ((p_values['p'] > limits[n]) & (p_values['p'] <= limits[n+1]))
        ax.scatter(p_values['lon'][p_ind], p_values['lat'][p_ind], c=colors[n])
        legend_elements.append(Line2D([0], [0], marker='o', color='w',markerfacecolor=colors[n], label=str(limits[n])+'-'+str(limits[n+1])))

    # legend table with customized values

    #print(legend_elements)
    
    ax.legend(handles=legend_elements, loc='center left', bbox_to_anchor=(1.1, 0.5))
    ax.set_title('P-values for T2m quantiles in ' + timestring, fontsize=20)

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

def main(p_values_file):
    
    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)
