#!/usr/bin/env python

'''
Written by Madeleine Ekblom, 2023-09-12
updates 2023-11-27:
- moved text from title to a new text box

Based on Jouni Räisänen's script for plotting rank histogram summart statistics 

Example:
$ python3 plot_rank_hist_sum_all_stations.py rh_summary_file nstat p01 p1 p5 max_freq_p max_freq_q

'''

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import (MultipleLocator, AutoMinorLocator)
import sys

def plot_rank_hist(p_freq, q_freq, number_of_stations, p01, p1, p5, max_freq_p, max_freq_q):
    '''
    Input: p_freq, q_freq, number of stations, p01, p1, p5, max freq p, max freq q
    Description: Plots p and q frequencies and saves the figure with figname
    '''
    # 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 
    
    # title
    title='SYNOP: Rank histogram summary statistics for '  + variable_shortnames[variable]
    
    # text box information
    text_box = 'Number of stations: ' + str(number_of_stations) + '\nFrequency, p<0.001: ' + str(p01) + '\nFrequency, p<0.01: ' + str(p1) + '\nFrequency, p<0.05: ' + str(p5)

    # these can also be input arguments if the number of bins are not constant
    number_of_p_bins=20
    number_of_q_bins=100

    # set up figure
    fig = plt.figure(figsize=(8,12))
    gs = fig.add_gridspec(3, hspace=0.5)
    ax1, ax2, ax3 = gs.subplots(sharex=False, sharey=False)
    
    # plotting details
    fs1=20 # fontsize for title
    fs2=14 # fontsize for labels, subplot titles, textbox
    fs3=10 # fontsize for text explanation

    # Plot p values
    ax1.set_title('Normalized p-value frequencies', fontsize=fs2)
    ax1.set_xlabel('p-value', fontsize=fs2)
    ax1.set_ylabel('Relative frequency', fontsize=fs2)
    ax1.set_ylim([0,max_freq_p])
    ax1.bar(np.arange(number_of_p_bins)+0.5, number_of_p_bins*p_freq) # add +0.5 to center the bars
    ax1.xaxis.set_major_locator(MultipleLocator(2))
    ax1.xaxis.set_minor_locator(MultipleLocator(1))
    p_xlabels = ax1.get_xticks()
    ax1.set_xticks(p_xlabels, p_xlabels/number_of_p_bins)
    ax1.set_xlim([0,number_of_p_bins])

    # Plot q values
    ax2.set_title('Normalized quantile frequencies', fontsize=fs2)
    ax2.set_xlabel('Quantile (%)', fontsize=fs2)
    ax2.set_ylabel('Relative frequency', fontsize=fs2)
    ax2.bar(np.arange(number_of_q_bins)+0.5, number_of_q_bins*q_freq) # add +0.5 to center the bars
    ax2.xaxis.set_major_locator(MultipleLocator(10))
    ax2.xaxis.set_minor_locator(MultipleLocator(5))
    ax2.set_ylim([0,max_freq_q])
    ax2.set_xlim([0,number_of_q_bins])

    # remove axis from ax3
    ax3.axis('off')
    
    # add text box with additional infromation:
    ax3.text(0.3, 0.6, text_box, bbox={'pad':10, 'facecolor':'white'}, fontsize=fs2)
    
    # explanation for figure
    text_pvalue= 'Normalized p-values: Frequency distribution of p-values at\nall individual stations, normalized to give an expected value of 1.\n'
    text_quantile='Normalized quantiles: Frequency distribution of model data\nrelative to quantiles of observations, averaged over all stations\nand divided by the expected value of 0.01.'
    
    text_expl = text_pvalue + text_quantile
    
    # add text box with explanation to the bottom of the figure
    ax3.text(0.1, -0.05, text_expl, bbox={'pad':10, 'facecolor':'white'}, fontsize=fs3)

    # add title to figure
    fig.suptitle(title, fontsize=fs1)
    
    # save figure
    figname='rank_hist_sumstats.png'
    plt.savefig(figname, dpi=300)

def read_data(rh_summary_file):
    '''
    Input: path/to/rh_summary_file
    Output: p_freq, q_freq
    Description: reads in p frequencies and q frequencies
    '''

    # first line contains a header and first column contains hour
    p_freq_all = np.loadtxt(rh_summary_file + '_p-freq.txt', skiprows=1)
    q_freq_all = np.loadtxt(rh_summary_file + '_q-freq.txt', skiprows=1)

    # p_freq and q_freq contain data 00, 06, 12, 18, 24 (the last is used in these plots)
    p_freq = p_freq_all[-1,1:]
    q_freq = q_freq_all[-1,1:]

    return p_freq, q_freq

if __name__=='__main__':
    rh_summary_file = sys.argv[1]
    number_of_stations = sys.argv[2]
    p01 = sys.argv[3]
    p1 = sys.argv[4]
    p5 = sys.argv[5]
    max_freq_p = float(sys.argv[6])
    max_freq_q = float(sys.argv[7])

    p_freq, q_freq = read_data(rh_summary_file)
    plot_rank_hist(p_freq, q_freq, number_of_stations, p01, p1, p5, max_freq_p, max_freq_q)

