import gzip
import pickle
#from coffea import hist
import coffea
import hist
import argparse
import numpy as np
from topcoffea.modules.histEFT import HistEFT
import mplhep as hep
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
from topeft.modules.axes import info as axes_info

parser = argparse.ArgumentParser(description="specify pkl file")
parser.add_argument("pkl_file_path", help="Path to the pkl file")
parser.add_argument("--hist",default='photon_pt2', help="Hist of your choice")
parser.add_argument("-y","--year",default='17',help='year')
parser.add_argument("-r", "--region",default='SR_ch', help='SR_ch or Zg_ch')
    
# Parse the command-line arguments
args = parser.parse_args()

pkl_file = args.pkl_file_path
hist = args.hist
year = args.year
region = args.region

hists={}
with gzip.open(pkl_file) as f:
    hin = pickle.load(f)

hist_pt = hin[hist]
hist_pt_sumw2 = hin[hist+'_sumw2']
print(hist_pt.axes['process'])
print(hist_pt.axes['channel'])
appl_of_interest = 'isSR_2lOS_ph'

Zg_ch = []
SR_of_ch = []
SR_sf_ch = []
for ch in hist_pt.axes['channel']:
    if "Zg" in ch:
        Zg_ch.append(ch)
    else:
        if "of" in ch:
            SR_of_ch.append(ch)
        elif "sf" in ch:
            SR_sf_ch.append(ch)
SR_ch = SR_of_ch+SR_sf_ch
if region=='SR':
    ch_of_interest = SR_ch
else:
    ch_of_interest = Zg_ch

central_sample = []
private_sample = []

for pro in hist_pt.axes['process']:
    if "central" in pro:
        central_sample.append(pro)
    else:
        private_sample.append(pro)

fig, (ax, rax) = plt.subplots(
        nrows=2,
        ncols=1,
        figsize=(10,10),
        gridspec_kw={"height_ratios": (3, 1)},
        sharex=True
)
fig.subplots_adjust(hspace=.07)

hep.style.use("CMS")
plt.sca(ax)

bins = hist_pt.axes[hist].edges
bins = np.append(bins, [bins[-1] + (bins[-1] - bins[-2])*0.3]) 
print("bins")
print(bins)
#bins = np.insert(bins,0,(bins[0]-bins[1])*0.3)
#bins = np.insert(bins,0,0)

#print(bins)
#SR_ch = ['2los_of_ph_1j']
if year == 'all':
    central_array = hist_pt.integrate('channel',ch_of_interest).integrate('appl',appl_of_interest).integrate('systematic','nominal').integrate('process',central_sample)[{'channel':sum,'process':sum}].eval({})[()][1:]
    central_sumw2_array = hist_pt_sumw2.integrate('channel',ch_of_interest).integrate('appl',appl_of_interest).integrate('systematic','nominal').integrate('process',central_sample)[{'channel':sum,'process':sum}].eval({})[()][1:]
    private_array = hist_pt.integrate('channel',ch_of_interest).integrate('appl',appl_of_interest).integrate('systematic','nominal').integrate('process',private_sample)[{'channel':sum,'process':sum}].eval({})[()][1:]

else:
    central_array = hist_pt.integrate('channel',ch_of_interest).integrate('appl',appl_of_interest).integrate('systematic','nominal').integrate('process',f'TTGamma_centralUL{year}')[{'channel':sum}].eval({})[()][1:]
    central_sumw2_array = hist_pt_sumw2.integrate('channel',ch_of_interest).integrate('appl',appl_of_interest).integrate('systematic','nominal').integrate('process',f'TTGamma_centralUL{year}')[{'channel':sum}].eval({})[()][1:]
    private_array = hist_pt.integrate('channel',ch_of_interest).integrate('appl',appl_of_interest).integrate('systematic','nominal').integrate('process',f'TTGamma_Dilept_privateUL{year}')[{'channel':sum}].eval({})[()][1:]

#print(central_array)
#brent_CF = np.array([1.30940008,1.15095172, 1.00803265, 0.90287607, 0.90437876,  0.92583374, 0.98851039, 0.983581, 0.97709902])
#private_array = private_array*brent_CF
#SR_CF = np.array([1.063,1.052,1.082,1.053,1.043,0.993,1.022,0.921,0.862])
#private_array = private_array*SR_CF
#common_CF = np.array([1.05729045,1.05561951,1.07377112,1.06411441,1.05234515,1.00253823, 1.05035979, 0.94874264, 0.83789775]) 
common_CF = np.array([1.05674389, 1.06924074, 1.08474516, 0.99113743])
private_array = private_array*common_CF

UL17_private_array = np.array([494.32974044,231.39988324,147.44290014,100.81294904,77.71776398,14.05260284,13.91907981,8.74200135,11.88909238])
UL17_central_array = np.array([525.32403097,243.36341114,159.56328075,106.17224186,81.06339636,13.94845177,14.22977723,8.0556855,10.24620589])
#private_array = private_array + UL17_private_array
#central_array=central_array+UL17_central_array

print("\nPrinting private array")
print(private_array)
print("\nPrinting central array")
print(central_array)
print("\nPrinting ratio of central to private array")
print(central_array/private_array)
print("")

hep.histplot(
    central_array,
    ax=ax,
    bins=bins,
    stack=False,
    density=False,
    label=f'UL{year} central skimmed tt$\gamma$',
    histtype='step',
)

hep.histplot(
    private_array,
    ax=ax,
    bins=bins,
    stack=False,
    density=False,
    label=f'UL{year} private skimmed tt$\gamma$',
    histtype='step',
)


hep.histplot(
    (central_array/private_array),
    yerr = np.sqrt(central_sumw2_array)/private_array,
    #error_opts = DATA_ERR_OPS,
    ax=rax,
    bins=bins,
    stack=False,
    density=False,
    #flow='show',
    histtype='errorbar',
)

ax.set_xlabel(None)
#ax.set_yscale('log')
rax.set_ylabel('Central/Private', loc='center',fontsize=17)
rax.set_ylim(0.95,1.15)
#rax.set_ylim(0.001,0.02)

rax.tick_params(axis='both',which='major',direction='in',labelsize=15, length=9,top=True,right=True)
rax.set_xlabel(axes_info[hist]['label'],loc='right',fontsize=20)
ax.tick_params(axis='both',which='major',direction='in',labelsize=15,length=9,top=True,right=True)
ax.set_ylabel('Events',loc='center',fontsize=20)
#ax.set_yscale('log')

# Minor ticks: no labels, just ticks
ax.minorticks_on()
rax.minorticks_on()
ax.tick_params(axis='both', which='minor', direction='in', length=3,top=True,right=True)
rax.tick_params(axis='both', which='minor', direction='in', length=3,top=True,right=True)
rax.axhline(1.0,linestyle="-",color="k",linewidth=1)
ax.legend()
#fig.savefig('central_vs_private_UL17_eftCoeffsNone_SRSelection.pdf')
dest_path = f'/afs/crc.nd.edu/user/a/abasnet/www/EFT/Spring2025/EFTttgamma_sample_validation_study/UL{year}/'
#fig.savefig(dest_path+'central_vs_private_UL17_SREventSelection_allWeights_eftCoeffsNone_skimmedSamples.pdf')
#fig.savefig(dest_path+'central_vs_private_UL17_noEventSelection_noEventSelectionMask_onlylumiXsecNormWeight_skimmedSamples.pdf')
#fig.savefig(dest_path+f'{hist}_UL17-private-vs-central-Unskimmed_SREventSelection_fullWeights_PrivateXSec2p9922_CentralXSec2p2471_GenMatchedPhotons_BrentCF.pdf')
#fig.savefig(dest_path+f'{hist}_UL17-private-vs-central-Skimmed_SREventSelection_allWeights_PrivateXSec2p9922_SowIntact_CentralIntact_GenMatchedPhotonsOnly_photon-pt-is-GenMatched.pdf')
#fig.savefig(dest_path+f'{hist}_UL17-private-vs-central-Skimmed_SREventSelection_allWeights_PrivateXSec2p9922_SowIntact_CentralIntact_GenMatchedPhotonsOnly_photon-pt-is-PureGen.pdf')
#fig.savefig(dest_path+f'{hist}_UL17-private-vs-central-Skimmed_SREventSelection_fullWeights_PrivateXSec2p9922_SowIntact_CentralIntact_OriginalpTBinning.pdf')
fig.savefig(dest_path+f'{hist}_UL{year}-private-vs-central-Skimmed_{region}EventSelection_fullWeights_RatioErrorCorrected_scaledBycommonCF_PrivateXSec2p9922_SowIntact_CentralIntact_purerecoPhotonpT_FinalphotonpTBinsForFitting.pdf')
#fig.savefig(dest_path+f'{hist}_UL{year}-private-vs-central-Skimmed_{region}EventSelection_fullWeights_RatioErrorCorrected_PrivateXSec2p9922_SowIntact_CentralIntact_purerecoPhotonpT.pdf')
print("\nSaved the plot")