"""
bias_corr_conusKF.py &
using ECMWF forecast and ERA5 verification data, implement Kalman-filter
type bias correction approach to estimate forecast bias.   Grid covers
CONUS, approximately.  Coded by Tom Hamill, May-Jun 2020.
tom.hamill@noaa.gov
"""

import pygrib
from dateutils import daterange, dateshift, dayofyear, splitdate
import os, sys
import numpy as np
import _pickle as cPickle

# =====================================================================

def decayavg_bias_simultaneousKF(Kalman_gain_4D, nlats, nlons, \
    obs, fcst, bias_estimate):

    # ---- update the bias correction using Kalman filter approach and 
    #      simultaneous use of all F and O across CONUS
    
    for i in range(nlons):
        for j in range(nlats):
            bias_estimate[j,i] = bias_estimate[j,i] - \
                np.sum(Kalman_gain_4D[j,i,:,:]* \
                (obs[:,:] - (forecast[:,:] - bias_estimate[:,:])))

    return bias_estimate

# =====================================================================

clead = sys.argv[1]  # lead time, e.g., 12, 72, 120 (in hours)
    
# --- read in the Kalman gain generated by invert_localized_covs_2019.py  
    
infile = 'Kalman_gain_4D_lead='+clead+'.cPick'
print (infile)
inf = open(infile, 'rb')
Kalman_gain_4D = cPickle.load(inf)
inf.close()

# --- other initialization stuff    
    
ilead = int(clead)
datadir = '/Users/Tom/python/ecmwf/'
cvariable = '2t'
tally_statistics = True

dateend = dateshift('2019123100',-ilead)
date_list_anal = daterange('2018110100',dateend,24) # initial time of the current forecast
ndates = len(date_list_anal)
date_list_fcst = []
for idate in range(ndates):
    date_list_fcst.append(dateshift(date_list_anal[idate],ilead)) # initial times of fcst

# ---- loop over dates and update bias estimates

for idate, datea in enumerate(date_list_anal):
    
    datef = date_list_fcst[idate]
    print ('------ processing analysis, forecast dates = ', datea, datef)
    if int(datea) >= 2019010100: tally_statistics = True

    # ---- read the ECMWF ERA5 reanalysis at the forecast date/time
    
    infile = datadir + 't2m_era5_halfdegree_'+datef+'.cPick'
    #print (infile)
    inf = open(infile, 'rb')
    analysis = cPickle.load(inf)
    if idate == 0:
        lats = cPickle.load(inf)
        lons = cPickle.load(inf)
        nlats, nlons = np.shape(lats)
        bias_decayavg = np.zeros((nlats, nlons), dtype=np.float32)
        bias_estimate = np.zeros((nlats, nlons), dtype=np.float64)            
    inf.close()
    
    # ---- read the ECMWF control forecast at this lead time and initial date
 
    infile = datadir + cvariable+'_'+datea+'_f'+clead+'.grib2'  
    #print (infile)
    grbfile = pygrib.open(infile) 
    grb = grbfile.select()[0] 
    forecast = grb.values
    grbfile.close()
    
    # ---- produce estimate of Kalman filter bias correction with seasonal variability.
        
    bias_estimate = decayavg_bias_simultaneousKF(Kalman_gain_4D, \
        nlats, nlons, analysis, forecast, bias_estimate)
        
    # ---- write bias estimates to file if in 2019.
    
    if tally_statistics == True:
        outfilename = datadir + 'bias_est_conusKF'+datef+'_f'+clead+'.cPick'
        #print ('       writing bias estimates to ', outfilename)
        ouf = open(outfilename, 'wb')
        cPickle.dump(bias_estimate, ouf)
        ouf.close()