#!/usr/bin/env python

'''
Written by Madeleine Ekblom, 2023-05-10
Updates 2023-12-01:
- visual changes of the plot: labels, legends, title, and colours updated
- text boxes added with rank histogram MSD and p-values 
- added: check if data available, if not add text box "No data available"

Based on Jouni Räisänen's script for plotting quantiles/time series data for SYNOP data, modified to plot only 00, 12 UTC data and rank histogram at 00, 12 UTC

Example: 
$ python3 plot_quantiles_rankhist.py example.cfg

example.cfg is a text file containing:
[data]
variable=39
station_id=102019
longitude=23.576000
latitude=68.602997
year_beg=2010
year_end=2012
month_beg=01
month_end=12

[input]
sim_data=time_series_102019_2010-2012.txt
quantile_sel=quantiles_102019.txt
rank_hists=rank_histogram_102019_2010-2012.txt

[output]
fig_name=standard_plot_102019_2010-2012_python.png

'''

import numpy as np
import matplotlib.pyplot as plt
import sys
import configparser
import datetime
from matplotlib.dates import drange, MonthLocator, DateFormatter
import matplotlib.lines as mlines
import matplotlib.patches as mpatches
import matplotlib.ticker as ticker

config = configparser.ConfigParser()

def plot_data(quantiles_data, rank_data, time_series_data, station_id, lon, lat, year_beg, year_end, month_beg, month_end, fig_name, savefig=False):
    '''
    Plot data at 00UTC and 12UTC: quantiles data give the lines and the time series data give the dots.
    '''
    # set up date array consisting of all days in beg year:
    days = np.arange(1,366)
    #    day1 = datetime.date(year_beg, 1, 1)
    #    day2 = datetime.date(year_beg+1, 1, 1)
    day1 = datetime.date(2001, 1, 1)
    day2 = datetime.date(2002, 1, 1)
    delta = datetime.timedelta(days=1)
    dates = drange(day1, day2, delta)
    
    
    
    # experiment id: now use 'test data' - to be replaced with real exp_id (as input argument)
    exp_id='test data'

    # arrays for quantile names, axes titles, hours to plot
    q_names = ['q01', 'q10', 'q25', 'q50', 'q75', 'q90', 'q99']
    sub_titles = ['00 UTC',  '12 UTC']
    hours = [0, 12]

    # calculate number of years:
    nyears = year_end - year_beg + 1

    # set up figure with 4x2 subplots -> last two rows of subplots for legend and textboxes
    fig = plt.figure(figsize=(10,10), layout='constrained')
    spec = fig.add_gridspec(4,2) 
    
    # plotting details:
    alpha= [0.5, 0.3, 0.1] # transparency for shaded area plotting
    color='blue' # color for shading
    linecolor='red' # color for solid line
    scatter_size=5 # size for scattered dots
    fs1=20 # fontsize for title
    fs2=14 # fontsize for labels, subplot titles
    fs3=12 # fontsize for legend
    fs4=10 # fontize for figure textbox
    ndigits=3 # round off to ndigits (e.g., lon, lat)          
                                             
    # legend elemnents for quantile data, one legend for model and one legend for observations
    legend_elements1 = [mlines.Line2D([],[], marker='.', linestyle='', color='k', label='Model data (' + exp_id + ')')]
    legend_elements2= [mpatches.Patch(color=color, alpha=alpha[0], label='25-50%'),
                                               mpatches.Patch(color=color, alpha=alpha[1], label='10-25%'),
                                               mpatches.Patch(color=color, alpha=alpha[2], label='1-10%'),
                                               mlines.Line2D([],[], color=linecolor, label='50%'),
                                               mpatches.Patch(color=color, alpha=alpha[0], label='50-75%'),
                                               mpatches.Patch(color=color, alpha=alpha[1], label='75-90%'),
                                               mpatches.Patch(color=color, alpha=alpha[2], label='90-99%')]
                                               
    # text box for empty data:
    text_no_data = "No data available" 
    
    # round lon, lat to three digits:
    lon = np.round(float(lon), ndigits)
    lat = np.round(float(lat), ndigits)
    
    # define lon/lat text strings for title of the form lat: XX N/S, lat: XX E/W
    if (lat > 0):
        # if lat positive -> N
        lat_text = ' (lat: ' + str(lat) + '$^\circ$N, '
    else:
        # if negative -> S
        lat_text = ' (lat ' + str(-lat) + '$^\circ$S, '
    
    if (0 <= lon) and (lon <= 180):
        # lon btw 0 an 180 -> E
        lon_text = ' lon: ' + str(lon) + '$^\circ$E)'
    else:
        if lon < 0:
            # lon < 0 --> - lon W
            lon_text = ' lon: ' + str(-lon) + '$^\circ$W)'
        else:
            # 180 < lon < 360 --> - (lon - 360) W
            lon = -(lon - 360)
            lon_text = ' lon: ' + str(lon) + '$^\circ$W)'
            
    # define figure title:
    fig_title = 'RADSOUND station: ' + station_id + ' ' + lat_text + lon_text + '\nModel data from ' + str(year_beg) + month_beg + '-' + str(year_end) + month_end
    fig.suptitle(fig_title, fontsize=20)

    # counter
    c = 0

    # plot quantiles/time series for times 00, 12. Loop from 0-1 to plot first row of subplots:
    for i in range(2):
        # set up axis: title, xaxis, x- and y-labels
        ax = fig.add_subplot(spec[0,i])
        ax.set_title(sub_titles[c], fontsize=fs2)
        ax.set_xlim(dates[0], dates[-1])
        ax.xaxis.set_major_locator(MonthLocator())
        ax.xaxis.set_major_formatter(DateFormatter('%b'))
        ax.set_xlabel('Time', fontsize=fs2)
        ax.set_ylabel('$\Delta$T = 2t - t850 ($^\circ$C)', fontsize=fs2)
    
        # quantile data:
        # find quantile data hour = 0, 12
        qh_ind = (quantiles_data['hour'][:] == hours[c])
            
        # get quantile data, qXX_data = quantile XX data at hour qh_ind
        q01_data = quantiles_data[q_names[0]][qh_ind]
        q10_data = quantiles_data[q_names[1]][qh_ind]
        q25_data = quantiles_data[q_names[2]][qh_ind]
        q50_data = quantiles_data[q_names[3]][qh_ind]
        q75_data = quantiles_data[q_names[4]][qh_ind]
        q90_data = quantiles_data[q_names[5]][qh_ind]
        q99_data = quantiles_data[q_names[6]][qh_ind]
            
        # solid line at q50:
        ax.plot(dates, q50_data, c=linecolor, ls='-')
            
        # shading will be the most opaque close to 50 and more transparent close to 1/99
          
        # fill area btw q25-q50, q50-q75: alpha = alpha[0]
        ax.fill_between(dates, q25_data[:], q50_data[:], color=color, alpha=alpha[0])
        ax.fill_between(dates, q50_data[:], q75_data[:], color=color, alpha=alpha[0])
            
        # fill area btw q10-q25, q75-q90
        ax.fill_between(dates, q10_data[:], q25_data[:], color=color, alpha=alpha[1])
        ax.fill_between(dates, q75_data[:], q90_data[:], color=color, alpha=alpha[1])

        # fill area btw q01-q10, q90-q99            
        ax.fill_between(dates, q01_data[:], q10_data[:], color=color, alpha=alpha[2])
        ax.fill_between(dates, q90_data[:], q99_data[:], color=color, alpha=alpha[2]) 
                
        # plot time series data:
        # find time series where hour = 0, 12
        th_ind = (time_series_data['hour'][:] == hours[c])
        t_data = time_series_data['value'][th_ind] # all years
        
        # check if data available for given station and time: check that the length of array t_data > 0
        if (len(t_data) > 0):
            for n in range(nyears):
                # plot data with black, small dots:
                ax.scatter(dates, t_data[n*365:(n+1)*365], s=scatter_size, c='k', marker='.')
        else:
            ax.text(0.5, 0.5, text_box, bbox={'pad':10, 'facecolor':'white'}, fontsize=fs3)
            
        # increase counter by 1   
        c = c + 1

    # plot rank histogram data, loop from 0-1 to plot last row of subplots:
    # rank_data[0,6:] contains 00UTC and rank_data[2,6:] contains 12UTC data --> plot rank_data[i*2, 6:]
    for i in range(2):
        ax2 = fig.add_subplot(spec[1,i])
        ax2.set_title('Rank histogram for station: ' + station_id + ' at ' + sub_titles[i], fontsize=fs2)
    	
    	# check if data available for given station and time:  if MSD value = -999 000 then no data. MSD value = rank_data[i*2=0, 12 UTC data, 4] > -99999
        if (rank_data[i*2,4] > -99999):
            ax2.bar(np.arange(100)+0.5, rank_data[i*2,6:]*100) # add + 0.5 to center bars
        else:
            ax2.text(0.5, 0.5, text_no_data, bbox={'pad':10, 'facecolor':'white'}, fontsize=fs3)
        ax2.set_xlim(0,100)
        ax2.axhline(y=1, c='k', ls='--') # horizontal black dashed line at y=0 
        ax2.set_xlabel('Percentile', fontsize=fs2)
        ax2.set_ylabel('Relative frequency', fontsize=fs2)
        ax2.xaxis.set_major_locator(ticker.FixedLocator([10,25,50,75,90])) 
    
    # third row of axis:
    # axis for legend, position: 2x0
    ax3 = fig.add_subplot(spec[2,0])
    ax3.axis('off')

    # add two legends: model data, and observations (with two columns)                            
    leg1 = ax3.legend(loc='upper center', handles=legend_elements1, fontsize=fs4)
    leg2 = ax3.legend(loc='lower center', handles=legend_elements2, ncols=2, fontsize=fs4, title='Quantiles for observations', title_fontsize=fs4)
    ax3.add_artist(leg1)
    ax3.add_artist(leg2)
     
    # axis for rank histogram text, positition 2x1
    ax4 = fig.add_subplot(spec[2,1])
    ax4.axis('off')
    # 'MSD:' + str(rank_data[0,4]) + ' p-value:' + str(rank_data[0,5])) # --> move to textbox 00 UTC
    ##rank_hist_00 =r'$\bf{00 UTC}$\n:' + r'$\bf{MSD}$:' + str(np.round(rank_data[0,4], ndigits)) + '\n' + r'$\bf{ p-value}:' + str(np.round(rank_data[0,5], ndigits))
    rank_hist_00 =r'$\bf{00 UTC}$: ' + r'$\bf{MSD}$: ' + str(np.round(rank_data[0,4], ndigits)) +' '+ r'$\bf{ p-value}$: ' + str(np.round(rank_data[0,5], ndigits))

    #'MSD:' + str(rank_data[2,4]) + ' p-value:' + str(rank_data[2,5])) --> move to textbox 12 UTC   
    ##rank_hist_12 =r'$\bf{12 UTC}$\n:' + r'$\bf{MSD}$:' + str(np.round(rank_data[2,4], ndigits)) + '\n' + r'$\bf{ p-value}:' + str(np.round(rank_data[2,5], ndigits))
    rank_hist_12 =r'$\bf{12 UTC}$: ' + r'$\bf{MSD}$: ' + str(np.round(rank_data[2,4], ndigits)) +' '+ r'$\bf{ p-value}$: ' + str(np.round(rank_data[2,5], ndigits))
    rank_hist_text = rank_hist_00 + '\n' + rank_hist_12
    # add text boxes with rank histogram data
    ax4.text(0.1,0.4, rank_hist_text, bbox={'pad':10, 'facecolor':'white'}, fontsize=fs3)
    
    # figure explanation:
    expl_data = 'Mean Square Deviation (' + r'$\bf{MSD}$' + ') relative to flat histogram\n' + r'$\bf{p-value}$' + ' of MSD relative to sampling distribution\n'
    expl_quantile='Quantiles are estimated based on observations at the station ' + station_id + ' (period: 1978-2020)\n'
    expl_model='Model data from experiment (' + exp_id + ')  interpolated to location of station ' + station_id
    
    fig_text = expl_data + expl_quantile+expl_model
    
    # axis for figure text:
    ax5 = fig.add_subplot(spec[3,:])
    ax5.axis('off')    
    # add text box with explanation to the bottom of the figure
    ax5.text(0.1, 0.4,fig_text, bbox={'pad':10, 'facecolor':'white'}, fontsize=fs4)
    
    if savefig:
        plt.savefig(fig_name, dpi=300)

def read_quantiles_data(quantile_sel_file):
    '''
    Reads quantile data from text file and returns a structured numpy array 
    Input: text file
    Output:a structured numpy array: sid, lon, lat, day_of_year, hour, q01, q10, q25, q50, q75, q90, q99
    '''
    # Header line contains:
    #station@hdr:integer longitude@hdr:real latitude@hdr:real day_of_year@hdr:integer hour@hdr:integer q01@body:real q10@body:real q25@body:real q50@body:real q75@body:real q90@body:real q99@body:real

    quantile_data = np.loadtxt(quantile_sel_file, skiprows=1, dtype={'names':('sid', 'lon', 'lat', 'day_of_year', 'hour', 'q01', 'q10', 'q25', 'q50', 'q75', 'q90', 'q99'), 'formats': ('i4', 'f4', 'f4', 'i4', 'i4', 'f4', 'f4', 'f4', 'f4', 'f4', 'f4', 'f4')})

    return quantile_data

def read_rank_data(rank_hists_file):
    '''
    Reads rank histogram data binned into a 100 bins: 0-1%, 1-2%, ..., 99-100%
    Input: text file
    Output: numpy array (unstructured) first columns contain station id, longitude, latitude, hour, msd, p_value, and then follows the bins f00 = 0-1%, ..., f99=99-100%
    '''

    #Header line contains:
    # station@hdr:integer longitude@hdr:real latitude@hdr:real hour@hdr:integer msd@body:real p_value@body:real f00@body:real f01@body:real f02@body:real f03@body:real f04@body:real f05@body:real f06@body:real f07@body:real f08@body:real f09@body:real f10@body:real f11@body:real f12@body:real f13@body:real f14@body:real f15@body:real f16@body:real f17@body:real f18@body:real f19@body:real f20@body:real f21@body:real f22@body:real f23@body:real f24@body:real f25@body:real f26@body:real f27@body:real f28@body:real f29@body:real f30@body:real f31@body:real f32@body:real f33@body:real f34@body:real f35@body:real f36@body:real f37@body:real f38@body:real f39@body:real f40@body:real f41@body:real f42@body:real f43@body:real f44@body:real f45@body:real f46@body:real f47@body:real f48@body:real f49@body:real f50@body:real f51@body:real f52@body:real f53@body:real f54@body:real f55@body:real f56@body:real f57@body:real f58@body:real f59@body:real f60@body:real f61@body:real f62@body:real f63@body:real f64@body:real f65@body:real f66@body:real f67@body:real f68@body:real f69@body:real f70@body:real f71@body:real f72@body:real f73@body:real f74@body:real f75@body:real f76@body:real f77@body:real f78@body:real f79@body:real f80@body:real f81@body:real f82@body:real f83@body:real f84@body:real f85@body:real f86@body:real f87@body:real f88@body:real f89@body:real f90@body:real f91@body:real f92@body:real f93@body:real f94@body:real f95@body:real f96@body:real f97@body:real f98@body:real f99@body:real

    rank_hists = np.loadtxt(rank_hists_file, skiprows=1)
    
    return rank_hists

def read_time_series_data(sim_data_file):
    '''
    Reads time series data 
    Input: text file
    Output: structured numpy array: sid, lon, lat, year, day_of_year, hour, value
    '''

    # Header line contains
    # station@hdr:integer longitude@hdr:real latitude@hdr:real year@hdr:integer day_of_year@hdr:integer hour@hdr:integer value@body:real

    time_series_data = np.loadtxt(sim_data_file, skiprows=1, dtype={'names': ('sid', 'lon', 'lat', 'year', 'day_of_year', 'hour', 'value'), 'formats': ('i4', 'f4', 'f4', 'i4', 'i4', 'i4', 'f4')})
    
    return time_series_data

def main(config_file):
    '''
    Main:
         reads config files
         reads input data: time_series_data, quantiles_data, rank_hist_data
         calls plotting function and saves figure with figure name given in config file
    '''
    ## Read from config file ##
    config.read(config_file)

    # data to be plotted
    variable = config['data']['variable']
    station_id = config['data']['station_id']
    lon = config['data']['longitude']
    lat = config['data']['latitude']
    year_beg = int(config['data']['year_beg'])
    year_end = int(config['data']['year_end'])
    month_beg = config['data']['month_beg']
    month_end = config['data']['month_end']    

    # input files
    sim_file = config['input']['sim_data']
    quantile_file = config['input']['quantile_sel']
    rank_hist_file = config['input']['rank_hists']
    
    # output files
    fig_name = config['output']['fig_name']
    
    ## Read input data ##
    time_series_data = read_time_series_data(sim_file)
    quantiles_data = read_quantiles_data(quantile_file)
    rank_hist_data = read_rank_data(rank_hist_file)

    ## Plot data ##
    plot_data(quantiles_data, rank_hist_data, time_series_data, station_id, lon, lat, year_beg, year_end, month_beg, month_end, fig_name, True)
    
if __name__=='__main__':
    if (len(sys.argv) < 2):
        sys.exit("Error: config file must be added as 2nd argument")
    elif (len(sys.argv) > 2):
        sys.exit("Error: only add config file as argument")
    config_file = sys.argv[1]
    
    main(config_file)
