#!/usr/bin/env python

## Just run code in recipes, use it carefully
## example: use netcdf4 to make cache file. Since python is faster than ncl.
#;DIAG_NORCPM; RUNTHESECODES: print('    no codes run here.')

import xarray as xa
import numpy as np
import sys,glob
from os.path import exists

years = range(1982,2018+1)
ny = len(years)
noleap_w = xa.DataArray([31,28,31,30,31,30,31,31,30,31,30,31])
noleap_w = noleap_w/sum(noleap_w)

reg = 'BaffinBay'
latb  = 66.
late  = 76.
lonb  = -80.
lone  = -50.
var = 'hi'
fns = glob.glob(f'./{var}_??.nc')
fns.sort()

## bad method, error when cross lon 0 line
if lonb < 0: lonb += 360
if lone < 0: lone += 360

gridds = xa.open_dataset(fns[0])
wgt = gridds.tarea
lat = gridds.TLAT
lon = gridds.TLON
lon = lon.where(lon>0, lon+360) ## 0-360

wgt = wgt.where(latb<lat ,0)
wgt = wgt.where(late>lat ,0)
wgt = wgt.where(lonb<lon ,0)
wgt = wgt.where(lone>lon ,0)
wgt.name = 'weighting'
wgt.to_netcdf(f'{reg}_weighting.nc')  ## save for plot

tsfns = list()
yrtsfns = list()
for fn in fns:
    n = fn[-5:-3]
    ofn = f'{var}_{n}_{reg}_ts.nc'
    yrofn = f'{var}_{n}_{reg}_yrts.nc'
    tsfns.append(ofn)
    yrtsfns.append(yrofn)
    if exists(ofn) and exists(yrofn): continue
    ds = xa.open_dataset(fn)
    dd = ds[var].weighted(wgt)
    ovar  = dd.mean(('nj','ni'),skipna=True)
    ovar.name = var
    ## yearly mean with noleap
    yovar = xa.DataArray(np.zeros(ny),dims=['year'], coords={'year':years})
    yovar.name = var
    for y in range(ny):
        yovar[y] = (ovar[12*y:12*y+1]*noleap_w).sum()
        
    ovar.to_netcdf(ofn)
    yovar.to_netcdf(yrofn)
    

tsfs = xa.open_mfdataset(tsfns,combine='nested',concat_dim='member')
outds = xa.Dataset({
    'avg': tsfs[var].mean(dim='member')
    ,'std': tsfs[var].std(dim='member')
    ,'min': tsfs[var].min(dim='member')
    ,'max': tsfs[var].max(dim='member')
    })
outds.to_netcdf(f'{var}_{reg}_ens_ts.nc')

tsfs = xa.open_mfdataset(yrtsfns,combine='nested',concat_dim='member')
outds = xa.Dataset({
    'avg': tsfs[var].mean(dim='member')
    ,'std': tsfs[var].std(dim='member')
    ,'min': tsfs[var].min(dim='member')
    ,'max': tsfs[var].max(dim='member')
    })
outds.to_netcdf(f'{var}_{reg}_ens_yrts.nc')