#!/usr/bin/env python
import xarray as xr, matplotlib.pyplot as plt
import geoxarray
try:
    import wystart
except:
    pass
plt.close('all')

#ds = xr.open_dataset('ts.historical_ssp245.1850-2100.ens14.nc')
ds = xr.open_dataset('ts.historical_ssp245.1850-2100.ens18.nc')
ds.load()
da = ds['ts'].groupby('time.year').mean('time').geo.fldmean().pipe(lambda x: x - x.sel(year=slice(1981,2010)).mean('year'))
da.mean('model').plot(color='C0', label='historical_ssp245')
for m in da.model.values: 
    da.sel(model=m).plot(color='C0', lw=0.5, alpha=0.5)

#ds_ = xr.open_dataset('ts.historical_ssp585.1850-2100.ens14.nc')
ds_ = xr.open_dataset('ts.historical_ssp585.1850-2100.ens18.nc')
ds_.load()
da_ = ds_['ts'].groupby('time.year').mean('time').geo.fldmean().pipe(lambda x: x - x.sel(year=slice(1981,2010)).mean('year'))
da_.sel(year=slice(2015, None)).mean('model').plot(color='C1', label='ssp585')
for m in da.model.values: 
    da_.sel(year=slice(2015, None)).sel(model=m).plot(color='C1', lw=0.5, alpha=0.5)

n_ens = da.model.size

plt.legend()
plt.ylabel('K')
plt.title(f'CMIP6 global mean TS anomaly: 1981-2010 clim; {n_ens} ens')

plt.savefig(f'fig.tsa.glbmean.historical_ssp.1850-2100.ens{n_ens}.png')