#!/usr/bin/env python

import os
from collections import OrderedDict
import cPickle as pickle
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib.cm import get_cmap
from matplotlib import style
from scipy import stats
from scipy import integrate

def plot_monthly_response(norm, pert):

    plot_grid = gridspec.GridSpec(4, 1, hspace=0.1)

    ax1 = plt.subplot(plot_grid[0])
    ax2 = plt.subplot(plot_grid[1])
    ax3 = plt.subplot(plot_grid[2])
    ax4 = plt.subplot(plot_grid[3])

    # Stomatal conductance
    ax1.plot(norm["Gtree"].values)
    ax1.plot(pert["Gtree"].values)

    # Leaf transpiration
    ax2.plot(norm["Etree"].values)
    ax2.plot(pert["Etree"].values)

    # Leaf assimilation
    ax3.plot(norm["Atree"].values)
    ax3.plot(pert["Atree"].values)

    ax4.plot(norm["LAItree"].values)
    ax4.plot(pert["LAItree"].values)
    ax4.plot(norm["LAIgrass"].values)
    ax4.plot(pert["LAIgrass"].values)

    plt.show()

    return 1

def main():

    data_dict = pickle.load(open(PKLPATH, 'rb'))

    year_agg = lambda x: x.groupby(level=['month', 'hour']).mean()

    data_mean_year = [year_agg(df) \
                      for df in OrderedDict(data_dict).values()]

    # **FOR LOOP WILL GO HERE
    plot_monthly_response(data_mean_year[3], data_mean_year[6])

    return 1

if __name__ == "__main__":

    FILEPATH = "~/Savanna/Data/HowardSprings_IAV/pickled/agg/mean_monthly_leaf.pkl"
    PKLPATH = os.path.expanduser(FILEPATH)

    main()

[((17, 16, 17, 51), 'matplotlib.gridspec.GridSpec', 'gridspec.GridSpec', (), '', True, 'import matplotlib.gridspec as gridspec\n'), ((19, 10, 19, 35), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(19, 22, 19, 34): 'plot_grid[0]'}, {}), '(plot_grid[0])', True, 'import matplotlib.pyplot as plt\n'), ((20, 10, 20, 35), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(20, 22, 20, 34): 'plot_grid[1]'}, {}), '(plot_grid[1])', True, 'import matplotlib.pyplot as plt\n'), ((21, 10, 21, 35), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(21, 22, 21, 34): 'plot_grid[2]'}, {}), '(plot_grid[2])', True, 'import matplotlib.pyplot as plt\n'), ((22, 10, 22, 35), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(22, 22, 22, 34): 'plot_grid[3]'}, {}), '(plot_grid[3])', True, 'import matplotlib.pyplot as plt\n'), ((41, 4, 41, 14), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((62, 14, 62, 42), 'os.path.expanduser', 'os.path.expanduser', ({(62, 33, 62, 41): 'FILEPATH'}, {}), '(FILEPATH)', False, 'import os\n'), ((52, 32, 52, 54), 'collections.OrderedDict', 'OrderedDict', ({(52, 44, 52, 53): 'data_dict'}, {}), '(data_dict)', False, 'from collections import OrderedDict\n')]