#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Created on Fri Sep  4 19:23:36 2020

#### Script Documentation Block  ####################################################
#
# Script Name: meteogram.py
#
# key subroutine for meteogram plot
#
# Abstract:
# plot the GEFS ensemble Meteogram plot
#
# Input Data Source: GEFS  /com/gens/prod/gefs.yyyymmdd/
#
# Output Destination: $DATA
#
# Software Package Used: Python/3.6.3
                         Matplotlib/2.2.2
#
# Script History Log: First Implemented Sep 2020
# Original Code Development Contact: Li Xu, CPC, 301-683-1548

"""
import matplotlib
matplotlib.use('Agg')

import numpy as np
import matplotlib.pyplot as plt
import datetime



def plot_meteorogram(fig,
                     data,
                     t1,
                     stninfo,
                     vname='2m Air Temp. (in degC)',
                     boxcolor='lightblue',
                     dt='20200201',
                     cyc='00',
                     pout='./all_plot/tmp2m'
                     ):
    '''data [ens,40], t1,t2 [40]'''

    t2=np.nanmean(data,axis=0)  # ensmeble mean
    iday=10

    dates = []
    base = datetime.datetime.strptime(dt, '%Y%m%d')
    for i in range(iday+1):  # one extra day annoate
        if i == 0:
            dates.append(
                (base + datetime.timedelta(days=i)).strftime('%d%h\n%Y'))

        else:
            dates.append((base + datetime.timedelta(days=i)).strftime('%d%h'))


    mx = np.nanmax(data, axis=0)
    mn = np.nanmin(data, axis=0)
    q1 = np.percentile(data, 25, axis=0)
    q3 = np.percentile(data, 75, axis=0)
    q2 = np.percentile(data, 50, axis=0)
    q90 = np.percentile(data, 90, axis=0)
    q10 = np.percentile(data, 10, axis=0)

    N = data.shape[-1]
    if cyc == '00' : istart=1
    if cyc == '06' : istart=2
    if cyc == '12' : istart=3
    if cyc == '18' : istart=4


    i = range(istart,istart+N)

    ax1 = fig.add_subplot(1, 1, 1)
    ax1.bar(i,
            q90 - q10,
            bottom=q10,
            width=0.4,
            color=boxcolor,
            edgecolor='black')
    ax1.bar(i,
            q3 - q1,
            bottom=q1,
            width=0.7,
            color=boxcolor,
            edgecolor='black')
    ax1.bar(i, q2 - q2, bottom=q2, width=0.7, color='black', edgecolor='black')
    ax1.errorbar(i,
                 q90,
                 yerr=[q90 - q90, mx - q90],
                 fmt='none',
                 color='k',
                 linewidth=1)
    ax1.errorbar(i,
                 q10,
                 yerr=[q10 - mn, q10 - q10],
                 fmt='none',
                 color='k',
                 linewidth=1)

    ax1.plot(i, t2, 'b-', label='ensemble mean')
    ax1.plot(i, t1, 'r--', label='control member')
    ax1.legend()

    iday = np.arange(0, N+4, 4)
    plt.xticks(iday, dates)
    plt.xticks(fontsize=8)

    import matplotlib.ticker as tic
    ax1.minorticks_on()
    ax1.xaxis.set_minor_locator(tic.AutoMinorLocator(n=4))

    plt.grid('on', linestyle=':')


    title = f'NOAA-NCEP GEFS {dt} {cyc}Z CYC 6hly Fcst for 10-day Period \n {vname} \n'
    ax1.set_title(title,fontsize=12)
    ltitle = f'{stninfo[1]} / {stninfo[2]}'
    ax1.set_title(ltitle,loc='left',fontsize=10)

    rtitle = f'Lat:{stninfo[3]}  Lon:{stninfo[4]}  Elev:{stninfo[5]}'
    ax1.set_title(rtitle,loc='right',fontsize=10)

    fig.text(0.2,-0,'Ensemble Meteogram:  20-member ensemble, control member and the ensemble mean',fontsize=8,fontweight='light')
    fig.text(0.2,-0.03,'Box-Whisker Attributes: minimum, maximum and 10th, 25th(Q1),50th(median), 75th(Q3), 90th percentiles',fontsize=8,fontweight='light')

    plt.savefig(f'{pout}_{stninfo[1]}.{stninfo[2]}.png',bbox_inches="tight")
    plt.clf()


# %%

def test():
    np.random.seed(19680801)

    # fake up some data
    t = []
    for i in range(40):
        t.append(np.random.randn(20) * i / 2 + 50)

    data = np.asarray(t).T
    t1 = np.nanmean(data, axis=0)
    t2 = np.random.randn(40) * 2 + 50

    base = datetime.datetime(2005, 2, 1)
    iday = 10
    stninfo=('00000','Mindelo','Cabo-Verde','16.83','-25.05','17.0','xgrid','ygrid')
    figsize=(10, 4)
    fig1 = plt.figure(figsize=figsize)
    plot_meteorogram(fig1,data, stninfo,cyc='18')


#%%
def clean(D, vmin=-999, vmax=1e8):
    '''clean D for setup to nan'''
    D[D <= vmin] = np.nan
    D[D >= vmax] = np.nan
    return D

def rbin(filename='m.bin', shape=(60, 150), dtype='<f4', offset=0):
    '''shape (-1,60,150) to read all the data
    dtype='>f4' for big endian
    offset is same as fromfile shape*4 for f4'''
    n = np.cumprod(shape)[-1]
    try:
        D = np.fromfile(filename, count=n, dtype=dtype).reshape(shape)
    except:
        print('Read error: ' + filename)

    return clean(D)


def read_stns(file='./all_stations.txt'):
    stns=[]
    with open(file,'r') as f:
        lines=f.readlines()
        for l in lines:
            if len(l) > 20: # no-empty line
                stns.append(l.split())
            else:
                print('ignore line:',l)

    return stns        # (stid, stname,country,lat,lon,elev,ygrid,xgrid

def get_xi_yi(lon,lat):
    '''convert lon, lat (str) to grid index'''
    lat=float(lat)
    lon=float(lon)

    if lon >=0:
        xi=int(round(lon/0.5))
    else:
        xi=int(round(360+lon)/0.5)

    yi=int(round((lat+90)/0.5))

    if xi >= 720 : xi=719
    if yi >= 361 : yi=360

    print('grid index (y,x):',yi,xi)
    return xi,yi