import FWCore.ParameterSet.Config as cms

# handle normal mixing or premixing
def getHcalDigitizer(process):
    if hasattr(process,'mixData'):
        return process.mixData
    if hasattr(process,'mix') and hasattr(process.mix,'digitizers') and hasattr(process.mix.digitizers,'hcal'):
        return process.mix.digitizers.hcal
    return None

def getHGCalDigitizer(process,section):
    if hasattr(process,'mix') and hasattr(process.mix,'digitizers'):
        if section == 'EE' and hasattr(process.mix.digitizers,'hgceeDigitizer'):
            return process.mix.digitizers.hgceeDigitizer
        elif section == 'FH' and hasattr(process.mix.digitizers,'hgchefrontDigitizer'):
            return process.mix.digitizers.hgchefrontDigitizer
        elif section == 'BH' and hasattr(process.mix.digitizers,'hgchebackDigitizer'):
            return process.mix.digitizers.hgchebackDigitizer
        elif section == 'HFNose' and hasattr(process.mix.digitizers,'hfnoseDigitizer'):
            return process.mix.digitizers.hfnoseDigitizer
    return None

# change assumptions about lumi rate
def setScenarioHLLHC(module,scenarioHLLHC):
    if scenarioHLLHC=="nominal":
        from CalibCalorimetry.HcalPlugins.HBHEDarkening_cff import _years_LHC, _years_HLLHC_nominal
        module.years = _years_LHC + _years_HLLHC_nominal
    elif scenarioHLLHC=="ultimate":
        from CalibCalorimetry.HcalPlugins.HBHEDarkening_cff import _years_LHC, _years_HLLHC_ultimate
        module.years = _years_LHC + _years_HLLHC_ultimate
    return module

# turnon = True enables default, False disables
# recalibration and darkening always together
def ageHB(process,turnon,scenarioHLLHC):
    if turnon:
        from CalibCalorimetry.HcalPlugins.HBHEDarkening_cff import HBDarkeningEP
        process.HBDarkeningEP = HBDarkeningEP
        process.HBDarkeningEP = setScenarioHLLHC(process.HBDarkeningEP,scenarioHLLHC)
    hcaldigi = getHcalDigitizer(process)
    if hcaldigi is not None: hcaldigi.HBDarkening = cms.bool(turnon)
    if hasattr(process,'es_hardcode'):
        process.es_hardcode.HBRecalibration = cms.bool(turnon)
    return process

def ageHE(process,turnon,scenarioHLLHC):
    if turnon:
        from CalibCalorimetry.HcalPlugins.HBHEDarkening_cff import HEDarkeningEP
        process.HEDarkeningEP = HEDarkeningEP
        process.HEDarkeningEP = setScenarioHLLHC(process.HEDarkeningEP,scenarioHLLHC)
    hcaldigi = getHcalDigitizer(process)
    if hcaldigi is not None: hcaldigi.HEDarkening = cms.bool(turnon)
    if hasattr(process,'es_hardcode'):
        process.es_hardcode.HERecalibration = cms.bool(turnon)
    return process

def ageHF(process,turnon):
    hcaldigi = getHcalDigitizer(process)
    if hcaldigi is not None: hcaldigi.HFDarkening = cms.bool(turnon)
    if hasattr(process,'es_hardcode'):
        process.es_hardcode.HFRecalibration = cms.bool(turnon)
    return process

def agedHFNose(process,algo=0):
    from SimCalorimetry.HGCalSimProducers.hgcalDigitizer_cfi import HFNose_setEndOfLifeNoise
    process = HFNose_setEndOfLifeNoise(process,byDose=True,byDoseAlgo=algo)
    return process

def agedHGCal(process,algo=0):
    from SimCalorimetry.HGCalSimProducers.hgcalDigitizer_cfi import HGCal_setEndOfLifeNoise
    process = HGCal_setEndOfLifeNoise(process,byDose=True,byDoseAlgo=algo)
    return process

def realisticHGCalStartup(process):
    from SimCalorimetry.HGCalSimProducers.hgcalDigitizer_cfi import HGCal_setRealisticStartupNoise
    process = HGCal_setRealisticStartupNoise(process)
    return process

# needs lumi to set proper ZS thresholds (tbd)
def ageSiPM(process,turnon,lumi):
    process.es_hardcode.hbUpgrade.doRadiationDamage = turnon
    process.es_hardcode.heUpgrade.doRadiationDamage = turnon

    # todo: determine ZS threshold adjustments

    # adjust PF thresholds for increased noise
    # based on: https://baylor.box.com/s/w32ja75krcbxcycyifexu28dwlgrj7wg
    hcal_lumis = [300, 1000, 3000, 4500, 1e10]
    hcal_thresholds = {
        300: {
            "seed": [0.5, 0.625, 0.75, 0.75],
            "rec": [0.4, 0.5, 0.6, 0.6],
        },
        1000: {
            "seed": [1.0, 1.5, 1.5, 1.5],
            "rec": [0.8, 1.2, 1.2, 1.2],
        },
        3000: {
            "seed": [1.25, 2.5, 2.5, 2.5],
            "rec": [1.0, 2.0, 2.0, 2.0],
        },
        4500: {
            "seed": [1.5, 3.0, 3.0, 3.0],
            "rec": [1.25, 2.5, 2.5, 2.5],
        },
    }
    ctmodules = ['calotowermaker','caloTowerForTrk','caloTowerForTrkPreSplitting','towerMaker','towerMakerWithHO']
    for ilumi, hcal_lumi in enumerate(hcal_lumis[:-1]):
        if lumi >= hcal_lumi and lumi < hcal_lumis[ilumi+1]:
            if hasattr(process,'particleFlowClusterHBHE'):
                process.particleFlowClusterHBHE.seedFinder.thresholdsByDetector[0].seedingThreshold              = hcal_thresholds[hcal_lumi]["seed"]
                process.particleFlowClusterHBHE.initialClusteringStep.thresholdsByDetector[0].gatheringThreshold = hcal_thresholds[hcal_lumi]["rec"]
                process.particleFlowClusterHBHE.pfClusterBuilder.recHitEnergyNorms[0].recHitEnergyNorm           = hcal_thresholds[hcal_lumi]["rec"]
                process.particleFlowClusterHBHE.pfClusterBuilder.positionCalc.logWeightDenominatorByDetector[0].logWeightDenominator = hcal_thresholds[hcal_lumi]["rec"]
                process.particleFlowClusterHBHE.pfClusterBuilder.allCellsPositionCalc.logWeightDenominatorByDetector[0].logWeightDenominator = hcal_thresholds[hcal_lumi]["rec"]
            if hasattr(process,'particleFlowClusterHCAL'):
                process.particleFlowClusterHCAL.pfClusterBuilder.allCellsPositionCalc.logWeightDenominatorByDetector[0].logWeightDenominator = hcal_thresholds[hcal_lumi]["rec"]
            if hasattr(process,'particleFlowRecHitHBHE'):
                process.particleFlowRecHitHBHE.producers[0].qualityTests[0].cuts[0].threshold = hcal_thresholds[hcal_lumi]["rec"]
            for ctmod in ctmodules:
                if hasattr(process,ctmod):
                    getattr(process,ctmod).HBThreshold1 = hcal_thresholds[hcal_lumi]["rec"][0]
                    getattr(process,ctmod).HBThreshold2 = hcal_thresholds[hcal_lumi]["rec"][1]
                    getattr(process,ctmod).HBThreshold = hcal_thresholds[hcal_lumi]["rec"][-1]
            break

    return process

def ageHcal(process,lumi,instLumi,scenarioHLLHC):
    hcaldigi = getHcalDigitizer(process)
    if hcaldigi is not None: hcaldigi.DelivLuminosity = cms.double(float(lumi))  # integrated lumi in fb-1

    # these lines need to be further activated by turning on 'complete' aging for HF
    if hasattr(process,'g4SimHits'):
        process.g4SimHits.HCalSD.InstLuminosity = cms.double(float(instLumi))
        process.g4SimHits.HCalSD.DelivLuminosity = cms.double(float(lumi))

    # recalibration and darkening always together
    if hasattr(process,'es_hardcode'):
        process.es_hardcode.iLumi = cms.double(float(lumi))

    # functions to enable individual subdet aging
    process = ageHB(process,True,scenarioHLLHC)
    process = ageHE(process,True,scenarioHLLHC)
    process = ageHF(process,True)
    process = ageSiPM(process,True,lumi)

    return process

def turn_on_HB_aging(process):
    process = ageHB(process,True,"")
    return process

def turn_off_HB_aging(process):
    process = ageHB(process,False,"")
    return process

def turn_on_HE_aging(process):
    process = ageHE(process,True,"")
    return process

def turn_off_HE_aging(process):
    process = ageHE(process,False,"")
    return process

def turn_on_HF_aging(process):
    process = ageHF(process,True)
    return process

def turn_off_HF_aging(process):
    process = ageHF(process,False)
    return process

def turn_off_SiPM_aging(process):
    process = ageSiPM(process,False,0.0)
    return process

def hf_complete_aging(process):
    if hasattr(process,'g4SimHits'):
        process.g4SimHits.HCalSD.HFDarkening = cms.untracked.bool(True)
    hcaldigi = getHcalDigitizer(process)
    if hcaldigi is not None: hcaldigi.HFDarkening = cms.untracked.bool(False)
    return process

def ageEcal(process,lumi,instLumi):
    if hasattr(process,'g4SimHits'):
        #these lines need to be further activiated by tuning on 'complete' aging for ecal
        process.g4SimHits.ECalSD.InstLuminosity = cms.double(instLumi)
        process.g4SimHits.ECalSD.DelivLuminosity = cms.double(float(lumi))

   # available conditions
    ecal_lumis = [300,1000,3000,4500]
    ecal_conditions = [
        ['EcalIntercalibConstantsRcd','EcalIntercalibConstants_TL{:d}_upgrade_8deg_v2_mc'],
        ['EcalIntercalibConstantsMCRcd','EcalIntercalibConstantsMC_TL{:d}_upgrade_8deg_v2_mc'],
        ['EcalLaserAPDPNRatiosRcd','EcalLaserAPDPNRatios_TL{:d}_upgrade_8deg_mc'],
        ['EcalPedestalsRcd','EcalPedestals_TL{:d}_upgradeTIA_8deg_mc'],
        ['EcalTPGLinearizationConstRcd','EcalTPGLinearizationConst_TL{:d}_upgrade_8deg_mc'],
    ]

    # update PF thresholds, based on https://indico.cern.ch/event/653123/contributions/2659235/attachments/1491385/2318364/170711_upsg_ledovskoy.pdf
    ecal_thresholds = {
        300 : 0.103,
        1000 : 0.175,
        3000 : 0.435,
        4500 : 0.707,
    }
    ecal_seed_multiplier = 2.5

    # try to get conditions
    if int(lumi) in ecal_lumis:
        if not hasattr(process.GlobalTag,'toGet'):
            process.GlobalTag.toGet=cms.VPSet()
        for ecal_condition in ecal_conditions:
            process.GlobalTag.toGet.append(cms.PSet(
                record = cms.string(ecal_condition[0]),
                tag = cms.string(ecal_condition[1].format(int(lumi))),
                connect = cms.string("frontier://FrontierProd/CMS_CONDITIONS")
                )
            )
        if hasattr(process,"particleFlowClusterECALUncorrected"):
            _seeds = process.particleFlowClusterECALUncorrected.seedFinder.thresholdsByDetector
            for iseed in range(0,len(_seeds)):
                if _seeds[iseed].detector.value()=="ECAL_BARREL":
                    _seeds[iseed].seedingThreshold = cms.double(ecal_thresholds[int(lumi)]*ecal_seed_multiplier)
            _clusters = process.particleFlowClusterECALUncorrected.initialClusteringStep.thresholdsByDetector
            for icluster in range(0,len(_clusters)):
                if _clusters[icluster].detector.value()=="ECAL_BARREL":
                    _clusters[icluster].gatheringThreshold = cms.double(ecal_thresholds[int(lumi)])

    return process

def ecal_complete_aging(process):
    if hasattr(process,'g4SimHits'):
        process.g4SimHits.ECalSD.AgeingWithSlopeLY = cms.untracked.bool(True)
    if hasattr(process,'ecal_digi_parameters'):
        process.ecal_digi_parameters.UseLCcorrection = cms.untracked.bool(False)
    return process

def customise_aging_300(process):
    process=ageHcal(process,300,5.0e34,"nominal")
    process=ageEcal(process,300,5.0e34)
    return process

def customise_aging_1000(process):
    process=ageHcal(process,1000,5.0e34,"nominal")
    process=turn_off_HE_aging(process) #avoid conflict between HGCal and Hcal in phase2 geom configuration
    process=ageEcal(process,1000,5.0e34)
    return process

def customise_aging_3000(process):
    process=ageHcal(process,3000,5.0e34,"nominal")
    process=turn_off_HE_aging(process) #avoid conflict between HGCal and Hcal in phase2 geom configuration
    process=ageEcal(process,3000,5.0e34)
    process=agedHGCal(process)
    process=agedHFNose(process)
    return process

def customise_aging_3000_ultimate(process):
    process=ageHcal(process,3000,7.5e34,"ultimate")
    process=turn_off_HE_aging(process) #avoid conflict between HGCal and Hcal in phase2 geom configuration
    process=ageEcal(process,3000,7.5e34)
    process=agedHGCal(process)
    process=agedHFNose(process)
    return process

def customise_aging_4500_ultimate(process):
    process=ageHcal(process,4500,7.5e34,"ultimate")
    process=turn_off_HE_aging(process) #avoid conflict between HGCal and Hcal in phase2 geom configuration
    process=ageEcal(process,4500,7.5e34)
    process=agedHGCal(process)
    process=agedHFNose(process)
    return process

[((66, 14, 66, 75), 'SimCalorimetry.HGCalSimProducers.hgcalDigitizer_cfi.HFNose_setEndOfLifeNoise', 'HFNose_setEndOfLifeNoise', (), '', False, 'from SimCalorimetry.HGCalSimProducers.hgcalDigitizer_cfi import HFNose_setEndOfLifeNoise\n'), ((71, 14, 71, 74), 'SimCalorimetry.HGCalSimProducers.hgcalDigitizer_cfi.HGCal_setEndOfLifeNoise', 'HGCal_setEndOfLifeNoise', (), '', False, 'from SimCalorimetry.HGCalSimProducers.hgcalDigitizer_cfi import HGCal_setEndOfLifeNoise\n'), ((76, 14, 76, 53), 'SimCalorimetry.HGCalSimProducers.hgcalDigitizer_cfi.HGCal_setRealisticStartupNoise', 'HGCal_setRealisticStartupNoise', ({(76, 45, 76, 52): 'process'}, {}), '(process)', False, 'from SimCalorimetry.HGCalSimProducers.hgcalDigitizer_cfi import HGCal_setRealisticStartupNoise\n'), ((41, 52, 41, 68), 'FWCore.ParameterSet.Config.bool', 'cms.bool', ({(41, 61, 41, 67): 'turnon'}, {}), '(turnon)', True, 'import FWCore.ParameterSet.Config as cms\n'), ((43, 46, 43, 62), 'FWCore.ParameterSet.Config.bool', 'cms.bool', ({(43, 55, 43, 61): 'turnon'}, {}), '(turnon)', True, 'import FWCore.ParameterSet.Config as cms\n'), ((52, 52, 52, 68), 'FWCore.ParameterSet.Config.bool', 'cms.bool', ({(52, 61, 52, 67): 'turnon'}, {}), '(turnon)', True, 'import FWCore.ParameterSet.Config as cms\n'), ((54, 46, 54, 62), 'FWCore.ParameterSet.Config.bool', 'cms.bool', ({(54, 55, 54, 61): 'turnon'}, {}), '(turnon)', True, 'import FWCore.ParameterSet.Config as cms\n'), ((59, 52, 59, 68), 'FWCore.ParameterSet.Config.bool', 'cms.bool', ({(59, 61, 59, 67): 'turnon'}, {}), '(turnon)', True, 'import FWCore.ParameterSet.Config as cms\n'), ((61, 46, 61, 62), 'FWCore.ParameterSet.Config.bool', 'cms.bool', ({(61, 55, 61, 61): 'turnon'}, {}), '(turnon)', True, 'import FWCore.ParameterSet.Config as cms\n'), ((180, 47, 180, 71), 'FWCore.ParameterSet.Config.untracked.bool', 'cms.untracked.bool', ({(180, 66, 180, 70): 'True'}, {}), '(True)', True, 'import FWCore.ParameterSet.Config as cms\n'), ((182, 52, 182, 77), 'FWCore.ParameterSet.Config.untracked.bool', 'cms.untracked.bool', ({(182, 71, 182, 76): 'False'}, {}), '(False)', True, 'import FWCore.ParameterSet.Config as cms\n'), ((188, 50, 188, 70), 'FWCore.ParameterSet.Config.double', 'cms.double', ({(188, 61, 188, 69): 'instLumi'}, {}), '(instLumi)', True, 'import FWCore.ParameterSet.Config as cms\n'), ((235, 53, 235, 77), 'FWCore.ParameterSet.Config.untracked.bool', 'cms.untracked.bool', ({(235, 72, 235, 76): 'True'}, {}), '(True)', True, 'import FWCore.ParameterSet.Config as cms\n'), ((237, 55, 237, 80), 'FWCore.ParameterSet.Config.untracked.bool', 'cms.untracked.bool', ({(237, 74, 237, 79): 'False'}, {}), '(False)', True, 'import FWCore.ParameterSet.Config as cms\n'), ((213, 36, 213, 47), 'FWCore.ParameterSet.Config.VPSet', 'cms.VPSet', ({}, {}), '()', True, 'import FWCore.ParameterSet.Config as cms\n'), ((216, 25, 216, 54), 'FWCore.ParameterSet.Config.string', 'cms.string', ({(216, 36, 216, 53): 'ecal_condition[0]'}, {}), '(ecal_condition[0])', True, 'import FWCore.ParameterSet.Config as cms\n'), ((218, 26, 218, 78), 'FWCore.ParameterSet.Config.string', 'cms.string', ({(218, 37, 218, 77): '"""frontier://FrontierProd/CMS_CONDITIONS"""'}, {}), "('frontier://FrontierProd/CMS_CONDITIONS')", True, 'import FWCore.ParameterSet.Config as cms\n')]