from __future__ import print_function
from multiprocessing import Pool
import os, glob, shutil
import numpy as np
import subprocess

# Strings to replace in crabConfig template
SAMPLE = '__SAMPLE__'
DATASET = '__DATASET__'
#EVTLIST = '__EVTLIST__'
#LUMIMASK = '__LUMIMASK__'
CAMPAIGN = '__CAMPAIGN__'
UNITSPERJOB = '__UNITSPERJOB__'

this_campaign = 'Era2017_17Dec2019_ggntuplev1'

crab_folder = 'crab_%s'%this_campaign
if not os.path.isdir(crab_folder):
    os.makedirs(crab_folder)

'''
run = 'DoubleEG_2017'
#job_units = 100
job_units = 200
job_units = 1000
#samples = {
#    '%sB'%run: '/DoubleEG/mandrews-Run2017B_17Nov2017-v1_AOD_slim-ext_v2-3bfee02a0afb4bfd03fd5261a90623cd/USER',
#    '%sC'%run: '/DoubleEG/mandrews-Run2017C_17Nov2017-v1_AOD_slim-ext_v2-3bfee02a0afb4bfd03fd5261a90623cd/USER',
#    '%sD'%run: '/DoubleEG/mandrews-Run2017D_17Nov2017-v1_AOD_slim-ext-3bfee02a0afb4bfd03fd5261a90623cd/USER',
#    '%sE'%run: '/DoubleEG/mandrews-Run2017E_17Nov2017-v1_AOD_slim-ext_v2-3bfee02a0afb4bfd03fd5261a90623cd/USER',
#    '%sF'%run: '/DoubleEG/mandrews-Run2017F_17Nov2017-v1_AOD_slim-ext_v2-964eedbb4080135606054ba835f474dc/USER',
#    }
# key: [<primary dset>, <secondary dset>]
samples = {
    '%sB'%run: '/DoubleEG/Run2017B-31Mar2018-v1/MINIAOD',
    '%sC'%run: '/DoubleEG/Run2017C-31Mar2018-v1/MINIAOD',
    '%sD'%run: '/DoubleEG/Run2017D-31Mar2018-v1/MINIAOD',
    '%sE'%run: '/DoubleEG/Run2017E-31Mar2018-v1/MINIAOD',
    '%sF'%run: '/DoubleEG/Run2017F-31Mar2018-v1/MINIAOD'
    }
'''
'''
run = 'Run2017'
job_units = 1000
samples = {
    '%sB-MINIAOD'%run: '/DoubleEG/Run2017B-31Mar2018-v1/MINIAOD',
    '%sC-MINIAOD'%run: '/DoubleEG/Run2017C-31Mar2018-v1/MINIAOD',
    '%sD-MINIAOD'%run: '/DoubleEG/Run2017D-31Mar2018-v1/MINIAOD',
    '%sE-MINIAOD'%run: '/DoubleEG/Run2017E-31Mar2018-v1/MINIAOD',
    '%sF-MINIAOD'%run: '/DoubleEG/Run2017F-31Mar2018-v1/MINIAOD'
    }
'''
#'''
job_units = 50
run = 'h24gamma_1j_1M'
samples = {
    '%s_100MeV'%run: '/h24gamma_01Nov2019-rhECAL/mandrews-h24gamma_1j_1M_100MeV_PU2017_MINIAODSIM_v2-919c80a76a70185609d372d13ecbc645/USER',
    '%s_400MeV'%run: '/h24gamma_01Nov2019-rhECAL/mandrews-h24gamma_1j_1M_400MeV_PU2017_MINIAODSIM_v2-919c80a76a70185609d372d13ecbc645/USER',
    '%s_1GeV'%run:   '/h24gamma_01Nov2019-rhECAL/mandrews-h24gamma_1j_1M_1GeV_PU2017_MINIAODSIM_v2-919c80a76a70185609d372d13ecbc645/USER',
    }
#'''
'''
#job_units = 100
job_units = 150
samples = {
    'DiPhotonJets':    '/DiPhotonJets_MGG-80toInf_13TeV_amcatnloFXFX_pythia8/RunIIFall17MiniAODv2-PU2017_12Apr2018_94X_mc2017_realistic_v14-v1/MINIAODSIM',
    'GJet_Pt-20to40':  '/GJet_Pt-20to40_DoubleEMEnriched_MGG-80toInf_TuneCP5_13TeV_Pythia8/RunIIFall17MiniAODv2-PU2017_12Apr2018_94X_mc2017_realistic_v14-v1/MINIAODSIM',
    'GJet_Pt-40toInf': '/GJet_Pt-40toInf_DoubleEMEnriched_MGG-80toInf_TuneCP5_13TeV_Pythia8/RunIIFall17MiniAODv2-PU2017_12Apr2018_94X_mc2017_realistic_v14-v2/MINIAODSIM',
    'QCD_Pt-30to40':   '/QCD_Pt-30to40_DoubleEMEnriched_MGG-80toInf_TuneCP5_13TeV_Pythia8/RunIIFall17MiniAODv2-PU2017_12Apr2018_94X_mc2017_realistic_v14-v1/MINIAODSIM',
    'QCD_Pt-40toInf':  '/QCD_Pt-40toInf_DoubleEMEnriched_MGG-80toInf_TuneCP5_13TeV_Pythia8/RunIIFall17MiniAODv2-PU2017_12Apr2018_94X_mc2017_realistic_v14-v1/MINIAODSIM',
    'GluGluHToGG':     '/GluGluHToGG_M125_13TeV_amcatnloFXFX_pythia8/RunIIFall17MiniAODv2-PU2017_12Apr2018_94X_mc2017_realistic_v14-v1/MINIAODSIM',
    }
'''

for s,dset in samples.iteritems(): #python3: samples.items()

    print('For sample:',s)

    assert os.environ['CMSSW_BASE'] != ''
    base_dir = '%s/src'%os.environ['CMSSW_BASE']

    # Read in crabConfig template
    with open('crabConfig_ggNtuplizer.py', "r") as template_file:
        file_data = template_file.read()

    # Replace crabConfig template strings with sample-specific values
    # NOTE: event list generated by edmPickEvents (pickevents_runEvents.txt) does not include LS numbers!!
    # This gives non-unique eventIds and so MUST be replaced with the original input `evt_list` above!
    file_data = file_data.replace(SAMPLE, s)
    file_data = file_data.replace(DATASET, dset)
    #file_data = file_data.replace(EVTLIST, evt_list)
    #file_data = file_data.replace(LUMIMASK, lumi_list)
    file_data = file_data.replace(CAMPAIGN, this_campaign)
    file_data = file_data.replace(UNITSPERJOB, str(job_units))

    # Write out sample-specific crabConfig
    with open('%s/crabConfig_ggNtuplizer_%s.py'%(crab_folder, s), "w") as sample_file:
        sample_file.write(file_data)
