#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@author: slillo
"""

import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime as dt,timedelta
import matplotlib.dates as mdates
import xarray as xr
import netCDF4 as nc
from LIM_CPC import driver
import data_retrieval
import os
import LIM_CPC
###from write_html_page import *
import warnings
warnings.filterwarnings('ignore')

def lat2strNodeg(x):
    deg = u"\u00B0"
    if x<0:
        return f'{abs(x)}S'
    else:
        return f'{x}N'

#%%

### BEGIN USER INPUT ###
##LIMpage_path = 
##RTdata_path = 
##getdataUSER = 
##getdataPASS = 
LIMpage_path = "/cpc/home/hwang/LIM_CPC/data_output"
RTdata_path = "/cpc/home/hwang/LIM_CPC/data_realtime"
getdataUSER = "hui.wang@noaa.gov"
getdataPASS = "ub4pbSAI"


fullVariance = True
DPI=120
credit='NOAA/PSL and University of Colorado/CIRES \nExperimental LIM Forecast'


### END USER INPUT ###

START = dt.now()
testdays = [dt.now().replace(hour=0,minute=0,second=0,microsecond=0)-timedelta(days=1+i) for i in range(10)]
FORECASTDAYS = sorted([t for t in testdays if not os.path.isdir(f'{LIMpage_path}/{t:%Y%m%d}')])


# %%===========================================================================
# INITIALIZE AND RUN BLENDED LIM FORECAST
# =============================================================================

print('\nInitializing...')
LIMdriver = driver.Driver('namelist.py')
LIMdriver.get_variables(read=True)
LIMdriver.get_eofs(read=True)
LIMdriver.prep_realtime_data(limkey=1) #dummy limkey just to get available times

FORECASTDAYS = sorted(list(set(FORECASTDAYS)&set(LIMdriver.RT_VARS['time'])))

#%%
print('\nRunning blended LIM forecasts...')
for T_INIT in FORECASTDAYS:
    START = dt.now()

    dirname = f'{T_INIT:%Y%m%d}'
    FCSTDIR = f'{LIMpage_path}/{dirname}'
            
    os.system(f'mkdir {FCSTDIR}')

    weekday = T_INIT.weekday()
    dayoffset = (4-weekday)%7

    try:
        LIMdriver.run_forecast_blend(t_init=T_INIT,lead_times=np.arange(0,29+dayoffset),fullVariance=fullVariance,\
                        save_netcdf_path=FCSTDIR)
    except:
        print(f'NO BLEND FORECAST FOR {T_INIT:%Y%m%d}')
        continue

    for LT in set([(0,),(14,),(21,),(28,),(21,28)]+[(21+dayoffset,),(28+dayoffset,),(21+dayoffset,28+dayoffset)]):
        print(f'\n{T_INIT:%Y%m%d} T2m {LT}')
        LIMdriver.plot_map(varname='T2m',t_init=T_INIT,lead_times=LT,fullVariance=fullVariance,gridded=True,\
                    prop={'levels':np.linspace(-5,5,21),'cbar_label':'$^oC$','dpi':DPI,'addtext':credit},save_to_path = FCSTDIR)
        print(f'\n{T_INIT:%Y%m%d} SLP {LT}')
        LIMdriver.plot_map(varname='SLP',t_init=T_INIT,lead_times=LT,fullVariance=fullVariance,add_offset='data_clim/SLP.JRA.1991-2020.nc',gridded=True,\
                    prop={'levels':np.linspace(-10,10,21).astype(int),'interpolate':1,'cbar_label':'$hPa$','dpi':DPI},save_to_path = FCSTDIR)
        print(f'\n{T_INIT:%Y%m%d} H500 {LT}')
        LIMdriver.plot_map(varname='H500',t_init=T_INIT,lead_times=LT,fullVariance=fullVariance,add_offset='data_clim/H500.JRA.1991-2020.nc',gridded=True,\
                    prop={'levels':np.linspace(-100,100,21).astype(int),'interpolate':1,'cbar_label':'$m$','dpi':DPI},save_to_path = FCSTDIR)
        print(f'\n{T_INIT:%Y%m%d} colIrr {LT}')
        LIMdriver.plot_map(varname='colIrr',t_init=T_INIT,lead_times=LT,fullVariance=fullVariance,add_offset='data_clim/colIrr.JRA.1991-2020.nc',gridded=True,\
                    prop={'cmap':{-2:'darkorange',-1:'sienna',-0.2:'w',0.2:'w',1:'seagreen',2:'turquoise'},\
                    'levels':(-200,200),'interpolate':1,'cbar_label':'$W/m^2$','figsize':(10,3.5),'drawstates':False,'latlon':True,'dpi':DPI},\
                    save_to_path = FCSTDIR)

    for bounds in [(-15,0),(-7.5,7.5),(0,15)]:
        LIMdriver.plot_timelon(varname='colIrr',t_init=T_INIT,lat_bounds=bounds,daysback=75,gridded=True,add_offset='data_clim/colIrr.JRA.1991-2020.nc',\
                                prop={'cmap':{-2:'darkorange',-1:'sienna',-0.2:'w',0.2:'w',1:'seagreen',2:'turquoise'},\
                                      'levels':(-200,200),'cbar_label':'$W/m^2$','dpi':DPI},\
                                      save_to_file=f'{FCSTDIR}/HOV_trop_{lat2strNodeg(bounds[0])}{lat2strNodeg(bounds[1])}.png')
    
    for bounds in [(20,40),(30,50),(40,60)]:
        LIMdriver.plot_timelon(varname='H500',t_init=T_INIT,lat_bounds=bounds,daysback=75,gridded=True,add_offset='data_clim/H500.JRA.1991-2020.nc',\
                                prop={'cmap':{-2:'darkorchid',-1:'dodgerblue',-0.2:'w',0.2:'w',1:'tomato',2:'firebrick'},\
                                      'levels':(-100,100),'cbar_label':'$m$','dpi':DPI},\
                                      save_to_file=f'{FCSTDIR}/HOV_500_{lat2strNodeg(bounds[0])}{lat2strNodeg(bounds[1])}.png')
    
#    LIMdriver.plot_teleconnection(T_INIT=T_INIT,gridded=True,daysback=60,prop={'dpi':DPI},save_to_path = FCSTDIR)
    
    FINISH = dt.now()
    ELAPSED = (FINISH-START).total_seconds()/60
    print(f'\n {ELAPSED:.2f} minutes to run forecast for {T_INIT:%Y%m%d}\n')