# -*- coding: utf-8 -*-
# @Author: Elena Graverini
# @Date:   2017-03-16 11:20:45
# @Last Modified by:   Elena Graverini
# @Last Modified time: 2017-05-10 16:52:40
# import matplotlib
from __future__ import division
import sys, os, gc
import matplotlib.pyplot as plt
from eosdata import MCMCDataFile  # , Plotter1D
import pickle


def main():
    input_file = "constantFF/samples-mu-2.hdf5"
    save_name = 'test-output.pdf'
    if len(sys.argv) > 1:
        input_file = str(sys.argv[1])
    else:
        print('Warning: no input / output file names specified.\nUsing %s and %s.' % (input_file, save_name))
    if len(sys.argv) > 2:
        save_name = str(sys.argv[2])
    pklfilename = save_name.replace('.pdf', '.pkl')
    suffix = 'prescaled_10'
    if not os.path.isfile(pklfilename):
        datafile = MCMCDataFile(input_file)
        # print(datafile.parameters[0])
        gc.collect()
        data = datafile.data()[:, 0]
        d = {}
        reduced_data = data[::10]
        d['reduced_data'] = reduced_data
        n, bins, patches = plt.hist(reduced_data, bins=100, normed=True, alpha=0.3)
        bin_centers = (bins[:-1] + bins[1:]) / 2
        d['bin_contents'] = n
        d['bin_centers'] = bin_centers
        pickle.dump(d, open(pklfilename, 'wb'))
    else:
        d = pickle.load(open(pklfilename, 'rb'))
        plt.scatter(d['bin_centers'], d['bin_contents'])
        suffix += '_scatter'
    plt.savefig(save_name.replace('.pdf', '_%s.pdf' % suffix))
    # # plt.figure()
    # p1d = Plotter1D(datafile, save_name)
    # p1d.histogram(0, xmin=None, xmax=None, kde=None)


def make_pickle(pkl_file, input_file):
    datafile = MCMCDataFile(input_file)
    gc.collect()
    data = datafile.data()[:, 0]
    d = {}
    reduced_data = data[::10]
    d['reduced_data'] = reduced_data
    n, bins, patches = plt.hist(reduced_data, bins=100, normed=True, alpha=0.3)
    bin_centers = (bins[:-1] + bins[1:]) / 2
    d['bin_contents'] = n
    d['bin_centers'] = bin_centers
    pickle.dump(d, open(pkl_file, 'wb'))
    plt.savefig(pkl_file.replace('.pkl', '.pdf'))


def superimpose(files, datafiles, labels, outname, lepton):
    # colors = ['#0078FF', '#FF6600', '#0AAFB6', '#FF3333', '#0000FF', '#00CC00', '#BF8040', '#FF33CC', '#FF7733']
    colors = ['#0078FF', '#FF6600', '#0AAFB6', '#FF3333', '#0000FF', '#00CC00', '#FF33CC']
    styles = ['-'] * len(colors)
    if len(files) > len(colors):
        styles.extend(['--'] * len(colors))
        colors.extend(colors)
    ds = []
    for i, f in enumerate(files):
        try:
            ds.append(pickle.load(open(f, 'rb')))
        except IOError:
            make_pickle(f, datafiles[i])
            ds.append(pickle.load(open(f, 'rb')))
    ax = plt.subplot(111)
    for i in range(len(files)):
        ax.plot(ds[i]['bin_centers'], ds[i]['bin_contents'], styles[i], label=labels[i], color=colors[i], linewidth=2.0)
    ax.set_xlabel('$q^2$ (GeV)')
    ax.set_ylabel(r'$d\Gamma\left(\Lambda_b\to\Lambda_c^\star\{lepton}\nu\right) / dq^2$'.format(lepton=lepton))
    # Shrink current axis by 10%
    box = ax.get_position()
    ax.set_position([box.x0, box.y0, box.width * 0.9, box.height])
    # plt.legend(loc='best')
    # Put a legend to the right of the current axis
    ax.legend(loc='center left', fancybox=True, shadow=True, bbox_to_anchor=(1, 0.5))
    plt.savefig(outname)
    plt.close("all")


def makeFileFromMarzia(x, y, fname):
    ds = {
        'bin_centers': x,
        'bin_contents': y,
    }
    if not os.path.isfile(fname):
        os.system('mkdir -p %s' % os.path.dirname(fname))
        pickle.dump(ds, open(fname, 'wb'))


if __name__ == '__main__':
    # main()
    marzia_f = 'Marzia/Marzia_mu_constantFF.pkl'
    marzia_ft = 'Marzia/Marzia_tau_constantFF.pkl'
    mx = [0.2, 0.4, 0.80, 1, 1.5, 2, 2.5, 3, 3.5, 4, 4.5, 5, 5.5, 6, 6.5, 7, 7.5, 8, 8.5]
    my = [0.147214, 0.14805, 0.147707, 0.147292, 0.145776, 0.143623, 0.140846, 0.137443, 0.133406, 0.12872, 0.123359, 0.117288, 0.110451, 0.102761, 0.0940879, 0.0842141, 0.0727531, 0.0589096, 0.0404804]
    mxt = [3, 3.5, 4, 4.5, 5, 5.5, 6, 6.5, 7, 7.5, 8, 8.5]
    myt = [0.00641833, 0.0194386, 0.0789258, 0.140122, 0.189971, 0.225482, 0.24678, 0.254695, 0.249817, 0.231917, 0.198965, 0.143346]
    # makeFileFromMarzia(mx, my, marzia_f)
    # makeFileFromMarzia(mxt, myt, marzia_ft)
    files_mu = ['constantFF/mcmc_gamma_muons.pkl',
                '1pole/mcmc_gamma_muons_1pole.pkl',
                '1pole_x_sqrtlambda/mcmc_gamma_muonss_1pole_x_sqrtlambda.pkl',
                # marzia_f,
                # 'linear/mcmc_gamma_muons_linear_rho-0_35.pkl',
                # 'linear/mcmc_gamma_muons_linear_rho-0_15.pkl',
                # 'linear/mcmc_gamma_muons_linear_rho-0_05.pkl',
                # 'linear/mcmc_gamma_muons_linear_rho0_00.pkl',
                # 'linear/mcmc_gamma_muons_linear_rho0_05.pkl',
                # 'HQT/mcmc_gamma_muons_HQT_rho-0_35_rho3b0_00.pkl',
                # 'HQT/mcmc_gamma_muons_HQT_rho-0_15_rho3b0_00.pkl',
                'HQT/mcmc_gamma_muons_HQT_rho-0_05_rho3b0_00.pkl',
                'HQT/mcmc_gamma_muons_HQT_rho0_00_rho3b0_00.pkl',
                'HQT/mcmc_gamma_muons_HQT_rho0_05_rho3b0_00.pkl',
                # 'HQT/mcmc_gamma_muons_HQT_rho-0_35_rho3b-0_05.pkl',
                # 'HQT/mcmc_gamma_muons_HQT_rho-0_15_rho3b-0_05.pkl',
                'HQT/mcmc_gamma_muons_HQT_rho-0_05_rho3b-0_05.pkl',
                'HQT/mcmc_gamma_muons_HQT_rho0_00_rho3b-0_05.pkl',
                'HQT/mcmc_gamma_muons_HQT_rho0_05_rho3b-0_05.pkl',
                # 'HQT/mcmc_gamma_muons_HQT_rho-0_35_rho3b0_05.pkl',
                # 'HQT/mcmc_gamma_muons_HQT_rho-0_15_rho3b0_05.pkl',
                'HQT/mcmc_gamma_muons_HQT_rho-0_05_rho3b0_05.pkl',
                'HQT/mcmc_gamma_muons_HQT_rho0_00_rho3b0_05.pkl',
                'HQT/mcmc_gamma_muons_HQT_rho0_05_rho3b0_05.pkl',
                ]
    datafiles_mu = ['constantFF/mcmc_gamma-mu.hdf5',
                    '1pole/mcmc_samples-mu.hdf5',
                    '1pole_x_sqrtlambda/mcmc_samples-mu-sqrtlambda.hdf5',
                    # None,
                    # 'linear/mcmc_samples-rho-0.35-mu.hdf5',
                    # 'linear/mcmc_samples-rho-0.15-mu.hdf5',
                    # 'linear/mcmc_samples-rho-0.05-mu.hdf5',
                    # 'linear/mcmc_samples-rho0.00-mu.hdf5',
                    # 'linear/mcmc_samples-rho0.05-mu.hdf5',
                    # 'HQT/mcmc_samples-HQT-rho-0.35-rho3b0.00-mu.hdf5',
                    # 'HQT/mcmc_samples-HQT-rho-0.15-rho3b0.00-mu.hdf5',
                    'HQT/mcmc_samples-HQT-rho-0.05-rho3b0.00-mu.hdf5',
                    'HQT/mcmc_samples-HQT-rho0.0-rho3b0.00-mu.hdf5',
                    'HQT/mcmc_samples-HQT-rho0.05-rho3b0.00-mu.hdf5',
                    # 'HQT/mcmc_samples-HQT-rho-0.35-rho3b-0.05-mu.hdf5',
                    # 'HQT/mcmc_samples-HQT-rho-0.15-rho3b-0.05-mu.hdf5',
                    'HQT/mcmc_samples-HQT-rho-0.05-rho3b-0.05-mu.hdf5',
                    'HQT/mcmc_samples-HQT-rho0.0-rho3b-0.05-mu.hdf5',
                    'HQT/mcmc_samples-HQT-rho0.05-rho3b-0.05-mu.hdf5',
                    # 'HQT/mcmc_samples-HQT-rho-0.35-rho3b0.05-mu.hdf5',
                    # 'HQT/mcmc_samples-HQT-rho-0.15-rho3b0.05-mu.hdf5',
                    'HQT/mcmc_samples-HQT-rho-0.05-rho3b0.05-mu.hdf5',
                    'HQT/mcmc_samples-HQT-rho0.0-rho3b0.05-mu.hdf5',
                    'HQT/mcmc_samples-HQT-rho0.05-rho3b0.05-mu.hdf5',
                    ]
    out_mu = 'comparison_mu.pdf'
    files_tau = ['constantFF/mcmc_gamma_taus.pkl',
                 '1pole/mcmc_gamma_taus_1pole.pkl',
                 '1pole_x_sqrtlambda/mcmc_gamma_taus_1pole_x_sqrtlambda.pkl',
                 # marzia_ft,
                 # 'linear/mcmc_gamma_taus_linear_rho-0_35.pkl',
                 # 'linear/mcmc_gamma_taus_linear_rho-0_15.pkl',
                 # 'linear/mcmc_gamma_taus_linear_rho-0_05.pkl',
                 # 'linear/mcmc_gamma_taus_linear_rho0_00.pkl',
                 # 'linear/mcmc_gamma_taus_linear_rho0_05.pkl',
                 # 'HQT/mcmc_gamma_taus_HQT_rho-0_35_rho3b0_00.pkl',
                 # 'HQT/mcmc_gamma_taus_HQT_rho-0_15_rho3b0_00.pkl',
                 'HQT/mcmc_gamma_taus_HQT_rho-0_05_rho3b0_00.pkl',
                 'HQT/mcmc_gamma_taus_HQT_rho0_00_rho3b0_00.pkl',
                 'HQT/mcmc_gamma_taus_HQT_rho0_05_rho3b0_00.pkl',
                 # 'HQT/mcmc_gamma_taus_HQT_rho-0_35_rho3b-0_05.pkl',
                 # 'HQT/mcmc_gamma_taus_HQT_rho-0_15_rho3b-0_05.pkl',
                 'HQT/mcmc_gamma_taus_HQT_rho-0_05_rho3b-0_05.pkl',
                 'HQT/mcmc_gamma_taus_HQT_rho0_00_rho3b-0_05.pkl',
                 'HQT/mcmc_gamma_taus_HQT_rho0_05_rho3b-0_05.pkl',
                 # 'HQT/mcmc_gamma_taus_HQT_rho-0_35_rho3b0_05.pkl',
                 # 'HQT/mcmc_gamma_taus_HQT_rho-0_15_rho3b0_05.pkl',
                 'HQT/mcmc_gamma_taus_HQT_rho-0_05_rho3b0_05.pkl',
                 'HQT/mcmc_gamma_taus_HQT_rho0_00_rho3b0_05.pkl',
                 'HQT/mcmc_gamma_taus_HQT_rho0_05_rho3b0_05.pkl',
                 ]
    datafiles_tau = ['constantFF/mcmc_gamma-tau.hdf5',
                     '1pole/mcmc_samples-tau.hdf5',
                     '1pole_x_sqrtlambda/mcmc_samples-tau-sqrtlambda.hdf5',
                     # None,
                     # 'linear/mcmc_samples-rho-0.35-tau.hdf5',
                     # 'linear/mcmc_samples-rho-0.15-tau.hdf5',
                     # 'linear/mcmc_samples-rho-0.05-tau.hdf5',
                     # 'linear/mcmc_samples-rho0.00-tau.hdf5',
                     # 'linear/mcmc_samples-rho0.05-tau.hdf5',
                     # 'HQT/mcmc_samples-HQT-rho-0.35-rho3b0.00-tau.hdf5',
                     # 'HQT/mcmc_samples-HQT-rho-0.15-rho3b0.00-tau.hdf5',
                     'HQT/mcmc_samples-HQT-rho-0.05-rho3b0.00-tau.hdf5',
                     'HQT/mcmc_samples-HQT-rho0.0-rho3b0.00-tau.hdf5',
                     'HQT/mcmc_samples-HQT-rho0.05-rho3b0.00-tau.hdf5',
                     # 'HQT/mcmc_samples-HQT-rho-0.35-rho3b-0.05-tau.hdf5',
                     # 'HQT/mcmc_samples-HQT-rho-0.15-rho3b-0.05-tau.hdf5',
                     'HQT/mcmc_samples-HQT-rho-0.05-rho3b-0.05-tau.hdf5',
                     'HQT/mcmc_samples-HQT-rho0.0-rho3b-0.05-tau.hdf5',
                     'HQT/mcmc_samples-HQT-rho0.05-rho3b-0.05-tau.hdf5',
                     # 'HQT/mcmc_samples-HQT-rho-0.35-rho3b0.05-tau.hdf5',
                     # 'HQT/mcmc_samples-HQT-rho-0.15-rho3b0.05-tau.hdf5',
                     'HQT/mcmc_samples-HQT-rho-0.05-rho3b0.05-tau.hdf5',
                     'HQT/mcmc_samples-HQT-rho0.0-rho3b0.05-tau.hdf5',
                     'HQT/mcmc_samples-HQT-rho0.05-rho3b0.05-tau.hdf5',
                     ]
    out_tau = 'comparison_tau.pdf'
    labels = ['Constant FF',
              'One pole',
              r'One pole $\times\; \sqrt{\lambda}$',
              # 'MB: constant FF',
              # r'linear, $\rho=-0.35$',
              # r'linear, $\rho=-0.15$',
              # r'linear, $\rho=-0.05$',
              # r'linear, $\rho=0.0$',
              # r'linear, $\rho=0.05$',
              # r'HQT, $\rho=-0.35$, $\rho_{3b}=0.00$',
              # r'HQT, $\rho=-0.15$, $\rho_{3b}=0.00$',
              r'HQT, $\rho=-0.05$, $\rho_{3b}=0.00$',
              r'HQT, $\rho=0.0$, $\rho_{3b}=0.00$',
              r'HQT, $\rho=0.05$, $\rho_{3b}=0.00$',
              # r'HQT, $\rho=-0.35$, $\rho_{3b}=-0.05$',
              # r'HQT, $\rho=-0.15$, $\rho_{3b}=-0.05$',
              r'HQT, $\rho=-0.05$, $\rho_{3b}=-0.05$',
              r'HQT, $\rho=0.0$, $\rho_{3b}=-0.05$',
              r'HQT, $\rho=0.05$, $\rho_{3b}=-0.05$',
              # r'HQT, $\rho=-0.35$, $\rho_{3b}=0.05$',
              # r'HQT, $\rho=-0.15$, $\rho_{3b}=0.05$',
              r'HQT, $\rho=-0.05$, $\rho_{3b}=0.05$',
              r'HQT, $\rho=0.0$, $\rho_{3b}=0.05$',
              r'HQT, $\rho=0.05$, $\rho_{3b}=0.05$',
              ]
    superimpose(files_mu, datafiles_mu, labels, out_mu, 'mu')
    superimpose(files_tau, datafiles_tau, labels, out_tau, 'tau')