#!/usr/bin/env python
import os, sys
import argparse
import ROOT
ROOT.PyConfig.IgnoreCommandLineOptions = True
from importlib import import_module
from PhysicsTools.NanoAODTools.postprocessing.framework.postprocessor import PostProcessor
from PhysicsTools.NanoSUSYTools.modules.LLObjectsProducer import *
from PhysicsTools.NanoSUSYTools.modules.qcdSFProducer import *
from PhysicsTools.NanoSUSYTools.modules.updateEvtWeightFastsim import *

def main(args):
    isdata = len(args.dataEra) > 0
    process = args.process

    mods = [
	LLObjectsProducer(args.era, args.sampleName, isData=isdata),
    ]

    if not isdata:
	mods.append(LLObjectsProducer(args.era, args.sampleName, isData=isdata, applyUncert="JESUp"))
	mods.append(LLObjectsProducer(args.era, args.sampleName, isData=isdata, applyUncert="JESDown"))
	mods.append(LLObjectsProducer(args.era, args.sampleName, isData=isdata, applyUncert="METUnClustUp"))
	mods.append(LLObjectsProducer(args.era, args.sampleName, isData=isdata, applyUncert="METUnClustDown"))

    if "QCD" in args.sampleName:
	mods.append(qcdSFProducer(args.era))

    if "SMS" in args.sampleName:
	mods.append(UpdateEvtWeightFastsim(isdata, args.nEvents, args.sampleName))

    files = []
    if len(args.inputfile) > 5 and args.inputfile[0:5] == "file:":
        #This is just a single test input file
        files.append(args.inputfile[5:])
    else:
        #this is a file list
        with open(args.inputfile) as f:
            files = [line.strip() for line in f]

    if process=="limits": p=PostProcessor(args.outputfile,files,cut="Pass_MET && Pass_NJets30", branchsel=None, outputbranchsel="keep_and_drop_limits.txt", modules=mods,provenance=False,maxEvents=args.maxEvents)
    else: 		  p=PostProcessor(args.outputfile,files,cut="Pass_MET && Pass_NJets30", branchsel=None, outputbranchsel="keep_and_drop_LL.txt", modules=mods,provenance=False,maxEvents=args.maxEvents)
    p.run()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='NanoAOD postprocessing.')
    parser.add_argument('-i', '--inputfile',
        default = "testing.txt",
        help = 'Path to the input filelist. To run with a single file instead of a file list prepend the filepath with \"file:\" (Default: testing.txt)')
    parser.add_argument('-o', '--outputfile',
                        default="./",
                        help = 'Path to the output file location. (Default: .)')
    parser.add_argument('-e', '--era',
        default = "2017", help = 'Year of production')
    parser.add_argument('-f', '--isFastSim', action="store",  default = False,
                        help = "Input file is fastsim (Default: false)")
    parser.add_argument('-D', '--isData',  type=str, default = "",
                        help = "Data era (B, C, D, ...).  Using this flag also switches the procesor to data mode. (Default: None, i.e. MC )")
    parser.add_argument('-d', '--dataEra',    action="store",  type=str, default = "",
                        help = "Data era (B, C, D, ...).  Using this flag also switches the procesor to data mode. (Default: None, i.e. MC )")
    parser.add_argument('-s', '--sampleName',    action="store",  type=str, default = "",
                        help = "Name of MC sample (from sampleSet file) (Default: )")
    parser.add_argument('-c', '--crossSection',
                        type=float,
                        default = 1,
                        help = 'Cross Section of MC to use for MC x-sec*lumi weight (Default: 1.0)')
    parser.add_argument('-n', '--nEvents',
                        type=float,
                        default = 1,
                        help = 'Number of events to use for MC x-sec*lumi weight (NOT the number of events to run over) (Default: 1.0)')
    parser.add_argument('-m', '--maxEvents',
                        type=int,
                        default = -1,
                        help = 'MAximum number of events to process (Default: all events)')
    parser.add_argument('-p', '--process', type=str, default = "",
                        help = "Type of QCD process to do (jetres or smear)")
    args = parser.parse_args()
    main(args)
