#!/usr/bin/env python
# vim:set ts=4 sts=4 sw=4 et:
"""
Pump for the Aanet data format.

This is undoubtedly the ugliest module in the entire framework.
If you have a way to read aanet files via the Jpp interface,
your pull request is more than welcome!
"""
from collections import defaultdict
import os.path

import numpy as np

from km3pipe.core import Pump, Blob
from km3pipe.dataclasses import Table
from km3pipe.logger import get_logger

log = get_logger(__name__)  # pylint: disable=C0103

__author__ = "Moritz Lotze and Tamas Gal"
__copyright__ = "Copyright 2016, Tamas Gal and the KM3NeT collaboration."
__credits__ = "Thomas Heid, Liam Quinn, Javier Barrios Martí"
__license__ = "MIT"
__maintainer__ = "Moritz Lotze and Tamas Gal"
__email__ = "tgal@km3net.de"
__status__ = "Development"


FITINF2NUM = {
    'JGANDALF_BETA0_RAD': 0,                # angular resolution [rad]
    'JGANDALF_BETA1_RAD': 1,                # angular resolution [rad]
    'JGANDALF_CHI2': 2,                     # chi2
    'JGANDALF_NUMBER_OF_HITS': 3,           # number of hits
    'JENERGY_ENERGY': 4,                    # uncorrected energy [GeV]
    'JENERGY_CHI2': 5,                      # chi2
    'JGANDALF_LAMBDA': 6,                   # control parameter
    'JGANDALF_NUMBER_OF_ITERATIONS': 7,     # number of iterations
    'JSTART_NPE_MIP': 8,                    # number of photo-electrons
                                            # up to the barycentre
    'JSTART_NPE_MIP_TOTAL': 9,              # number of photo-electrons
                                            # along the whole track
    'JSTART_LENGTH_METRES': 10,             # distance between first
                                            # and last hits in metres
    'JVETO_NPE': 11,                        # number of photo-electrons
    'JVETO_NUMBER_OF_HITS': 12,             # number of hits
    'JENERGY_MUON_RANGE_METRES': 13,        # range of a muon with the
                                            # reconstructed energy [m]
    'JENERGY_NOISE_LIKELIHOOD': 14,         # log likelihood of every hit
                                            # being K40
    'JENERGY_NDF': 15,                      # number of degrees of freedom
    'JENERGY_NUMBER_OF_HITS': 16,           # number of hits
    'JCOPY_Z_M': 17,                        # true vertex position along
                                            # track [m]
}


RECO2NUM = {
    'JMUONBEGIN': 0,        	  # Start muon fit applications
    'JMUONPREFIT': 1,             # JPrefit.cc
    'JMUONSIMPLEX': 2,            # JSimplex.cc
    'JMUONGANDALF': 3,            # JGandalf.cc
    'JMUONENERGY': 4,             # JEnergy.cc
    'JMUONSTART': 5,              # JStart.cc
    'JMUONEND': 6,                # Termination muon fit applications
    'LineFit': 7,                 # An angular reco guess.
                                  # It could be a seed for JPrefit

    'JSHOWERBEGIN': 100,       	  # Start shower fit applications
    'JSHOWERPREFIT': 101,         # JShowerPrefit.cc
    'JSHOWEREND': 102,            # Termination shower fit applications

    'JPP_REC_TYPE': 4000,         # Jpp reconstruction type for AAnet

    'JUSERBEGIN': 1000,           # Start of user applications
    'JMUONVETO': 1001,            # JVeto.cc
    'JPRESIM': 1002,              # JPreSim_HTR.cc
    'JMUONPATH': 1003,            # JPath.cc
    'JMCEVT': 1004,               # JMCEvt.cc

    'KM3DeltaPos': 10000,         # This is not a fit this gives
                                  # position information only
}


FITINF2NAME = {v: k for k, v in FITINF2NUM.items()}
RECO2NAME = {v: k for k, v in RECO2NUM.items()}


class AanetPump(Pump):
    """A pump for binary Aanet files.

    Parameters
    ----------
    filename: str, optional
        Name of the file to open. If this parameter is not given, ``filenames``
        needs to be specified instead.
    ignore_hits: bool, optional [default=False]
        If true, don't read our the hits/mchits.
    """

    def configure(self):
        self.filename = self.require('filename')
        self.ignore_hits = bool(self.get('ignore_hits'))
        self.header = None
        self.blobs = self.blob_generator()
        self.group_id = 0

    def get_blob(self, index):
        NotImplementedError("Aanet currently does not support indexing.")

    def blob_generator(self):
        """Create a blob generator."""
        # pylint: disable:F0401,W0612
        import aa           # pylint: disablF0401        # noqa
        from ROOT import EventFile           # pylint: disablF0401

        filename = self.filename
        log.info("Reading from file: {0}".format(filename))
        if not os.path.exists(filename):
            log.warning(filename + " not available: continue without it")

        try:
            event_file = EventFile(filename)
        except Exception:
            raise SystemExit("Could not open file")

        log.info("Generating blobs through new aanet API...")
        self.header = self._parse_header(event_file.header)
        for event in event_file:
            log.debug('Reading event...')
            blob = self._read_event(event, filename)
            log.debug('Reading header...')
            blob["AaHeader"] = self.header
            self.group_id += 1
            yield blob
        del event_file

    def _parse_eventinfo(self, event):
        event_id = event.frame_index
        mc_id = event.frame_index - 1
        # run_id = self._get_run_id()
        wgt1, wgt2, wgt3, wgt4 = self._parse_wgts(event.w)
        info = Table({
            'event_id': event_id,
            'mc_id': mc_id,
            # 'run_id': run_id,         # TODO
            'weight_w1': wgt1,
            'weight_w2': wgt2,
            'weight_w3': wgt3,
            'weight_w4': wgt4,
            'group_id': self.group_id,
        }, h5loc='/event_info', name='EventInfo')
        return info

    @staticmethod
    def _parse_wgts(wgt):
        if len(wgt) == 3:
            wgt1, wgt2, wgt3 = wgt
            wgt4 = np.nan
        elif len(wgt) == 4:
            # what the hell is w4?
            wgt1, wgt2, wgt3, wgt4 = wgt
        else:
            wgt1 = wgt2 = wgt3 = wgt4 = np.nan
        return wgt1, wgt2, wgt3, wgt4

    def _parse_tracks(self, tracks):
        out = defaultdict(list)
        for i, trk in enumerate(tracks):
            self.log.debug('Reading Track #{}...'.format(i))
            trk_type = trk.rec_type
            try:
                trk_name = RECO2NAME[trk_type]
            except KeyError:
                trk_name = "Generic_Track_#{}".format(i)
                self.log.warn(
                        "Unknown Reconstruction type! "
                        "Setting to '{}'".format(trk_name)
                )
            trk_dict = self._read_track(trk)
            out[trk_name].append(Table(
                    trk_dict,
                    h5loc='/reco/{}'.format(trk_name.lower()),
                    name=trk_name)
            )
        for key in out:
            name = out[key][0].name
            h5loc = out[key][0].h5loc
            out[key] = Table(
                    np.concatenate(out[key]),
                    name=name,
                    h5loc=h5loc,
            )
        self.log.debug(out)
        return out

    def _read_track(self, trk):
        out = {}
        out['pos_x'] = trk.pos.x
        out['pos_y'] = trk.pos.y
        out['pos_z'] = trk.pos.z
        out['dir_x'] = trk.dir.x
        out['dir_y'] = trk.dir.y
        out['dir_z'] = trk.dir.z
        out['id'] = trk.id
        out['energy'] = trk.E
        out['time'] = trk.t
        out['length'] = trk.len
        out['likelihood'] = trk.lik
        out['rec_type'] = trk.rec_type
        out['group_id'] = self.group_id
        # TODO: hit_ids,
        # TODO: rec_stages,
        self.log.debug('Reading fitinf...')
        fitinf = self._parse_fitinf(trk.fitinf)
        out.update(fitinf)
        return out

    def _parse_fitinf(self, fitinf):
        # iterating empty ROOT vector causes segfaults!
        if len(fitinf) == 0:
            self.log.debug("Found empty fitinf, skipping...")
            return {}

        out = {}
        for i, elem in enumerate(fitinf):
            name = FITINF2NAME[i]
            self.log.debug("Reading fitinf #{} ('{}')...".format(i, name))
            out[name] = elem
        return out

    def _parse_mctracks(self, mctracks):
        out = defaultdict(list)
        for trk in mctracks:
            out['dir_x'].append(trk.dir.x)
            out['dir_y'].append(trk.dir.y)
            out['dir_z'].append(trk.dir.z)
            out['pos_x'].append(trk.pos.x)
            out['pos_y'].append(trk.pos.y)
            out['pos_z'].append(trk.pos.z)
            out['energy'].append(trk.E)
            out['time'].append(trk.t)
            out['type'].append(trk.type)
            out['id'].append(trk.id)
            out['length'].append(trk.len)
            out['bjorkeny'].append(trk.getusr('by'))
            out['interaction_channel'].append(trk.getusr('ichan'))
            out['is_cc'].append(trk.getusr('cc'))
        out['group_id'] = self.group_id
        return Table(out, name='McTracks', h5loc='/mc_tracks')

    def _parse_mchits(self, mchits):
        out = defaultdict(list)
        for hit in mchits:
            out['a'].append(hit.a)
            out['origin'].append(hit.origin)
            out['pmt_id'].append(hit.pmt_id)
            out['time'].append(hit.t)
        out['group_id'] = self.group_id
        return Table(out, name='McHits', h5loc='/mc_hits', split_h5=True)

    def _parse_hits(self, hits):
        out = defaultdict(list)
        for hit in hits:
            out['channel_id'].append(hit.channel_id)
            out['dom_id'].append(hit.dom_id)
            out['time'].append(hit.t)
            out['tot'].append(hit.tot)
            out['triggered'].append(hit.trig)
        out['group_id'] = self.group_id
        return Table(out, name='Hits', h5loc='/hits', split_h5=True)

    @staticmethod
    def _parse_header(header):
        tags = {}
        for key, taglist in header._hdr_dict():
            tags[key] = [k for k in taglist]
        out = {}
        for i, (key, entries) in enumerate(header):
            out[key] = {}
            for j, elem in enumerate(entries.split()):
                if key in tags:
                    try:
                        elem_name = tags[key][j]
                    except IndexError:
                        elem_name = '{}_{}'.format(key, j)
                        log.info("Can't infer field name, "
                                 "setting to '{}'...".format(elem_name))
                else:
                    elem_name = '{}_{}'.format(key, j)
                    log.info("Can't infer field name, "
                             "setting to '{}'...".format(elem_name))
                out[key][elem_name] = elem
        return out

    def _read_event(self, event, filename):
        blob = Blob()
        if self.ignore_hits:
            self.log.debug('Skipping Hits...')
        else:
            self.log.debug('Reading Hits...')
            blob['Hits'] = self._parse_hits(event.hits)
            self.log.debug('Reading McHits...')
            blob['McHits'] = self._parse_mchits(event.mc_hits)
        self.log.debug('Reading McTracks...')
        blob['McTracks'] = self._parse_mctracks(event.mc_trks)
        self.log.debug('Reading EventInfo...')
        blob['EventInfo'] = self._parse_eventinfo(event)
        self.log.debug('Reading Tracks...')
        blob.update(self._parse_tracks(event.trks))
        return blob

    def process(self, blob=None):
        return next(self.blobs)

    def __iter__(self):
        return self

    def __next__(self):
        return next(self.blobs)
