#!/usr/bin/env python3

'''
Written by Madeleine Ekblom, 2023-05-10
updates 2023-11-27:
- visual changes of the plot: labels, legends , title, and colours updated

Based on Jouni Räisänen's script for plotting quantiles/time series data (00, 06, 12, 18 UTC) and rank histogram (24UTC=combination of 00,06,12,and 18 UTC data)

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, variable, station_id, lon, lat, year_beg, year_end, month_beg, month_end, figname, savefig=False):
    '''
    Plot data: quantiles data give the lines and the time series data give the dots
    '''
    # 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
    
    # experiment id: now use 'test data' -  to be replaced with real exp_id (as input argument)
    exp_id = 'test data' 
        
    # 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)
    
    # arrays for quantile names, axes titles, hours to plot
    q_names = ['q01', 'q10', 'q25', 'q50', 'q75', 'q90', 'q99']
    sub_titles = ['00 UTC', '06 UTC', '12 UTC', '18 UTC']
    hours = [0, 6, 12, 18]
    
    # 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
    fs4=10 # fontsize for extra text box, legend
    ndigits = 3 # round off to ndigits (e.g., lon, lat)
    
    # use variable to get variable and corresponding unit, used as ylabel in plots:
    ylabel = variable_shortnames[variable] + ' (' + variable_units[variable] + ')'
        
    # 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%')]
    
    # 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 = 'SYNOP station: ' + station_id + lat_text + lon_text + '\nModel data from ' + str(year_beg) + month_beg + '-' + str(year_end) + month_end 
    
    # calculate number of years:
    nyears = year_end - year_beg + 1

    # set up figure
    fig = plt.figure(figsize=(10,10), layout='constrained')
    spec = fig.add_gridspec(4,2)
    fig.suptitle(fig_title, fontsize=fs1)

    # counter
    c = 0

    # plot quantiles/time series for times 00, 06, 12, 18:
    for i in range(2):
        for j in range(2):
            # set up axis: title, xaxis, x- and y-labels
            ax = fig.add_subplot(spec[i,j])
            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(ylabel, fontsize=fs2)
            
            # quantile data:
            # find quantile data hour = 0, 6, 12, 18
            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, alpha=alpha[1]
            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, alpha=alpha[2]            
            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, 6, 12, 18
            th_ind = (time_series_data['hour'][:] == hours[c])
            t_data = time_series_data['value'][th_ind] # all years
            
            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='.') 
            c = c + 1

    # plot rank histogram data:
    ax2 = fig.add_subplot(spec[2,0])
    ax2.set_title('Rank histogram', fontsize=fs2)
    ax2.bar(np.arange(100)+0.5, rank_data[4,6:]*100) # add + 0.5 to center bars
    ax2.set_xlim(0,100)
    ax2.set_xlabel('Percentile', fontsize=fs2)
    ax2.set_ylabel('Relative frequency ', fontsize=fs2)
    ax2.axhline(y=1, c='k', ls='--') # horizontal line at y=1.0 
    ax2.xaxis.set_major_locator(ticker.FixedLocator([10,25,50,75,90]))
    
    # add a sixth axis for legends and text data:
    ax3 = fig.add_subplot(spec[2,1])
    ax3.axis('off')
        
    # legend with two columns at the upper centre of the figure:                                   
    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)

    # add text boxes with (1) rank histogram data and (2) explanation
    rank_hist_text = r'$\bf{MSD}$:' + str(np.round(rank_data[4,4], ndigits)) + '\n' + r'$\bf{ p-value}$:' + str(np.round(rank_data[4,5], ndigits))
   
    # add text boxes to axis in right lower corner
    ax3.text(0.35, -0.3, rank_hist_text, bbox={'pad': 10, 'facecolor':'white'}, fontsize=fs3)
    
    # figure explanation: explaining what quantiles and model data are
    expl_msd_pvalue = '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: 2010-2021)\n' 
    expl_model='Model data from experiment (' + exp_id + ') interpolated to location of this station '
    fig_text = expl_msd_pvalue + expl_quantile + expl_model
    
    # add subplot at the bottom of the figure, and add figure explanation as text to this subplot
    ax4 = fig.add_subplot(spec[3,:])
    ax4.axis('off')
    ax4.text(0.1, 0.4, fig_text, bbox={'pad':10, 'facecolor':'white'}, fontsize=fs4)
    
    # save figure:
    if savefig:
        plt.savefig(figname, 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
    figname = 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, variable, station_id, lon, lat, year_beg, year_end, month_beg, month_end, figname, 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)
