#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Nov  3 01:56:59 2022
@author: li.xu

contact:
    Li Xu
    CPC/NCEP/NOAA
    (301)683-1548
    li.xu@noaa.gov

"""
#### Script Documentation Block  ####################################################
#
# Script Name: plt_skewt.py
#
#
# Abstract:
# plot the GFS SkewT plot for CPCI
#
# Input Data Source: ./stn/stn_i.txt
#
# Output Destination: ../plot
#
# Software Package Used: Python/3.6.3
#                        Matplotlib/2.2.2
#
# Script History Log: First Implemented Oct 2020
# Original Code Development Contact: Li Xu, CPC, 301-683-1548

# modified 19Jan2022 add the Central Africa stations
#
#
####################################################################################

import numpy as np

import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
import pandas as pd

import metpy.calc as mpcalc
from metpy.cbook import get_test_data
from metpy.plots import add_metpy_logo, Hodograph, SkewT
from metpy.units import units


def fdate(dt, cyc, fhh):
    '''compute the forecast date and hour'''
    from datetime import datetime
    from datetime import timedelta
    today = datetime.strptime(dt + cyc, '%Y%m%d%H')
    print('dt cyc:', today, fhh)

    delta = today + timedelta(hours=int(fhh))
    return delta.strftime('%Y%m%d%H')
# Create a new figure. The dimensions here give a good aspect ratio

def plot(fin='stn/stn_1.txt',fout='stn/stn_1.png'):

    with open(fin, 'r') as f:
        ttl = f.readline()
        pres, hght, temp, dwpt, drct, sknt = np.loadtxt(f, unpack=True)

    ctitle = f'NOAA-NCEP GFS {dt} {cyc}Z CYC {fhh}hr Fcst valid {hr_fcst}Z {dt_fcst} \n\n {ttl}'
    print(ctitle)

    hght = hght * units.hPa
    p = pres * units.hPa
    T = temp * units.degC
    Td = dwpt * units.degC
    wind_speed = sknt * units.knots
    wind_dir = drct * units.degrees
    u, v = mpcalc.wind_components(wind_speed, wind_dir)


    fig = plt.figure(figsize=(9, 9))

    # Grid for plots
    skew = SkewT(fig, rotation=45)

    # Plot the data using normal plotting functions, in this case using
    # log scaling in Y, as dictated by the typical meteorological plot
    skew.plot(p, T, 'r')
    skew.plot(p, Td, 'g')
    skew.plot_barbs(p, u, v)
    skew.ax.set_ylim(1000, 100)

    # Add the relevant special lines
    skew.plot_dry_adiabats()
    skew.plot_moist_adiabats()
    skew.plot_mixing_lines()

    # Good bounds for aspect ratio
    skew.ax.set_xlim(-50, 60)
    skew.ax.set_title(ctitle)

    # Create a hodograph
    ax_hod = inset_axes(skew.ax, '40%', '40%', loc=1)
    h = Hodograph(ax_hod, component_range=80.)
    h.add_grid(increment=20)
    h.plot_colormapped(u, v, hght)

    fig.savefig(fout,dpi=100,bbox_inches='tight')

    plt.close()



import sys,glob
if len(sys.argv) <= 1:
    dt = '20200912'
    cyc = '00'
    fhh = '136'
    print( 'input: 20200912 00 136'  )
else:
    dt = sys.argv[1]
    cyc = sys.argv[2]
    fhh = sys.argv[3]

fcst = fdate(dt, cyc, fhh)
dt_fcst = fcst[:8]
hr_fcst = fcst[-2:]


total=83
for i in range(1, total + 1):  # index start from 1
    plot(fin=f'./stn/stn_{i}.txt',fout=f'../plot/{int(fhh):02d}_stn{i}.png')

# modified 19Jan2022 add the Central Africa stations
#
for i in range(50):  # index start from 0
    plot(f'./stn/cf_stn_{i}.txt',fout=f'../plot/{int(fhh):02d}_cf_stn{i}.png')