#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri Aug 21 16:16:33 EDT 2020
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
#more imports
import seaborn as sns
import statsmodels.formula.api as smf
from misc.landmask import flagland
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ifile = 'q2m.validate.2019.nc'
year = ifile.split('.')[-2]
ds = xr.open_dataset(ifile)
# land ocean flags (land:1; ocean:0)
ds['landocean'] = flagland(ds)
df = ds.to_dataframe() * 1000 # units from kg/kg to g/kg
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    region = ['land', 'ocean', 'land+ocean'][1]
    figname = __file__.replace('.py', f'_{region}_{tt.today()}.png')
    plt.close()
    xlim = (-1.5, 1.5)
    ylim = xlim
    if region == 'ocean':
        df_ = df[df['landocean']==0] #ocean grids only
    elif region == 'land':
        df_ = df[df['landocean']>0] #land grids only
    else:
        df_ = df #both land and ocean
    g = sns.jointplot(x='anom', y='anom_true', data=df_, kind='hex', xlim=xlim, ylim=ylim)

    ax = g.ax_joint
    ax.grid('on')
    ax.plot(xlim, ylim, color='lightgray')

    #regression model
    r = smf.ols('anom_true ~ anom -1', data=df_).fit()
    y_ = np.array(xlim)*r.params.anom
    ax.plot(xlim, y_)
    s = f'{region} \ny = {r.params.anom:.2g}x \nR$^2$ = {r.rsquared:.2g}'
    ax.text(0.02, 0.98, s, ha='left', va='top', transform=ax.transAxes) 
    ax.set_xlabel(f'ERA5 {year} q2m anom [g/kg]: constant RH')
    ax.set_ylabel(f'ERA5 {year} q2m anom [g/kg]: true')

    plt.savefig(figname)
    print('[saved]:', figname)
    
    tt.check(f'**Done**')
    plt.show()