import numpy as np
import pandas as pd
from matplotlib import pyplot as plt

"""
Script originally written by Tianlu Yuan
Recall that timing is one of the most important factors

"""

def fhit_vs_z(dinfo, nstr=4, tvtx=0, title=False):
    """plot fhit vs z for the *nstr* brightest doms in *dinfo*, where
    *dinfo* is the pickled dom data from extract.py
    """
    if isinstance(dinfo, str):
        df = pd.read_pickle(dinfo)
    else:
        df = dinfo
    print df
    plt.figure()
    topstrs = df.groupby('string').sum().sort_values('qdom', ascending=False)[:nstr]

    color_cycle = plt.rcParams['axes.prop_cycle'].by_key()['color']
    exp_labels = []
    exp_styles = []
    for idx, string in enumerate(topstrs.index):
        df_string = df.loc[df['string'] == string].sort_values('zdom')
        qsct = plt.scatter(df_string['fhit']-tvtx, df_string['zdom'], s=df_string['qdom']/100,
                           label='Str. {}'.format(string))
        if 'phit' in df_string:
            plt.plot(df_string['phit']-tvtx, df_string['zdom'], label='',
                     color=color_cycle[idx%len(color_cycle)])
            plt.fill_betweenx(df_string['zdom'], df_string['phit_low']-tvtx,
                              df_string['phit_high']-tvtx, alpha=0.3,
                              label='', color=color_cycle[idx%len(color_cycle)])
            if idx == 0:
                exp_labels.append('Predicted')
                exp_styles.append('-')
        if 'mnphit' in df_string:
            plt.plot(df_string['mnphit']-tvtx, df_string['zdom'], ':', label='',
                     color=color_cycle[idx%len(color_cycle)])
            if idx == 0:
                exp_labels.append('mnphit')
                exp_styles.append(':')
        if 'mlphit' in df_string:
            plt.plot(df_string['mlphit']-tvtx, df_string['zdom'], '--', label='',
                     color=color_cycle[idx%len(color_cycle)])
            if idx == 0:
                exp_labels.append('mlphit')
                exp_styles.append('--')

    for label, linestyle in zip(exp_labels, exp_styles):
        plt.axvline(np.nan, color='k', linestyle=linestyle, label=label)
    if tvtx ==0:
        plt.xlabel(r'$t_{DOM}^{first}\ [\mu s]$', fontsize=14)
    else:
        plt.xlim(0.2, 2.25)
        plt.xlabel(r'$t-t_0 [\mu s]$', fontsize=14)
    plt.text(0,0.95,'IceCube Preliminary', fontsize=14,
             transform=plt.gca().transAxes, color='r')
    if title:
        plt.title('{}.{}'.format(df['run'].iloc[0], df['event'].iloc[0]))
    plt.ylabel(r'z', fontsize=14)
    plt.xticks(fontsize=14)
    plt.yticks(fontsize=14)
    plt.legend()
    plt.tight_layout()
    plt.savefig('fhit_vs_z.pdf')


if __name__=='__main__':
    import argparse
    parser = argparse.ArgumentParser(description='Plot fhit vs z')
    parser.add_argument('dinfo', type=str, help='input dom-info pickled file')
    parser.add_argument('-n', type=int, default=4, help='number of strings to plot (brightest to least bright)')
    parser.add_argument('-t', type=float, default=0, help='vertex time offset')
    args = parser.parse_args()
    fhit_vs_z(args.dinfo, args.n, args.t, True)