"""
==================
DOM hits.
==================

This example shows how to create DOM hits statistics to estimate track
distances.
"""

from collections import defaultdict, Counter

import km3pipe as kp
import pandas as pd

import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm

from km3modules import StatusBar
from km3pipe.tools import pld3
import km3pipe.style

km3pipe.style.use('default')


filename = "data/km3net_jul13_90m_muatm50T655.km3_v5r1.JTE_r2356.root.0-499.h5"
geo = kp.Geometry(filename="data/km3net_jul13_90m_r1494_corrected.detx")


class MuonFilter(kp.Module):
    """Write all muons from MCTracks to Muons."""
    def process(self, blob):
        tracks = blob['McTracks']
        muons = [t for t in blob['McTracks'] if t.type == 5]
        blob["Muons"] = kp.dataclasses.TrackSeries(muons, tracks.event_id)
        return blob


class DOMHits(kp.Module):
    """Create histogram with n_hits and distance of hit to track."""
    def configure(self):
        self.hit_statistics = defaultdict(list)

    def process(self, blob):
        hits = blob['Hits']
        muons = blob['Muons']

        highest_energetic_muon = max(muons, key=lambda x: x.energy)
        muon = highest_energetic_muon

        triggered_hits = hits.triggered_hits
        dom_hits = Counter(triggered_hits.dom_id)
        for dom_id, n_hits in dom_hits.iteritems():
            distance = pld3(geo.detector.dom_positions[dom_id],
                            muon.pos,
                            muon.dir)
            self.hit_statistics['n_hits'].append(n_hits)
            self.hit_statistics['distance'].append(distance)
        return blob

    def finish(self):
        df = pd.DataFrame(self.hit_statistics)
        sdf = df[(df['distance'] < 200) & (df['n_hits'] < 50)]
        plt.hist2d(sdf['distance'], sdf['n_hits'], cmap='plasma',
                   bins=(max(sdf['distance'])-1, max(sdf['n_hits'])-1),
                   norm=LogNorm())
        plt.xlabel('Distance between hit and muon track [m]')
        plt.ylabel('Number of hits on DOM')
        plt.show()


pipe = kp.Pipeline()
pipe.attach(kp.io.HDF5Pump, filename=filename)
pipe.attach(StatusBar, every=100)
pipe.attach(MuonFilter)
pipe.attach(DOMHits)
pipe.drain()

Add example to show how to create dom hit statistics
Add example to show how to create dom hit statistics
