#!/usr/bin/env python

################################################################################################################
#### This is a python script that takes my numu.p infile
#### and plots a mc energy 1D histogram and a mc energy 
#### vs mc zenith 2D histogram. 
####
#### NOTE: Constants and Functions defined here are just
####       my way to plot. You are free to change whatever
####       you want and plot however you want. 
####
####       All you need to have are:
####
####          --- 1D histogram of true and reco energy plotted together 
####                    -> errorbars
####                    -> printed text
####                    -> colorblind proof
####          --- ratios of true and reco energy histograms & the errorbars
####          --- 2D histogram of true cos_zenith VS true_energy
####                    -> printed text
####                    -> colorblind proof
####                    -> colorbar
####          --- appropriate x, y axes / title / legend / etc for all plots
####
#### python command line:
#### python plot_distributions_template.py --outfile "out file name (.pdg/.png/etc)" --infile numu.p
################################################################################################################

###########################################
#### import tools & set up basic variables
###########################################
#### Basic tools
from optparse import OptionGroup, OptionParser 
import numpy as np
import cPickle
import matplotlib
matplotlib.use('Agg')

import matplotlib.pyplot as plt
#### Set latex fonts
plt.rc ('text', usetex=True)
plt.rc ('font', family='serif')
plt.rc ('font', serif='Computer Modern Roman')
#### Format major ticksizes
plt.rcParams['xtick.major.size'] = 3
plt.rcParams['ytick.major.size'] = 3
#### Customize multiple subplots
import matplotlib.gridspec as gridspec
#### Customize prints on plots
from mpl_toolkits.axes_grid.anchored_artists import AnchoredText

# ==============================================================
# = Constants for plotting 
# ==============================================================
#### Set dimensions of output plot (in inches)
f_height, f_width = 5.5, 15
#### Set limits/labels/titles/bins for 1D histogram
elimit, elabel = (1,3), r'log E (GeV)'
etitle  = r'log energy distributions'
ebins = 20
#### Set limits/labels/bins/vmax for 2D histogram
labels = {'e':r'log true E (GeV)', 'z':r'cos true $\theta$' }
pedges = {'e':np.linspace(0.70,1.90,11), 'z':np.linspace(-1.,1.,11)[::-1] }  ## for plotting histograms
hedges = {'e':10**pedges['e'], 'z':np.arccos(np.linspace(-1.,1.,11))[::-1] } ## for making histograms
vmax = 700. ## color bar max limits

# ==============================================================
# = Other constants
# ==============================================================
#### Set seconds per year ( expected_num_in_n_years = weights * n_years * seconds_per_year )
seconds_per_year = 3600.*24.*365.
#### Set default paths ( will be parsed from user )
infile = '/home/icecube/i3/data/plotting_lesson/numu.p'
outfile = '/data/i3home/elims/bootcamp/test.pdf'

###########################################
#### parse user's options
###########################################
usage = "%prog [options]"

parser = OptionParser(usage=usage)
parser.add_option( "-o", "--outfile", type="string", default=outfile,
                   help = "output file (path/name.pdf)")
parser.add_option( "-i", "--infile", type="string", default=infile,
                   help = "input pickled dictionary")
(options, args) = parser.parse_args()

outfile = options.outfile
infile = options.infile

###########################################
#### Functions to build 1D histograms
####
#### hints: 
####   --- np.histogram (data)
####   --- plot ( xvalues, histogram )
####   --- errorbar ( xvalues, histogram )
###########################################
def plot_1D_histo ( axis, xvalues, weights, xmin, xmax, nbin, color, label, linestyle ):

    ''' Plot 1D histogram + errorbars '''

    #### build & plot histograms

    #### calculate & plot errorbars

    #### print statement 'IceCube Preliminary'

    return

def plot_1D_ratio ( axis, histos, edges, xmin, xmax, nbin ):

    ''' Plot ratio of 2 1D histograms + errorbars '''

    #### calculate ratios & their errors

    #### plot ratios & their errorbars

    return

def format_1D_axis ( axis, xmin, xmax, xLabel=r'', yLabel=r'', Title=r'', ratio_plot=False ):

    ''' format 1D energy histogram axes '''

    #### format y axis: label, limit, tick size, grid lines, invisible tick/label?

    #### format x axis: label, limit, tick size, grid lines, invisible tick/label?

    #### set title and legend

    return 

###########################################
#### Functions to build 2D histograms
####
#### hints: 
####   --- np.histogram2d (data)
####   --- np.histogramdd (data)
####   --- pcolormesh ( x bin edges, y bin edges, histogram )
####   --- annotate ( text )
###########################################
def plot_2D_histo ( energy, zenith, weight ):
    
    ''' Plot log_energy VS cos_zenith histogram '''
    
    #### build 2d histogram 

    #### plot 2d histogram

    #### print number of events per bin

    return

def format_2D_axis ( axis, Title, grid_color='black' ):

    ''' Format log_energy VS cos_zenith histogram axes '''

    #### set x, y labels and title

    #### set x, y ticks

    #### set grid lines

    #### set x, y limits

    return

#############################################
#### Everything starts here ! :)
##############################################
#### open pickled infile
with open (infile, 'rb') as f:
    numu = cPickle.load (f)
f.close ()

#### to access data through numu
## log_true_energy = np.log10 (numu['true_e'])
## log_reco_energy = np.log10 (numu['reco_e'])
## cos_true_zenith = np.cos (numu['true_z'])
## weight = numu['weight'] * seconds_per_year

######################################### define canvas subplots ########################################
## h = plt.figure ( figsize=(f_width, f_height) )
## gs = gridspec.GridSpec ( n_rows, m_columns )

########################################### plot 1D histogram ###########################################
## ax0 = h.add_subplot (gs[0])

## 1) reco & mc histogram & errors

## 2) ratios & their errors

########################################### plot 2D histogram ###########################################
## ax1 = h.add_subplot (gs[1])

## 3) true coszen VS true energy histogram

## 4) colorbar

######################################### done plotting & save ##########################################
## plt.suptitle ( your_title )
## h.savefig ( outfile )
## plt.close ( 'all' )

#### set super title; save figure; close all canvas