#!/usr/bin/env python
# coding: utf-8

# In[1]:


get_ipython().run_line_magic('matplotlib', 'inline')
import xarray as xr
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
import intake


# In[2]:


catalog = intake.open_esm_datastore('https://cmip6-pds.s3.amazonaws.com/pangeo-cmip6.json')


# In[26]:


# Global attributes are described in https://docs.google.com/document/d/1h0r8RZr_f3-8egBMMh7aqLwy3snpD6_MrDz1q8n5XUk/edit

# Variables are described in this spreadsheet
# http://proj.badc.rl.ac.uk/svn/exarch/CMIP6dreq/tags/latest/dreqPy/docs/CMIP6_MIP_tables.xlsx
variable_id = 'tas' # Surface Air Temperature
table_id = 'Amon' # Monthly data from Atmosphere

grid = 'gn' #

# Records for Institution, experiment, and source_id are stored in https://github.com/WCRP-CMIP/CMIP6_CVs
experiment_ids = ['historical', 'ssp126', 'ssp245', 'ssp370', 'ssp585']
activity_ids = ['ScenarioMIP', 'CMIP'] # Search Scenarios & Historical data only
institution_id = 'NASA-GISS'
source_id = ['GISS-E2-1-G']

# Member ID is formed from subexperiment id & variant label
# variant label is
# variant_label: a label constructed from 4 indices stored as global attributes:  
# variant_label = r<k>i<l>p<m>f<n>
#       where
# k = realization_index
# l = initialization_index
# m = physics_index
# n = forcing_index
# For a given experiment, the realization_index, initialization_index, physics_index, and forcing_index are used to uniquely identify each simulation of an ensemble of runs contributed by a single model.

member_id = 'r1i1p3f1'


# In[27]:


res = catalog.search(activity_id=activity_ids, institution_id=institution_id, experiment_id=experiment_ids, source_id=source_id, table_id=table_id, variable_id=variable_id, grid_label=grid, member_id=member_id)
display(res.df)


# In[28]:


datasets = res.to_dataset_dict(zarr_kwargs={'consolidated': True, 'decode_times': True})


# In[29]:


display(datasets)


# In[7]:


# Put the air temperature into units that we understand better

for i in datasets:
    ds = datasets[i]
    if ds.tas.units == 'K':
        ds.tas.values = (ds.tas.values - 273.15) * 9.0 / 5.0 + 32.0
        ds['tas', 'units'] = 'F'


# In[8]:


display(datasets['ScenarioMIP.NASA-GISS.GISS-E2-1-G.ssp126.Amon.gn'])


# In[9]:


# Lets look at how the air temperature changed over the period of this simulation. We will extract the air temperature at the end from the air temperature at the beginning
# We take an average over the last twelve months minus the first twelve months to even out seasonality

# Lets look at SSP585, one of the more extreme emissions scenarios where little is done to curb greenhouse gas emissions

tas = datasets['ScenarioMIP.NASA-GISS.GISS-E2-1-G.ssp585.Amon.gn']['tas'].isel(time=np.arange(1032-12,1032)).mean(dim='time') - datasets['ScenarioMIP.NASA-GISS.GISS-E2-1-G.ssp585.Amon.gn']['tas'].isel(time=np.arange(0,12)).mean(dim='time')
tas.compute()


# In[10]:


tas.plot()


# In[11]:


# Compute the mean of the differences
mean_difference = tas.mean().compute()


# In[12]:


mean_difference


# In[13]:


# That shows that on average Earth warms 8 F over the next 100 years in the SSP585 scenario according to the NASA GISS model!


# In[22]:


# We can also look at how the temperature evolves over a single location for all of the scenarios too.
# Here we will look at Colorado Springs

location_name = 'Colorado Springs'
location_lon = -104.82
location_lat = 38.83

# convert to 0-360 for longitudes
if location_lon < 0:
     location_lon = 360 + location_lon


# In[23]:


ds_scenarios = xr.Dataset()


# interate through the locations and create a dataset
# containing the temperature values for each location
for i in datasets:
    ds = datasets[i]

    ds2 = ds.sel(lon=location_lon, lat=location_lat, method='nearest')
    ds2 = ds2.rename({'tas' : i}).drop(('lat', 'lon', 'height', 'lat_bnds', 'lon_bnds', 'time_bnds', 'member_id'))
    ds_scenarios = xr.merge([ds_scenarios, ds2], compat='override')

ds_scenarios


# In[24]:


# Compute the extraction of points and get them into a dataframe
df_f = ds_scenarios.to_dataframe()
df_f.describe()


# In[25]:


# Lets make a time series plot here so we can see the results.

ax = df_f.plot(figsize=(40, 10), title="NASA GISS", grid=1)
ax.set(xlabel='Date', ylabel='2-m Air Temperature (deg F)')
plt.show()


# In[ ]: