import os
import pandas as pd


organs = ['logicalHead', 'logicalSkull', 'logicalBrain', 'logicalTrunk', 'logicalLeftLeg', 'logicalRightLeg', 'logicalLeftArmBone', 'logicalRightArmBone', 'logicalLeftLegBone', 'logicalRightLegBone', 'logicalUpperSpine', 'logicalLeftScapula', 'logicalRightScapula', 'logicalLeftAdrenal', 'logicalRightAdrenal', 'logicalThymus', 'logicalLeftClavicle', 'logicalRightClavicle', 'logicalSmallIntestine', 'logicalRibCage', 'logicalMiddleLowerSpine', 'logicalPelvis', 'logicalStomach', 'logicalUpperLargeIntestine', 'logicalLowerLargeIntestine', 'logicalSpleen', 'logicalPancreas', 'logicalLiver', 'logicalLeftKidney', 'logicalRightKidney', 'logicalUrinaryBladder', 'logicalLeftLung', 'logicalRightLung', 'logicalThyroid', 'logicalLeftOvary', 'logicalRightOvary', 'logicalUterus', 'logicalLeftBreast', 'logicalRightBreast', 'logicalMaleGenitalia', 'logicalLeftTeste', 'logicalRightTeste']
particles = ['TE', 'TP', 'GCR_H', 'GCR_He', 'GCR_Li', 'GCR_Be', 'GCR_B', 'GCR_C', 'GCR_N', 'GCR_O', 'GCR_F', 'GCR_Ne', 'GCR_Na', 'GCR_Mg', 'GCR_Al', 'GCR_Si', 'GCR_P', 'GCR_S', 'GCR_Cl', 'GCR_Ar', 'GCR_K', 'GCR_Ca', 'GCR_Sc', 'GCR_Ti', 'GCR_V', 'GCR_Cr', 'GCR_Mn', 'GCR_Fe', 'GCR_Co', 'GCR_Ni']

edepRes = {}
edep2Res = {}
rmsRes = {}

files = os.listdir('result')
for file in files:
    f = open('result/' + file, 'r', encoding='utf-8')
    data = f.readlines()
    data = [[s.split(': ')[1].strip().replace('\n', '').replace(' MeV', '') for s in row.split(' | ')] for row in data]

    edep = {}
    edep2 = {}
    rms = {}
    for row in data:
        organ = row[0]
        edep[organ] = float(row[1])
        edep2[organ] = float(row[2])
        rms[organ] = float(row[3])

    particle = file.replace('edep_', '').replace('.txt', '')
    edepRes[particle] = pd.Series(edep, index=organs)
    edep2Res[particle] = pd.Series(edep2, index=organs)
    rmsRes[particle] = pd.Series(rms, index=organs)

edepRes = pd.DataFrame(edepRes)
edep2Res = pd.DataFrame(edep2Res)
rmsRes = pd.DataFrame(rmsRes)

edepRes['Organ Sum'] = edepRes.apply(lambda x: x.sum(), axis=1)
edepRes.loc['Particle Sum'] = edepRes.apply(lambda x: x.sum())

edep2Res['Organ Sum'] = edep2Res.apply(lambda x: x.sum(), axis=1)
edep2Res.loc['Particle Sum'] = edep2Res.apply(lambda x: x.sum())

rmsRes['Organ Sum'] = rmsRes.apply(lambda x: x.sum(), axis=1)
rmsRes.loc['Particle Sum'] = rmsRes.apply(lambda x: x.sum())

edepRes.to_csv('result/eDep.csv')
edepRes.to_csv('result/eDep2.csv')
edepRes.to_csv('result/rsm.csv')