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

# Q: How does one create an Iris cube from scratch and save it as a NetCDF file?

# I was asked specifically if it was possible to make a cube of orthogonal time components (year, month, day), so I'll start there:

# In[3]:


import iris.cube
import iris.coords


# In[4]:


import numpy as np


# First, we need an array of data...

# In[5]:


data = np.empty(shape=[30, 12, 4, 500, 1000], dtype=np.float64)


# Now we create the cube to contain that data:

# In[6]:


cube = iris.cube.Cube(data)
print(cube)


# Clearly this dataset needs some metadata...

# In[7]:


cube.rename('air_temperature')
cube.units = 'kelvin'
cube.attributes['source'] = 'Data from xyz'
print(cube)


# And some coordinate metadata:

# In[8]:


year = iris.coords.DimCoord(points=np.arange(1980, 2010),
                            long_name='year', units='1')
cube.add_dim_coord(year, data_dim=0)
print(cube)


# In[9]:


month = iris.coords.DimCoord(points=np.arange(1, 13),
                             long_name='month', units='1')
cube.add_dim_coord(month, data_dim=1)

day = iris.coords.DimCoord(points=[1, 9, 17, 25],
                           long_name='day', units='1')
cube.add_dim_coord(day, data_dim=2)


# In[10]:


lat = iris.coords.DimCoord(points=np.linspace(-90, 90, 500),
                           standard_name='latitude', units='degrees')
cube.add_dim_coord(lat, data_dim=3)

lon = iris.coords.DimCoord(points=np.linspace(-180, 180, 1000),
                           standard_name='longitude', units='degrees')
cube.add_dim_coord(lon, data_dim=4)


# In[11]:


print(cube)


# We may also want to combine the date information into a single auxilliary time coordinate. This will be a 3d coordinate of shape 30x12x4.

# In[12]:


import cf_units
import netcdftime

time_unit = cf_units.Unit('hours since 2000-01-01')
times = np.empty([30, 12, 4])

for i_y, i_m, i_d in np.ndindex(*times.shape):
    dt = netcdftime.datetime(year.points[i_y], month.points[i_m], day.points[i_d])
    times[i_y, i_m, i_d] = time_unit.date2num(dt)

time_coord = iris.coords.AuxCoord(times, standard_name='time', units=time_unit)
cube.add_aux_coord(time_coord, data_dims=[0, 1, 2])


# In[13]:


print(cube)


# In[14]:


iris.save(cube, '/data/pelson/delme.nc')


# In[15]:


get_ipython().system('ncdump -h /data/pelson/delme.nc')


# ## Single time dimension approach
# 
# Instead of having (orthogonal) dimensions for each of the time components, we can combine them into a single dimension, and use the API to pick out the values of interest:

# In[16]:


data = np.empty(shape=[30 * 12 * 4, 500, 1000], dtype=np.float64)

cube2 = iris.cube.Cube(data)
print(cube2)


# In[17]:


cube2.rename('air_temperature')
cube2.units = 'kelvin'
cube2.attributes['source'] = 'Data from xyz'

lat = iris.coords.DimCoord(points=np.linspace(-90, 90, 500),
                           standard_name='latitude', units='degrees')
cube2.add_dim_coord(lat, data_dim=1)

lon = iris.coords.DimCoord(points=np.linspace(-180, 180, 1000),
                           standard_name='longitude', units='degrees')
cube2.add_dim_coord(lon, data_dim=2)

print(cube2)


# In[18]:


time_unit = cf_units.Unit('days since 2000-01-01')
times = []
for year in np.arange(1980, 2010):
    for month in np.arange(1, 13):
        for day in [1, 9, 17, 25]:
            times.append(netcdftime.datetime(year, month, day))

time = iris.coords.DimCoord(points=time_unit.date2num(times),
                            standard_name='time', units=time_unit)
cube2.add_dim_coord(time, data_dim=0)
print(cube2)


# But how do we access useful things, like the first day of each month? First, let's add some extra coordinate information:

# In[19]:


import iris.coord_categorisation as iccat

iccat.add_day_of_month(cube2, 'time')
iccat.add_month(cube2, 'time')
iccat.add_year(cube2, 'time')
print(cube2)


# What do these coordinates look like, and how can we use them?

# In[20]:


print(cube2.coord('day_of_month').points[0:10])


# In[21]:


print(cube2.extract(iris.Constraint(day_of_month=1)))


# In[22]:


print(cube2.coord('month').points[0:10])


# In[23]:


print(cube2.extract(iris.Constraint(month='Jan')))


# In[24]:


print(cube2.coord('year').points[0:10])


# In[25]:


print(cube2.extract(iris.Constraint(year=1980)))


# We can combine this together to get the first of the month for all March, April & May in the 1980s...

# In[26]:


first_mam_80s = iris.Constraint(day_of_month=1, month=['Mar', 'Apr', 'May'],
                                year=lambda c: 1980 <= c.point < 1990)
print(cube2.extract(first_mam_80s))
print(cube2.extract(first_mam_80s).coord('time'))