import json
import os
import ROOT
from Gaudi.Configuration import *
from GaudiKernel import SystemOfUnits as units

from Configurables import ApplicationMgr, EICDataSvc, PodioOutput, GeoSvc

from Configurables import PodioInput
from Configurables import Jug__Digi__CalorimeterHitDigi as CalorimeterHitDigi
from Configurables import Jug__Reco__ImagingPixelReco as ImagingPixelReco
from Configurables import Jug__Reco__ImagingTopoCluster as ImagingTopoCluster
from Configurables import Jug__Reco__ImagingClusterReco as ImagingClusterReco


# input arguments from calibration file
with open('config/emcal_barrel_calibration.json') as f:
    calib_data = json.load(f)['electron']

kwargs = dict()
kwargs['sf'] = float(calib_data['sampling_fraction_img'])

# input arguments through environment variables
kwargs['input'] = os.environ.get('CB_EMCAL_SIM_FILE', '../topside/barrel_pion0_5GeV.edm4hep.root')
kwargs['output'] = os.environ.get('CB_EMCAL_REC_FILE', 'barrel_pion0_5GeV_cluster.root')
kwargs['compact'] = os.environ.get('CB_EMCAL_COMPACT_PATH', '../topside/test.xml')
kwargs['nev'] = int(os.environ.get('CB_EMCAL_NUMEV', 100))

if kwargs['nev'] < 1:
    f = ROOT.TFile(kwargs['input'])
    kwargs['nev'] = f.events.GetEntries()

print(kwargs)

geo_service = GeoSvc("GeoSvc", detectors=kwargs['compact'].split(','), OutputLevel=INFO)
podioevent = EICDataSvc("EventDataSvc", inputs=kwargs['input'].split(','), OutputLevel=DEBUG)
out = PodioOutput("out", filename=kwargs['output'])

podioinput = PodioInput("PodioReader", collections=["MCParticles", "EcalBarrelHits"], OutputLevel=DEBUG)

# use the same daq_setting for digi/reco pair
daq_setting = dict(
        dynamicRangeADC=3*units.MeV,
        capacityADC=32767,
        pedestalMean=400,
        pedestalSigma=50)   # 50/32767*3 MeV ~ 5 keV
imcaldigi = CalorimeterHitDigi("imcal_digi",
        # OutputLevel=DEBUG,
        inputHitCollection="EcalBarrelHits",
        outputHitCollection="DigiEcalBarrelImagingHits",
        energyResolutions=[0., 0.02, 0.],
        **daq_setting)
imcalreco = ImagingPixelReco("imcal_reco",
        # OutputLevel=DEBUG,
        inputHitCollection=imcaldigi.outputHitCollection,
        outputHitCollection="RecoEcalBarrelImagingHits",
        readoutClass="EcalBarrelHits",
        layerField="layer",
        sectorField="module",
        **daq_setting)

imcalcluster = ImagingTopoCluster("imcal_cluster",
        # OutputLevel=DEBUG,
        inputHitCollection=imcalreco.outputHitCollection,
        outputProtoClusterCollection="EcalBarrelImagingProtoClusters",
        localDistXY=[2.*units.mm, 2*units.mm],
        layerDistEtaPhi=[10*units.mrad, 10*units.mrad],
        neighbourLayersRange=2,
        sectorDist=3.*units.cm,
        minClusterNhits=5)
clusterreco = ImagingClusterReco("imcal_clreco",
        # OutputLevel=DEBUG,
        inputProtoClusterCollection=imcalcluster.outputProtoClusterCollection,
        outputLayerCollection="EcalBarrelImagingClustersLayers",
        outputClusterCollection="EcalBarrelImagingClusters",
        mcHits="EcalBarrelHits",
        samplingFraction=sf)

out.outputCommands = ["keep *"]

ApplicationMgr(
    TopAlg=[podioinput, imcaldigi, imcalreco, imcalcluster, clusterreco, out],
    EvtSel='NONE',
    EvtMax=kwargs['nev'],
    ExtSvc=[podioevent],
    OutputLevel=DEBUG
)