#!/usr/bin/env python3

'''
Written by Madeleine Ekblom, 2023-05-10

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

config = configparser.ConfigParser()

#plot_data(quantiles_data, rank_hist_data, time_series_data, fig_title, fig_name, year_beg, year_end, month_beg, month_end, True)
def plot_data(quantiles_data, rank_data, time_series_data, fig_title, fig_name, year_beg, year_end, month_beg, month_end, savefig=False):
    '''
    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)

    # 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]

    # 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(3,2)
    fig.suptitle(fig_title, fontsize=20)

    # 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
            ax = fig.add_subplot(spec[i,j])
            ax.set_title(sub_titles[c])
            ax.set_xlim(dates[0], dates[-1])
            ax.xaxis.set_major_locator(MonthLocator())
            ax.xaxis.set_major_formatter(DateFormatter('%b'))
    
            # quantile data:
            # find quantile data hour = 0, 6, 12, 18
            qh_ind = (quantiles_data['hour'][:] == hours[c])
            for q in q_names:
                q_data = quantiles_data[q][qh_ind]
                ax.plot(dates, q_data[:365], 'k-') # --> change q only contains one year of data!!
            # 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):
                ax.scatter(dates, t_data[n*365:(n+1)*365], marker='.')
            c = c + 1

    # plot rank histogram data:
    ax2 = fig.add_subplot(spec[2,:])
    ax2.set_title('Rank histogram\n' + 'MSD:' + str(rank_data[4,4]) + ' p-value:' + str(rank_data[4,5]))
    ax2.bar(np.arange(100), rank_data[4,6:]*100)
    ax2.set_xlim(0,100)
    ax2.axhline(y=1)
    
    if savefig:
        plt.savefig(fig_name, dpi=300)

    #plt.savefig(fig_name, dpi=300)

#    plt.show( )


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']
    longitude = config['data']['longitude']
    latitude = 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 ##
    fig_title = 'Station: ' + station_id + ', ' + str(year_beg) + month_beg + '-' + str(year_end) + month_end + '\nLat=' + latitude + ' Lon=' + longitude
    plot_data(quantiles_data, rank_hist_data, time_series_data, fig_title, fig_name, year_beg, year_end, month_beg, month_end, 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)
