events = ur.open('{}/{}.root:tree'.format(settings['inPath'], settings['inFile']))


# cluster daten
print('loading cluster data...')
clsCharge = getData(events, 'PXDClusters/PXDClusters.m_clsCharge')
iterLength = 0
for i in range(len(clsCharge)):
    iterLength += len(clsCharge[i])
if settings['clusterCharge'] is True or settings['eventNumbers'] is True:
    clsChargeFlat = flatten(clsCharge)
    eventNumbers = genEventNumbers(clsCharge)
if settings['seedCharge'] is True:
    seedCharge = getData(events, 'PXDClusters/PXDClusters.m_seedCharge')
    seedChargeFlat = flatten(seedCharge)
if settings['clusterSize'] is True:
    clsSize = getData(events, 'PXDClusters/PXDClusters.m_clsSize')
    clsSizeFlat = flatten(clsSize)
if settings['uvSize'] is True:
    uSize = getData(events, 'PXDClusters/PXDClusters.m_uSize')
    vSize = getData(events, 'PXDClusters/PXDClusters.m_vSize')
    uSizeFlat = flatten(uSize)
    vSizeFlat = flatten(vSize)
if settings['uvPosition'] is True or settings['coordinates'] is True:
    uPosition = getData(events, 'PXDClusters/PXDClusters.m_uPosition')
    vPosition = getData(events, 'PXDClusters/PXDClusters.m_vPosition')
    uPositionFlat = flatten(uPosition)
    vPositionFlat = flatten(vPosition)
if settings['sensorID'] is True or settings['coordinates'] is True or settings['ladder'] or settings['layer']:
    sensorID = getData(events, 'PXDClusters/PXDClusters.m_sensorID')
    sensorIDFlat = flatten(sensorID)


# koordinaten, layer und ladder
if settings['coordinates'] is True:
    print('calculating coordinates...')
    xcoords, ycoords, zcoords = getCartesianFlattened(uPosition, vPosition, sensorID, transformation)
if settings['layer'] is True or settings['ladder'] is True:
    layers, ladders = getLayers(sensorIDFlat, layersLadders)


# digits laden
if settings['uvCellID'] is True or settings['cellCharge'] is True or settings['pxdClusters'] is True:
    print('loading cell data...')
    uCellID = getData(events, 'PXDDigits/PXDDigits.m_uCellID')
    vCellID = getData(events, 'PXDDigits/PXDDigits.m_vCellID')
    charge = getData(events, 'PXDDigits/PXDDigits.m_charge')
    digitClusters = getData(events, 'PXDClustersToPXDDigits/m_elements/m_elements.m_to')

    # digits umorganisieren
    uCellIDFlat = []
    vCellIDFlat = []
    chargeFlat = []
    for i in range(len(digitClusters)):
        udigits, vdigits, cdigits = getEventData(digitClusters[i], uCellID[i], vCellID[i], charge[i])
        uCellIDFlat.append(udigits)
        vCellIDFlat.append(vdigits)
        chargeFlat.append(cdigits)

    # digits flachen
    uCellIDFlat = flatten(uCellIDFlat)
    vCellIDFlat = flatten(vCellIDFlat)
    chargeFlat = flatten(chargeFlat)


# adc matrix bestimmen
if settings['pxdClusters'] is True:
    print('generating pixel cluster')
    clustersFlat = getClustersFlattened(uCellID, vCellID, charge, digitClusters)


# mc laden
if settings['pdg'] or settings['momenta'] or settings['clusterNumbers']:
    print('loading mc info...')
    pdg = getData(events, 'MCParticles/MCParticles.m_pdg')
    momentumX = getData(events, 'MCParticles/MCParticles.m_momentum_x')
    momentumY = getData(events, 'MCParticles/MCParticles.m_momentum_y')
    momentumZ = getData(events, 'MCParticles/MCParticles.m_momentum_z')
    mcClusters = getData(events, 'PXDClustersToMCParticles/m_elements/m_elements.m_to')
    mcfromClusters = getData(events, 'PXDClustersToMCParticles/m_elements/m_elements.m_from')

    # mc umorganisieren
    momentumXList = []
    momentumYList = []
    momentumZList = []
    pdgList = []
    clusterNumbersList = []
    for i in range(len(mcClusters)):
        fullClusterReferences = fillMCList(mcfromClusters[i], mcClusters[i], len(clsCharge[i]))
        clusterNumbersList.append(fullClusterReferences)
        pdgs, xmom, ymom, zmom = getMCData(fullClusterReferences, pdg[i], momentumX[i], momentumY[i], momentumZ[i])
        momentumXList.append(xmom)
        momentumYList.append(ymom)
        momentumZList.append(zmom)
        pdgList.append(pdgs)

    # mc flachen
    momentumXFlat = flatten(momentumXList)
    momentumYFlat = flatten(momentumYList)
    momentumZFlat = flatten(momentumZList)
    pdgFlat = flatten(pdgList)
    clusterNumbersFlat = flatten(clusterNumbersList)




if settings['saveFormat'] == 'structured':
    print('preparing data array...')
    # daten type definieren
    dataType = []
    if settings['eventNumbers'] is True:
        dataType.append(('event', 'i4'))
    if settings['clusterNumbers'] is True:
        dataType.append(('clsNumber', 'i4'))
    if settings['pxdClusters'] is True:
        dataType.append(('cluster', 'i4', (9,9)))
    if settings['clusterCharge'] is True:
        dataType.append(('clsCharge', 'i4'))
    if settings['seedCharge'] is True:
        dataType.append(('seedCharge', 'i4'))
    if settings['clusterSize'] is True:
        dataType.append(('clsSize', 'i4'))
    if settings['uvSize'] is True:
        dataType.append(('uSize', 'i4'))
        dataType.append(('vSize', 'i4'))
    if settings['uvCellID'] is True:
        dataType.append(('uCellID', 'object'))
        dataType.append(('vCellID', 'object'))
    if settings['cellCharge'] is True:
        dataType.append(('charge', 'object'))
    if settings['coordinates'] is True:
        dataType.append(('xPosition', 'f8'))
        dataType.append(('yPosition', 'f8'))
        dataType.append(('zPosition', 'f8'))
    if settings['uvPosition'] is True:
        dataType.append(('uPosition', 'f8'))
        dataType.append(('vPosition', 'f8'))
    if settings['sensorID'] is True:
        dataType.append(('sensorID', 'i8'))
    if settings['layer'] is True:
        dataType.append(('layer', 'i4'))
    if settings['ladder'] is True:
        dataType.append(('ladder', 'i4'))
    if settings['momenta'] is True:
        dataType.append(('xMomentum', 'f8'))
        dataType.append(('yMomentum', 'f8'))
        dataType.append(('zMomentum', 'f8'))
    if settings['pdg'] is True:
        dataType.append(('pdg', 'i4'))

    # structured array definieren
    forArray = []
    for i in range(iterLength):
        toAppend = []
        if settings['eventNumbers'] is True:
            toAppend.append(eventNumbers[i])
        if settings['clusterNumbers'] is True:
            toAppend.append(clusterNumbersFlat[i])
        if settings['pxdClusters'] is True:
            toAppend.append(clustersFlat[i])
        if settings['clusterCharge'] is True:
            toAppend.append(clsChargeFlat[i])
        if settings['seedCharge'] is True:
            toAppend.append(seedChargeFlat[i])
        if settings['clusterSize'] is True:
            toAppend.append(clsSizeFlat[i])
        if settings['uvSize'] is True:
            toAppend.append(uSizeFlat[i])
            toAppend.append(vSizeFlat[i])
        if settings['uvCellID'] is True:
            toAppend.append(uCellIDFlat[i])
            toAppend.append(vCellIDFlat[i])
        if settings['cellCharge'] is True:
            toAppend.append(chargeFlat[i])
        if settings['coordinates'] is True:
            toAppend.append(xcoords[i])
            toAppend.append(ycoords[i])
            toAppend.append(zcoords[i])
        if settings['uvPosition'] is True:
            toAppend.append(uPositionFlat[i])
            toAppend.append(vPositionFlat[i])
        if settings['sensorID'] is True:
            toAppend.append(sensorIDFlat[i])
        if settings['layer'] is True:
            toAppend.append(layers[i])
        if settings['ladder'] is True:
            toAppend.append(ladders[i])
        if settings['momenta'] is True:
            toAppend.append(momentumXFlat[i])
            toAppend.append(momentumYFlat[i])
            toAppend.append(momentumZFlat[i])
        if settings['pdg'] is True:
            toAppend.append(pdgFlat[i])
        toAppend = tuple(toAppend)
        forArray.append(toAppend)
    saveArr = np.array(forArray, dtype=dataType)

elif settings['saveFormat'] == 'array' or settings['saveFormat'] == 'txt':
    print('preparing data array...')
    # structured array definieren
    forArray = []
    for i in range(iterLength):
        toAppend = np.array([0])
        if settings['eventNumbers'] is True:
            toAppend = np.concatenate((toAppend, [eventNumbers[i]]))
        if settings['clusterNumbers'] is True:
            toAppend = np.concatenate((toAppend, [clusterNumbersFlat[i]]))
        if settings['pxdClusters'] is True:
            toAppend = np.concatenate((toAppend, clustersFlat[i].flatten()))
        if settings['clusterCharge'] is True:
            toAppend = np.concatenate((toAppend, [clsChargeFlat[i]]))
        if settings['seedCharge'] is True:
            toAppend = np.concatenate((toAppend, [seedChargeFlat[i]]))
        if settings['clusterSize'] is True:
            toAppend = np.concatenate((toAppend, [clsSizeFlat[i]]))
        if settings['coordinates'] is True:
            toAppend = np.concatenate((toAppend, [xcoords[i]]))
            toAppend = np.concatenate((toAppend, [ycoords[i]]))
            toAppend = np.concatenate((toAppend, [zcoords[i]]))
        if settings['uvPosition'] is True:
            toAppend = np.concatenate((toAppend, [uPositionFlat[i]]))
            toAppend = np.concatenate((toAppend, [vPositionFlat[i]]))
        if settings['sensorID'] is True:
            toAppend = np.concatenate((toAppend, [sensorIDFlat[i]]))
        if settings['layer'] is True:
            toAppend = np.concatenate((toAppend, [layers[i]]))
        if settings['ladder'] is True:
            toAppend = np.concatenate((toAppend, [ladders[i]]))
        if settings['momenta'] is True:
            toAppend = np.concatenate((toAppend, [momentumXFlat[i]]))
            toAppend = np.concatenate((toAppend, [momentumYFlat[i]]))
            toAppend = np.concatenate((toAppend, [momentumZFlat[i]]))
        if settings['pdg'] is True:
            toAppend = np.concatenate((toAppend, [pdgFlat[i]]))
        forArray.append(toAppend[1:])

    saveArr = np.array(forArray)



# array speichern
if settings['saveFormat'] == 'structured' or settings['saveFormat'] == 'array':
    print('saving data array as {} to {}'.format(settings['outFile'], settings['outPath']))
    np.save('{}/{}.npy'.format(settings['outPath'], settings['outFile']), saveArr)

if settings['saveFormat'] == 'txt':
    print('saving txt data as {} to {}'.format(settings['outFile'], settings['outPath']))
    with open('{}/{}.txt'.format(settings['outPath'], settings['outFile']), 'w') as f:
        for line in saveArr:
            for item in line:
                f.write(str(item))
                f.write(' ')
            f.write('\n')