#!/usr/bin/env python

'''
Written by Madeleine Ekblom, Oct 2023
Updates by M.E, 2023-11-29
- split title into title and textbox with additional information

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

Example:
$ python3 plot_rank_hist_sum_temp.py rh_summary_file nstat p01_00 p1_00 p5_00 p01_12 p1_12 p5_12 p_ylim q_ylim figname

'''

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

def plot_rank_hist(p_freq_00, p_freq_12, q_freq_00, q_freq_12, number_of_stations, p01_00, p1_00, p5_00, p01_12, p1_12, p5_12, p_ylim, q_ylim, figname):
    '''
    Input: p_freq 00/12, q_freq 00/12, number of stations, p01, p1, p5, (00/12), p_ylim, q_ylim, figname
    Description: Plots p and q frequencies and saves the figure with figname
    '''
    # title text
    title='RADSOUND: Rank histogram summary statistics\nfor $\Delta$T = 2t - t850'
    
    # can also be input arguments if the number of p and q bins are not constant
    number_of_p_bins=20
    number_of_q_bins=100

    # set up figure with 2x2 subplots
    ##fig = plt.figure(figsize=(9,12))
    fig = plt.figure(figsize=(10,12))
    spec = fig.add_gridspec(3,2, hspace=0.5)
    
    # plotting details:
    fs1 = 20 # fontsize for title
    fs2 = 14 # fontsize for labels, subplot titles
    fs3 = 10 # fontsize for figure explanations

    # Plot p values, 00 UTC
    ax1 = fig.add_subplot(spec[0,0])
    ax1.set_title('00 UTC\nNormalized p-value frequencies', fontsize=fs2)
    ax1.set_xlabel('p-value', fontsize=fs2)
    ax1.set_ylabel('Relative frequency', fontsize=fs2)
    ax1.set_ylim([0,p_ylim])
    ax1.bar(np.arange(number_of_p_bins)+0.5, number_of_p_bins*p_freq_00) # 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 p values, 12 UTC
    ax2 = fig.add_subplot(spec[0,1])
    ax2.set_title('12 UTC\nNormalized p-value frequencies', fontsize=fs2)
    ax2.set_xlabel('p-value', fontsize=fs2)
    ax2.set_ylabel('Relative frequency', fontsize=fs2)
    ax2.set_ylim([0,p_ylim])
    ax2.bar(np.arange(number_of_p_bins)+0.5, number_of_p_bins*p_freq_12) # add +0.5 to center the bars
    ax2.xaxis.set_major_locator(MultipleLocator(2))
    ax2.xaxis.set_minor_locator(MultipleLocator(1))
    p_xlabels = ax2.get_xticks()
    ax2.set_xticks(p_xlabels, p_xlabels/number_of_p_bins)
    ax2.set_xlim([0,number_of_p_bins])

    # Plot q values, 00 UTC
    ax3 = fig.add_subplot(spec[1,0])
    ax3.set_title('00 UTC\nNormalized quantile frequencies', fontsize=fs2)
    ax3.set_xlabel('Quantile (%)', fontsize=fs2)
    ax3.set_ylabel('Relative frequency', fontsize=fs2)
    ax3.bar(np.arange(number_of_q_bins)+0.5, number_of_q_bins*q_freq_00) # add +0.5 to center the bars
    ax3.xaxis.set_major_locator(MultipleLocator(10))
    ax3.xaxis.set_minor_locator(MultipleLocator(5))
    ax3.set_ylim([0,q_ylim])
    ax3.set_xlim([0,number_of_q_bins])
    ax3.axhline(y=1, c='k', ls='--') # horizontal black dashed line at y=0 

    # Plot q values, 12 UTC
    ax4 = fig.add_subplot(spec[1,1])
    ax4.set_title('12 UTC\nNormalized quantile frequencies', fontsize=fs2)
    ax4.set_xlabel('Quantile (%)', fontsize=fs2)
    ax4.set_ylabel('Relative frequency', fontsize=fs2)
    ax4.bar(np.arange(number_of_q_bins)+0.5, number_of_q_bins*q_freq_12) # add +0.5 to center the bars
    ax4.xaxis.set_major_locator(MultipleLocator(10))
    ax4.xaxis.set_minor_locator(MultipleLocator(5))
    ax4.set_ylim([0,q_ylim])
    ax4.set_xlim([0,number_of_q_bins])
    ax4.axhline(y=1, c='k', ls='--') # horizontal black dashed line at y=0 
    
    # add text box with additional information:
    ax5 = fig.add_subplot(spec[2,:])
    ax5.axis('off')
    
    # textbox text
    textbox = '                         Number of stations: ' + str(number_of_stations) + '\nFrequency, p<0.001, 00 UTC: ' + str(p01_00) + '     Frequency, p<0.001, 12 UTC: ' + str(p01_12) + '\nFrequency, p<0.01,   00 UTC: ' + str(p1_00) + '     Frequency, p<0.01,   12 UTC: ' + str(p1_12) + '\nFrequency, p<0.05,   00 UTC: ' + str(p5_00) + '     Frequency, p<0.05,   12 UTC: ' + str(p5_12)
    
    ##ax5.text(0.3, 0.6, textbox, bbox={'pad':10, 'facecolor':'white'}, fontsize=fs2)
    ax5.text(0.01, 0.6, textbox, 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
    ax5.text(0.20, -0.05, text_expl, bbox={'pad':10, 'facecolor':'white'}, fontsize=fs3)
    
    # add title to figure
    fig.suptitle(title, fontsize=fs1)
    
    plt.savefig(figname)

def read_data(rh_summary_file):
    '''
    Input: path/to/rh_summary_file
    Output: p_freq_00, p_freq_12, q_freq_00, q_freq_12
    Description: reads in p frequencies and q frequencies
    '''
    #print(rh_summary_file)

    # 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, 12, 24
    p_freq_00 = p_freq_all[0,1:]
    p_freq_12 = p_freq_all[1,1:]
    q_freq_00 = q_freq_all[0,1:]
    q_freq_12 = q_freq_all[1,1:]

    return p_freq_00, p_freq_12, q_freq_00, q_freq_12

if __name__=='__main__':
    rh_summary_file = sys.argv[1]
    number_of_stations = sys.argv[2]
    p01_00 = sys.argv[3]
    p1_00 = sys.argv[4]
    p5_00 = sys.argv[5]
    p01_12 = sys.argv[6]
    p1_12 = sys.argv[7]
    p5_12 = sys.argv[8]
    p_ylim = float(sys.argv[9])
    q_ylim = float(sys.argv[10])
    figname = sys.argv[11]

    p_freq_00, p_freq_12, q_freq_00, q_freq_12 = read_data(rh_summary_file)
    plot_rank_hist(p_freq_00, p_freq_12, q_freq_00, q_freq_12, number_of_stations, p01_00, p1_00, p5_00, p01_12, p1_12, p5_12, p_ylim, q_ylim, figname)

