import ROOT
import os.path

ROOT.gROOT.ProcessLine('.L '+os.path.join(os.path.dirname(__file__), 'root.C+g'))

import ROOT
import event

class RootWriter(object):
    def __init__(self, filename):
        self.filename = filename
        self.file = ROOT.TFile(filename, 'RECREATE')

        self.T = ROOT.TTree('T', 'Chroma events')
        self.ev = ROOT.Event()
        self.T.Branch('ev', self.ev)

    def write_event(self, pyev):
        '''Write an event.Event object to the ROOT tree as a ROOT.Event object.'''
        self.ev.event_id = pyev.event_id

        self.ev.mc.particle = pyev.particle_name
        self.ev.mc.gen_position.SetXYZ(*pyev.gen_position)
        self.ev.mc.gen_direction.SetXYZ(*pyev.gen_direction)
        self.ev.mc.gen_total_energy = pyev.gen_total_energy
        self.ev.mc.nphoton = pyev.nphoton

        if pyev.photon_start is not None:
            photons = pyev.photon_start
            ROOT.fill_photons(self.ev, True,
                              len(photons.positions), 
                              np.ravel(photons.positions),
                              np.ravel(photons.directions),
                              np.ravel(photons.polarizations),
                              photons.wavelengths, photons.times)
        if pyev.photon_stop is not None:
            photons = photon_stop
            ROOT.fill_photons(self.ev, True,
                              len(photons.positions), 
                              np.ravel(photons.positions),
                              np.ravel(photons.directions),
                              np.ravel(photons.polarizations),
                              photons.wavelengths, photons.times)

        self.ev.mc.subtrack.resize(0)
        if pyev.subtracks is not None:
            self.ev.mc.subtrack.resize(len(pyev.subtracks))
            for i, subtrack in enumerate(pyev.subtracks):
                self.ev.mc.subtrack[i].name = subtrack.particle_name
                self.ev.mc.subtrack[i].position.SetXYZ(*subtrack.position)
                self.ev.mc.subtrack[i].direction.SetXYZ(*subtrack.direction)
                self.ev.mc.subtrack[i].start_time = subtrack.start_time
                self.ev.mc.subtrack[i].total_energy = subtrack.total_energy

        ROOT.fill_hits(self.ev, len(pyev.hits.t), pyev.hits.t, pyev.hits.q, pyev.hits.histories)
        self.T.Fill()

    def close(self):
        self.T.Write()
        self.file.Close()