%matplotlib inline
import numpy as np
import scipy.io as spio
import matplotlib.pyplot as plt

matfile = '../../data/CA2009.mat'
matdata = spio.loadmat(matfile)
print matdata.keys()

def print_mat_nested(d, indent=0, nkeys=0):
    """Pretty print nested structures from .mat files   
    Inspired by: `StackOverflow <http://stackoverflow.com/questions/3229419/pretty-printing-nested-dictionaries-in-python>`_
    """
    
    # Subset dictionary to limit keys to print.  Only works on first level
    if nkeys>0:
        d = {k: d[k] for k in d.keys()[:nkeys]}  # Dictionary comprehension: limit to first nkeys keys.

    if isinstance(d, dict):
        for key, value in d.iteritems():         # iteritems loops through key, value pairs
          print '\t' * indent + 'Key: ' + str(key)
          print_mat_nested(value, indent+1)

    if isinstance(d,np.ndarray) and d.dtype.names is not None:  # Note: and short-circuits by default
        for n in d.dtype.names:    # This means it's a struct, it's bit of a kludge test.
            print '\t' * indent + 'Field: ' + str(n)
            print_mat_nested(d[n], indent+1)

print_mat_nested(matdata, nkeys=10)

t = matdata['T_CA065_090603_010']   # Just look at one set of data

# Here we define a simple function to print the data type and shape.  It just saves us some repetition.
def dtype_shape_str(x):
    """ Return string containing the dtype and shape of x."""
    return str(x.dtype) + " " + str(x.shape)

print "t['data'] is a: " + dtype_shape_str(t['data'])
print "t['data'][0,0] is a: " + dtype_shape_str(t['data'][0,0])
print "t['data'][0,0]['time'] is a: " + dtype_shape_str(t['data'][0,0]['time'])
print "t['data'][0,0]['time'][0,0] is a: " + dtype_shape_str(t['data'][0,0]['time'][0,0])

time = t['data'][0,0]['time'][0,0]
temp = t['data'][0,0]['temp'][0,0]
plt.plot(time,temp,'ko')
plt.xlabel('time')
plt.ylabel('temp')

def loadmat(filename):
    '''
    this function should be called instead of direct spio.loadmat
    as it cures the problem of not properly recovering python dictionaries
    from mat files. It calls the function check keys to cure all entries
    which are still mat-objects
    
    from: `StackOverflow <http://stackoverflow.com/questions/7008608/scipy-io-loadmat-nested-structures-i-e-dictionaries>`_
    '''
    data = spio.loadmat(filename, struct_as_record=False, squeeze_me=True)
    return _check_keys(data)

def _check_keys(dict):
    '''
    checks if entries in dictionary are mat-objects. If yes
    todict is called to change them to nested dictionaries
    '''
    for key in dict:
        if isinstance(dict[key], spio.matlab.mio5_params.mat_struct):
            dict[key] = _todict(dict[key])
    return dict        

def _todict(matobj):
    '''
    A recursive function which constructs from matobjects nested dictionaries
    '''
    dict = {}
    for strg in matobj._fieldnames:
        elem = matobj.__dict__[strg]
        if isinstance(elem, spio.matlab.mio5_params.mat_struct):
            dict[strg] = _todict(elem)
        else:
            dict[strg] = elem
    return dict

matdata = loadmat(matfile)

print_mat_nested(matdata, nkeys=6)

# Now that the data has a hierarchy, we can start combining things.

def combine_lpdata(d):
    """Combine the low-pass and regular data to a single dictionary.  
    The input is the dictionary that contains the 'data' and 'lpdata' keys.
    
    The function works in-place, but returns a reference to the dictionary as well."""
    
    # For regular data, just pop it off the dictionary, keeping the same name.
    # The `pop` method returns the value of key `k` and removes it from the dictionary at the same time.
    for k in d['data'].keys():
        d[k] = d['data'].pop(k)
    d.pop('data')

    # For low-passed data, prepend 'lp' to the front of the key.
    for k in d['lpdata'].keys():
        d['lp'+k] = d['lpdata'].pop(k)
    d.pop('lpdata')

    return d

# Test the function
from copy import deepcopy
test = deepcopy(matdata['T_CA042_090603_010'])  # Copy the data for the test.
print 'keys before:', test.keys()
combine_lpdata(test)
print 'keys after:', test.keys()  # Note how we print the same thing, but the keys have changed.


# Now we'll do it on all of matdata
# The next line ignores the special variables starting with '__'
# You can't have matlab variables starting with '__' so this should only remove extra keys created by loadmat.
variables = [k for k in matdata.keys() if not k.startswith('__')]   # This is a list comprehension. 

# Records is now just a list, where each element is a string corresponding to a key of the dictionary matdata.
print variables

# We can loop through all of it, running our function
for v in variables:
    combine_lpdata(matdata[v])

# Or we can loop through just part of it to print the results
for v in variables[:4]:
    print matdata[v]

# The split method generates a list for each name, while the list comprehension 
# generates a super-list containing all the sub-lists.
print [v.split('_') for v in variables]  
v_array = np.array([v.split('_') for v in variables])

# Now we find the unique stations, parameters, and depths
parameters = np.unique(v_array[:,0])
stations = np.unique(v_array[:,1])
depths = np.unique(v_array[:,3])

print 'Parameters: ', parameters
print 'Stations: ', stations
print 'Depths: ', depths  # Note that they are strings

# First, we'll combine the data sets for each unique combination of instruments, station, and depth.
matdata2 = {}  # Creates an empty dictionary

for k,v in matdata.iteritems():  # Here's way to loop over key,value pairs of dictionaries.
    if  k.startswith('__'):  # Skip the special keys
        continue             # Continue goes to the next loop iteration.
    param, station, date, depth = k.split('_')  # This splits up the string, assigning each of the parts to different variables.
    
    # The if statements here just check to make sure that level of the dictionary has already been created.
    # If all levels are present, we can just assign a new key.  If not, we have to create a new dictionary
    # at that level.
    if station in matdata2:
        if param in matdata2[station]:
            if depth in matdata2[station][param]:
                matdata2[station][param][depth][date] = v  # We don't need to check depth b/c we're assigning it.
            else:
                matdata2[station][param][depth] = {date: v}
        else:
            matdata2[station][param] = {depth: {date: v}}
    else:
        matdata2[station] = {param: {depth: {date: v}}}

matdata3 = deepcopy(matdata2)

for station in matdata3.itervalues():  # Loops over values (in this case, the sub-dictionaries)
    for params in station.itervalues():  # OK, the param dict is the one we're going to modify
        for depth in params.itervalues():
            for date in sorted(depth.keys()):  # Sort so that the dates order properly
                for v in depth[date].keys():
                    if v not in depth:
                        depth[v] = [depth[date].pop(v)]  # Initialize the list with the first value
                    else:
                        depth[v].append(np.array(np.nan)) # Put nans between lists to break-up timeseries
                        depth[v].append(depth[date].pop(v))
                depth.pop(date)
            
            # Now we have to convert our lists of ndarrays into a single array
            for v in depth.keys():  # depths no longer contains the dates, but the variables now
                depth[v] = np.hstack(tuple(depth[v])) # hstack expects a tuple, not a list

from matplotlib.dates import num2date   # This converts matplotlib date-nums to Python date objects

for name, station in matdata3.iteritems():
    T = station['T']   # Just get the temperatures
    station_mint = np.inf
    station_maxt = -np.inf
    print name + ': '
    
    # Loop through depths find start/stop times and mean dt
    for depth in sorted(T.keys()):
        time = T[depth]['time']
        mint = np.floor(np.nanmin(time))
        maxt = np.ceil(np.nanmax(time))
        station_mint = min(mint,station_mint)  # Note difference in min functions (Also, np.fmin)
        station_maxt = max(maxt,station_maxt)  # Note difference in min function
        
        dt   = np.nanmean(np.diff(time)) *24*60  # Convert days to minutes
        print u'\t {}: {} to {}.  ∆t = {:f} min'.format(depth, num2date(mint).strftime('%Y-%m-%d'),
                                                               num2date(maxt).strftime('%Y-%m-%d'), dt)

    print '{}: {} to {}'.format(name, num2date(station_mint).strftime('%Y-%m-%d'), 
                                      num2date(station_maxt).strftime('%Y-%m-%d'))

# Finally, we'll interpolate to a common timestep so that each station and depth can use the same timestep
matdata4 = deepcopy(matdata3)
from matplotlib.dates import datestr2num

# Create the time array
time =np.arange(datestr2num('2010-04-30'), datestr2num('2010-10-09'), 2.0/24./60.)  # 2 minute intervals
nt = len(time)

for key, station in matdata4.iteritems():
    T = station['T']  # Only looping through temps for example
    depths = sorted(T.keys())
    nd = len(depths)
    data = np.empty((nd,nt))  # Pre-allocate space for the array.  Note the double (()) - the argument is a tuple!
    lpdata = np.ones((nd,nt))  # Pre-allocate, slightly slower.

    # Note that in numpy, interp is (xi, x, y) vs. matlab (x, y, xi) 
    # where x,y is the original and xi is where to interpolate the y to.
    for i,d in enumerate(depths):
        data[i,:] = np.interp(time, T[d]['time'], T[d]['temp'])
        lpdata[i] = np.interp(time, T[d]['time'], T[d]['lptemp'])  # NB: for 2D arrays: [i]  = [i][:] = [i,:]
    
    # This is one way to convert a list of strings to floating-point array
    depth_as_floats = np.array([float(i) for i in depths])

    # Now we reassign the data
    matdata4[key] = {'time':time, 'temp':data, 'lptemp':lpdata, 'depth':depth_as_floats}

from matplotlib.dates import MonthLocator

f,ax = plt.subplots(3,1)
f.set_size_inches(11,8.5)

for a, s in zip(ax, sorted(matdata4.keys())):
    a.plot(matdata4[s]['time'], matdata4[s]['lptemp'].T,'-')
    a.set_ylabel('T [$^\circ$C]')
    a.set_title(s, loc='left')
    # This sticks the legend outside of the axes. Note that if you don't adjust
    # the axis width, the legend will be 'off the page'.
    L =a.legend(matdata4[s]['depth'], bbox_to_anchor=(1.05, 1), loc=2) 
    L.set_title('Depth [m]')
    f.subplots_adjust(right=0.8) # adjust axis width
    
ax[-1].xaxis_date()
for a in ax[:-1]:
    a.set_xticks(ax[-1].get_xticks())
    a.set_xticklabels([])