#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri Sep 27 16:59:27 EDT 2019
from datetime import datetime
import os.path, sys, os
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
import geoxarray
 
ifiles = ['chirps.p05.Mozambique.masked.nc',
    'FLOR.CTL1860_newdiag_tigercpu_intelmpi_18_576PE.precip.Mozambique.masked.nc',
    'FLOR.CTL1990_newdiag_tigercpu_intelmpi_18_576PE.precip.Mozambique.masked.nc']

def do_fldmean(ifile):
    #ifile = 'chirps.p05.Mozambique.masked.nc'
    ofile = ifile.replace('.nc', f'.fldmean.nc')
    if os.path.exists(ofile):
        print('[exists]:', ofile)
        return
    # read
    da_i = xr.open_dataarray(ifile)

    # calculate
    da = da_i.geo.fldmean()

    # save
    da.attrs = da_i.attrs
    da.to_dataset(name=da_i.name) \
        .to_netcdf(ofile, unlimited_dims='time')
    print('[saved]:', ofile)
    
if __name__ == '__main__':
    tformat = '%Y-%m-%d %H:%M:%S'
    t0 = datetime.now()
    print('[start]:', t0.strftime(tformat))

    for ifile in ifiles:
        print(ifile)
        do_fldmean(ifile)

    t1 = datetime.now()
    print('[end]:', t1.strftime(tformat))
    print('[total time]:', f'{(t1-t0).seconds:,} seconds')