#!/usr/bin/env python

'''
Written by Madeleine Ekblom, 2023-09-12

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, figname):
    '''
    Input: p_freq, q_freq, number of stations, p01, p1, p5, max freq p, max freq q, figname 
    Description: Plots p and q frequencies and saves the figure with figname
    '''
    title='Rank histogram summary statistics\n 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)

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

    fig = plt.figure(figsize=(8,12))
    gs = fig.add_gridspec(2, hspace=0.4)
    ax1, ax2 = gs.subplots(sharex=False, sharey=False)

    # Plot p values
    ax1.set_title('Normalized p-value frequencies')
    ax1.set_xlabel('p-value')
    ax1.set_ylim([0,max_freq_p])
    ax1.bar(np.arange(number_of_p_bins)+0.5, number_of_p_bins*p_freq)
    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')
    ax2.set_xlabel('Quantile (%)')
    ax2.bar(np.arange(number_of_q_bins)+0.5, number_of_q_bins*q_freq)
    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])

    fig.suptitle(title)
    plt.savefig(figname)

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])
    figname = sys.argv[8]

    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, figname)

