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

# I've implemented an Interpolator class in ``iris.analysis.interpolator`` which allows one to define simple subclasses for the implementation of various interpolation schemes. 

# In[1]:


from iris.analysis.interpolator import Interpolator


# There is already a specialised version of this which limits the interpolation to 1D monotonic coordinates, named ``RectilinearInterpolator``.

# In[2]:


from iris.analysis.interpolator import RectilinearInterpolator


# With this we can implement some of the simpler interpolation schemes, such as NDLinear interpolation:

# In[3]:


class LinearInterpolator(RectilinearInterpolator):
    def _build_interpolator(self, coord_points, data,
                            extrapolation_mode=None):
        """
        Given the coordinate points of the source grid, and
        some data which is of the appropriate shape for those coord
        points, construct the interpolator.

        """
        from iris.experimental.regrid import _RegularGridInterpolator
        self._interpolator = _RegularGridInterpolator(coord_points, data,
                                                      fill_value=None,
                                                      bounds_error=False)
        return self._interpolator

    def _interpolate_data_at_coord_points(self, data, coord_points):
        """
        Given some coord points (already prepared into the correct
        form) and the data to interpolate, do the interpolation.
        
        """
        # Update the interpolator's data to the data passed through.
        self._interpolator.values = data
        return self._interpolator(coord_points)


# In[4]:


import iris

fname = iris.sample_data_path('air_temp.pp')
temperature = iris.load_cube(fname)
print temperature


# We construct the Interpolator with the cube to interpolate, and the interpolation coordinates:

# In[5]:


interpolator = LinearInterpolator(temperature, ['latitude', 'longitude'])


# With this, we can use the orthogonal_cube method:

# In[6]:


print interpolator.orthogonal_cube([['latitude', 51],
                                    ['longitude', 0]])


# Notice how we didn't need to override the ``orthogonal_cube`` method - it could potentially be an external interface if there is commonality with the regrid interface. Perhaps even:
# 
# ```
# InterpolatorRegridder(Regridder):
#    def __init__(self, interpolator):
#        ...
# 
# ```
# 
# The implementation of the Interpolator class wouldn't need to change much to accommodate this.

# ## Other interpolators
# 
# So the real benefit of building the Interpolator class is that it is easy to implement other interpolation schemes with little code (and virtually no repeating of validation):

# In[7]:


class Spline1DInterpolator(RectilinearInterpolator):
    def __init__(self, *args, **kwargs):
        super(Spline1DInterpolator, self).__init__(*args, **kwargs)
        assert len(self.coord_dims) == self.ndims == 1
    
    def _build_interpolator(self, coord_points, mock_data,
                            extrapolation_mode=None):
        # N.B. There is no need to build an interpolator here - spline is
        # heavily dependent on the data and so data cannot be switched in and
        # out efficiently. As a result, we need to construct the interpolators
        # as we go. So just store the coord points for construction of the
        # interpolator when we need it.
        self._data_coord_points = coord_points

    def _interpolate_data_at_coord_points(self, data, coord_points):
        assert data.ndim == 1
        import scipy.interpolate
        interp = scipy.interpolate.InterpolatedUnivariateSpline(
                                                self._data_coord_points, data)
        return interp(coord_points[:, 0])


# In[8]:


import iris
e1 = iris.load_cube(iris.sample_data_path('E1_north_america.nc'))[:, :, 20]
print e1


# In[9]:


times_to_interpolate = np.linspace(e1.coord('time').points[4],
                                   e1.coord('time').points[7],
                                   30)
interpolator = Spline1DInterpolator(e1, ['time'])

interpolated_cube = interpolator.orthogonal_cube([['time', times_to_interpolate]])

print interpolated_cube


# In[10]:


import matplotlib.pyplot as plt

plt.plot(e1.coord('time').points[4:8],
         e1.data[4:8, 20], 'xr', label='actual')

plt.plot(times_to_interpolate, interpolated_cube[:, 20].data, label='Interpolated')
plt.legend()
plt.show()


# ### Triangulation based interpolation

# In[11]:


from numpy.lib.stride_tricks import as_strided
import scipy.interpolate


class TriangulatedLinearInterpolator(Interpolator):
    def _build_interpolator(self, coord_points, data,
                            extrapolation_mode=None):
        data = as_strided(np.array(0, dtype=float),
                               shape=(np.product(data.shape), ), strides=[0])
        # Flatten the coord points for this interpolator.
        coord_points = np.array([points.flatten() for points in coord_points]).T
        return scipy.interpolate.LinearNDInterpolator(coord_points, data) 
        
    def interpolated_dtype(self, dtype):
        return np.float

    def _interpolate_data_at_coord_points(self, data, coord_points):
        # Cast the data to float64 - that is all it can handle.
        self._interpolator.values = data.reshape(-1, 1).astype(np.float64)
        return self._interpolator(coord_points)


# In[12]:


fname = '/data/local/dataZoo/NetCDF/ORCA1/NEMO_ORCA1_CF.nc'
sst = iris.load_cube(fname, 'sea_surface_temperature')
sst.data[sst.data == 0] = np.nan 
sst.attributes.clear()
print sst


# In[13]:


interpolator = TriangulatedLinearInterpolator(sst, ['latitude', 'longitude'])
arctic_circle = interpolator.orthogonal_cube([['longitude', np.linspace(-180, 180, 1000)],
                                              ['latitude', 60]])

print arctic_circle


# In[14]:


import iris.quickplot as qplt

qplt.plot(arctic_circle[0].coord('longitude'), arctic_circle[0])
plt.show()


# In[15]:


regridded = interpolator.orthogonal_cube([['longitude', np.linspace(-180, 180, 1000)],
                                          ['latitude', np.linspace(-90, 90, 500)]])
qplt.contourf(regridded[0], coords=['longitude', 'latitude'])
plt.gca().coastlines()
plt.show()


# In[15]: