#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Sep 18 23:01:01 EDT 2019
#import os.path, sys, os
import matplotlib.pyplot as plt
import xarray as xr, numpy as np, pandas as pd
import geoxarray

figname = 'fig.map.flor1860.resample.png'
figsize = None#(8,4)
savefig = [False, True][1]
ylabel = None#'mm/day'
title = None#'Mozambique mean annual precip max from FLOR1860 [mm/day]'

plt.figure(figsize=figsize)

# data 
scale = 24*3600
ifile = 'data/FLOR.CTL1860_newdiag_tigercpu_intelmpi_18_576PE.precip.Mozambique.masked.nc'
da = xr.open_dataarray(ifile)#.sel(time=slice('0101', '1000'))
years_start = range(101, 902, 100)
das = []
for year_start in years_start:
    years = slice(f'{year_start:04d}', f'{year_start+99:04d}')
    print(years)
    da_ = da.sel(time=years) \
        .groupby('time.year').max('time') \
        .mean('year') \
        .pipe(lambda x: x*scale)
    das.append(da_)
year_ranges = [f'{year_s:04d}-{year_s+99:04d}' for year_s in years_start]
da = xr.concat(das, dim=pd.Index(year_ranges, name='year'))
da.plot(col='year', col_wrap=3, cmap='Spectral')

if title is not None:
    plt.title(title)
if ylabel is not None:
    plt.ylabel(ylabel)

#plt.tight_layout()
if savefig:
    plt.savefig(figname, dpi=128)