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

# Produces an interactive heatmap with x and y as lattitude and longitude, the colour represents the value of the selected tracer variable, and two sliders control the depth and timestep.
# 
# Modified version of this script:
# http://stackoverflow.com/questions/6697259/interactive-matplotlib-plot-with-two-sliders

# In[3]:


from pylab import *
from matplotlib.widgets import Slider, Button, RadioButtons

import netCDF4 as nc
import numpy as np
import matplotlib.pyplot as plt


# In[4]:


grid_t = nc.Dataset('/data/jpetrie/MEOPAR/NEMO-3.6-code/NEMOGCM/CONFIG/mygyre/EXP00/GYRE_5d_00010101_00011230_grid_T.nc')
#/data/jpetrie/MEOPAR/SalishSea/results/0bc30428-1611-11e6-9e0c-0025909a8461/SS5x5_1h_20041019_20041023_ptrc_T.nc')
#/data/eolson/MEOPAR/SS36runs/runNoBioPISCES/SOG2D_1h_20041019_20041019_grid_T.nc')


lons = grid_t.variables['nav_lon']
lats = grid_t.variables['nav_lat']
keys = grid_t.variables.keys()


# In[5]:


# Looks for all variables with a 4 dimensional input. 
# Should probably check dimensions match as well, but doesn't currently
plottable_vars = list() 
for key in keys:
    if np.size(grid_t.variables[key].shape) == 4:
        plottable_vars.append(key)


# In[6]:


# Global variable used to share information between the update functions about what tracer is currently being plotted
current_var = plottable_vars[0]


# In[31]:


# Function used to update the heatmap when depth/time sliders are moved
def update(val):
    time_index = stime.val
    depth_index = sdepth.val
    # set_array() requires a 1D array. ravel() is used to get around this
    # The [:-1,:-1] is to counteract the weird indexing that pcolormesh does
    #   See: http://stackoverflow.com/questions/29009743/using-set-array-with-pyplot-pcolormesh-ruins-figure
    mesh.set_array(grid_t.variables[current_var][round(time_index),round(depth_index),:,:][:-1,:-1].ravel())
    draw()
    


# In[32]:


def updateVar(variable_name):
    global current_var
    time_index = stime.val
    depth_index = sdepth.val
    # set_array() requires a 1D array. ravel() is used to get around this
    # The [:-1,:-1] is to counteract the weird indexing that pcolormesh does
    #   See: http://stackoverflow.com/questions/29009743/using-set-array-with-pyplot-pcolormesh-ruins-figure
    mesh.set_array(grid_t.variables[variable_name][round(time_index),round(depth_index),:,:][:-1,:-1].ravel())
    current_var = variable_name
    draw()


# In[33]:


fig, ax = plt.subplots(1, 1)
subplots_adjust(left=0.25, bottom=0.25)

mesh = ax.pcolormesh(lats, lons,grid_t.variables[plottable_vars[0]][0,0,:,:])
fig.colorbar(mesh)


# In[34]:


# Dimensions for each of the sliders
axfreq = axes([0.25, 0.1,0.65, 0.03])
axamp  = axes([0.25, 0.15, 0.65, 0.03])
vax = axes([0.025, 0.5, 0.15, 0.15])


stime = Slider(axfreq, 'Time', 1, grid_t.dimensions['time_counter'].size -1)
sdepth = Slider(axamp, 'Depth', 1,grid_t.dimensions['deptht'].size - 1)
radio = RadioButtons(vax, plottable_vars, active=0) # creates a list of checkboxes with the different tracer options


# In[35]:


stime.on_changed(update)
sdepth.on_changed(update)
radio.on_clicked(updateVar)


# In[36]:


show()


# In[ ]:





# In[ ]: