# Filename: calib.py
# pylint: disable=locally-disabled
"""
Calibration.

"""
import numpy as np

from .db import DBManager
from .core import Module
from .hardware import Detector
from .dataclasses import Table
from .tools import istype
from .logger import get_logger

__author__ = "Tamas Gal"
__copyright__ = "Copyright 2016, Tamas Gal and the KM3NeT collaboration."
__credits__ = ["Thomas Heid"]
__license__ = "MIT"
__maintainer__ = "Tamas Gal"
__email__ = "tgal@km3net.de"
__status__ = "Development"

log = get_logger(__name__)

try:
    import numba as nb
except (ImportError, OSError):
    HAVE_NUMBA = False
    jit = lambda f: f
    log.warning(
        "No numba detected, consider `pip install numba` for more speed!"
    )
else:
    try:
        from numba.typed import Dict
    except ImportError:
        log.warning("Please update numba (0.43+) to have dictionary support!")
        HAVE_NUMBA = False
        jit = lambda f: f
    else:
        HAVE_NUMBA = True
        from numba import jit


class Calibration(Module):
    """A very simple, preliminary Module which gives access to the calibration.

    Parameters
    ----------
    apply: bool, optional [default=True]
        Apply the calibration to the hits (add position/direction/t0)?
    filename: str, optional [default=None]
        DetX file with detector description.
    det_id: int, optional
        .detx ID of detector (when retrieving from database).
    t0set: optional
        t0set (when retrieving from database).
    calibset: optional
        calibset (when retrieving from database).
    """
    __name__ = 'Calibration'
    name = 'Calibration'

    def configure(self):
        self._should_apply = self.get('apply', default=True)
        self.filename = self.get('filename')
        self.det_id = self.get('det_id')
        self.run = self.get('run')
        self.t0set = self.get('t0set')
        self.calibset = self.get('calibset')
        self.detector = self.get('detector')
        self._pos_dom_channel = None
        self._dir_dom_channel = None
        self._t0_dom_channel = None
        self._pos_pmt_id = None
        self._dir_pmt_id = None
        self._t0_pmt_id = None
        self._lookup_tables = None    # for Numba

        if self.det_id and self.run:
            self.cprint(
                "Grabbing the calibration for Det ID {} and run {}".format(
                    self.det_id, self.run
                )
            )
            raw_detx = DBManager().detx_for_run(self.det_id, self.run)
            self.detector = Detector(string=raw_detx)
            self._create_dom_channel_lookup()
            self._create_pmt_id_lookup()
            return

        # TODO: deprecation
        if self.get('calibration'):
            self.log.warning(
                "The parameter 'calibration' has been renamed "
                "to 'calibset'. The 'calibration' parameter will be removed "
                "in the next version of KM3Pipe"
            )
            self.calibset = self.get('calibration')

        if self.filename or self.det_id:
            if self.filename is not None:
                self.detector = Detector(filename=self.filename)
            if self.det_id:
                self.detector = Detector(
                    det_id=self.det_id,
                    t0set=self.t0set,
                    calibration=self.calibset
                )

        if self.detector is not None:
            self.log.debug("Creating lookup tables")
            self._create_dom_channel_lookup()
            self._create_pmt_id_lookup()
        else:
            self.log.critical("No detector information loaded.")

    def process(self, blob, key='Hits', outkey='CalibHits'):
        if self._should_apply:
            blob[outkey] = self.apply(blob[key])
        return blob

    def get_detector(self):
        """Return the detector"""
        return self.detector

    def apply_t0(self, hits):
        """Apply only t0s"""
        if HAVE_NUMBA:
            apply_t0_nb(
                hits.time, hits.dom_id, hits.channel_id, self._lookup_tables
            )
        else:
            n = len(hits)
            cal = np.empty(n)
            lookup = self._calib_by_dom_and_channel
            for i in range(n):
                calib = lookup[hits['dom_id'][i]][hits['channel_id'][i]]
                cal[i] = calib[6]
            hits.time += cal
        return hits

    def apply(self, hits, no_copy=False, correct_slewing=False):
        """Add x, y, z, t0 (and du, floor if DataFrame) columns to the hits.

        """
        if not no_copy:
            hits = hits.copy()

        if istype(hits, 'DataFrame'):
            # do we ever see McHits here?
            hits = Table.from_template(hits, 'Hits')

        if hasattr(hits, 'dom_id') and hasattr(hits, 'channel_id'):
            dir_x, dir_y, dir_z, du, floor, pos_x, pos_y, pos_z, t0 = _get_calibration_for_hits(
                hits, self._calib_by_dom_and_channel
            )
        elif hasattr(hits, 'pmt_id'):
            dir_x, dir_y, dir_z, du, floor, pos_x, pos_y, pos_z, t0 = _get_calibration_for_mchits(
                hits, self._calib_by_pmt_id
            )
        else:
            raise TypeError(
                "Don't know how to apply calibration to '{0}'. "
                "We need at least 'dom_id' and 'channel_id', or "
                "'pmt_id'.".format(hits.name)
            )

        if hasattr(hits, 'time'):
            if hits.time.dtype != t0.dtype:
                time = hits.time.astype('f4') + t0.astype('f4')
                hits = hits.drop_columns(['time'])
                hits = hits.append_columns(['time'], [time])
            else:
                hits.time += t0

        hits_data = {}
        for colname in hits.dtype.names:
            hits_data[colname] = hits[colname]
        calib = {
            'dir_x': dir_x,
            'dir_y': dir_y,
            'dir_z': dir_z,
            'du': du.astype(np.uint8),
            'floor': du.astype(np.uint8),
            'pos_x': pos_x,
            'pos_y': pos_y,
            'pos_z': pos_z,
            't0': t0,
        }
        hits_data.update(calib)
        if correct_slewing:
            hits_data['time'] -= slew(hits_data['tot'])
        return Table(
            hits_data,
            h5loc=hits.h5loc,
            split_h5=hits.split_h5,
            name=hits.name
        )

    def _create_dom_channel_lookup(self):
        if HAVE_NUMBA:
            from numba.typed import Dict
            from numba import types
            data = Dict.empty(
                key_type=types.i8, value_type=types.float64[:, :]
            )
        else:
            data = {}
        for pmt in self.detector.pmts:
            if pmt.dom_id not in data:
                data[pmt.dom_id] = np.zeros((31, 9))
            data[pmt.dom_id][pmt.channel_id] = np.asarray([
                pmt.pos_x, pmt.pos_y, pmt.pos_z, pmt.dir_x, pmt.dir_y,
                pmt.dir_z, pmt.t0, pmt.du, pmt.floor
            ],
                                                          dtype=np.float64)
        self._calib_by_dom_and_channel = data
        if HAVE_NUMBA:
            self._lookup_tables = [(dom, cal) for dom, cal in data.items()]

    def _create_pmt_id_lookup(self):
        if HAVE_NUMBA:
            from numba.typed import Dict
            from numba import types
            data = Dict.empty(key_type=types.i8, value_type=types.float64[:])
        else:
            data = {}
        for pmt in self.detector.pmts:
            data[pmt.pmt_id] = np.asarray([
                pmt.pos_x,
                pmt.pos_y,
                pmt.pos_z,
                pmt.dir_x,
                pmt.dir_y,
                pmt.dir_z,
                pmt.t0,
                pmt.du,
                pmt.floor,
            ],
                                          dtype=np.float64)
        self._calib_by_pmt_id = data

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        return "Calibration: det_id({0})".format(self.det_id)


@jit
def apply_t0_nb(times, dom_ids, channel_ids, lookup_tables):
    """Apply t0s using a lookup table of tuples (dom_id, calib)"""
    dom_id = 0
    lookup = np.empty((31, 9))
    for i in range(len(times)):
        cur_dom_id = dom_ids[i]
        if cur_dom_id != dom_id:
            dom_id = cur_dom_id
            for (d, m) in lookup_tables:
                if d == dom_id:
                    np.copyto(lookup, m)
        t0 = lookup[channel_ids[i]][6]
        times[i] += t0


@jit
def _get_calibration_for_hits(hits, lookup):
    """Append the position, direction and t0 columns and add t0 to time"""
    n = len(hits)
    cal = np.empty((n, 9))
    for i in range(n):
        calib = lookup[hits['dom_id'][i]][hits['channel_id'][i]]
        cal[i] = calib
    dir_x = cal[:, 3]
    dir_y = cal[:, 4]
    dir_z = cal[:, 5]
    du = cal[:, 7]
    floor = cal[:, 8]
    pos_x = cal[:, 0]
    pos_y = cal[:, 1]
    pos_z = cal[:, 2]

    t0 = cal[:, 6]

    return [dir_x, dir_y, dir_z, du, floor, pos_x, pos_y, pos_z, t0]


@jit
def _get_calibration_for_mchits(hits, lookup):
    """Append the position, direction and t0 columns and add t0 to time"""
    n_hits = len(hits)
    cal = np.empty((n_hits, 9))
    for i in range(n_hits):
        cal[i] = lookup[hits['pmt_id'][i]]
    dir_x = cal[:, 3]
    dir_y = cal[:, 4]
    dir_z = cal[:, 5]
    du = cal[:, 7]
    floor = cal[:, 8]
    pos_x = cal[:, 0]
    pos_y = cal[:, 1]
    pos_z = cal[:, 2]
    t0 = cal[:, 6]

    return [dir_x, dir_y, dir_z, du, floor, pos_x, pos_y, pos_z, t0]


class CalibrationService(Module):
    """A service which provides calibration routines for hits

    Parameters
    ----------
    filename: str, optional [default=None]
        DetX file with detector description.
    det_id: int, optional
        .detx ID of detector (when retrieving from database).
    t0set: optional
        t0set (when retrieving from database).
    calibset: optional
        calibset (when retrieving from database).
    detector: kp.hardware.Detector, optional
    """
    __name__ = 'Calibration'
    name = 'Calibration'

    def configure(self):
        self.filename = self.get('filename')
        self.det_id = self.get('det_id')
        self.t0set = self.get('t0set')
        self.calibset = self.get('calibset')

        self._detector = self.get('detector')

        if self._detector is not None:
            self._calibration = Calibration(detector=self._detector)

        self._calibration = None

        self.expose(self.calibrate, "calibrate")
        self.expose(self.get_detector, "get_detector")
        self.expose(self.get_calibration, "get_calibration")
        self.expose(self.load_calibration, "load_calibration")
        self.expose(self.correct_slewing, "correct_slewing")

        self.expose(self.detector_deprecation, "detector")

    def load_calibration(
        self, filename=None, det_id=None, t0set=None, calibset=None
    ):
        """Load another calibration"""
        self.filename = filename
        self.det_id = det_id
        self.t0set = t0set
        self.calibset = calibset
        self._detector = None
        self._calibration = None

    @property
    def detector_deprecation(self):
        self.log.deprecation(
            "The service 'detector' of the CalibrationService has been "
            "deprecated. Please use 'get_detector()' instead in future."
        )

    def calibrate(self, hits):
        return self.calibration.apply(hits)

    @property
    def detector(self):
        if self._detector is None:
            self._detector = self.calibration.detector
        return self._detector

    def get_detector(self):
        """Extra getter to be as lazy as possible (expose triggers otherwise"""
        return self.detector

    @property
    def calibration(self):
        if self._calibration is None:
            self._calibration = Calibration(
                filename=self.filename,
                det_id=self.det_id,
                t0set=self.t0set,
                calibset=self.calibset
            )
        return self._calibration

    def get_calibration(self):
        """Extra getter to be as lazy as possible (expose triggers otherwise"""
        return self.calibration

    def correct_slewing(self, hits):
        """Apply time slewing correction to the hit times"""
        hits.time -= slew(hits.tot)


@jit
def slew(tot):
    """Calculate the time slewing of a PMT response for a given ToT


    Parameters
    ----------
    tot: int or np.array(int)
      Time over threshold value of a hit

    Returns
    -------
    time: int or np.array(int)
      Time slewing, which has to be subtracted from the original hit time.
    """

    # First parametrisation
    #p0 = 7.70824
    #p1 = 0.00879447
    #p2 = -0.0621101
    #p3 = -1.90226

    # Second parametrisation
    # p0 =  13.6488662517;
    # p1 =  -0.128744123166;
    # p2 =  -0.0174837749244;
    # p3 =  -4.47119633965;

    # corr = p0 * np.exp(p1 * np.sqrt(tot) + p2 * tot) + p3
    #      this->push_back(  8.01,     #
    # return corr
    corr = np.array([
        8.01, 7.52, 7.05, 6.59, 6.15, 5.74, 5.33, 4.95, 4.58, 4.22, 3.89, 3.56,
        3.25, 2.95, 2.66, 2.39, 2.12, 1.87, 1.63, 1.40, 1.19, 0.98, 0.78, 0.60,
        0.41, 0.24, 0.07, -0.10, -0.27, -0.43, -0.59, -0.75, -0.91, -1.08,
        -1.24, -1.41, -1.56, -1.71, -1.85, -1.98, -2.11, -2.23, -2.35, -2.47,
        -2.58, -2.69, -2.79, -2.89, -2.99, -3.09, -3.19, -3.28, -3.37, -3.46,
        -3.55, -3.64, -3.72, -3.80, -3.88, -3.96, -4.04, -4.12, -4.20, -4.27,
        -4.35, -4.42, -4.49, -4.56, -4.63, -4.70, -4.77, -4.84, -4.90, -4.97,
        -5.03, -5.10, -5.16, -5.22, -5.28, -5.34, -5.40, -5.46, -5.52, -5.58,
        -5.63, -5.69, -5.74, -5.80, -5.85, -5.91, -5.96, -6.01, -6.06, -6.11,
        -6.16, -6.21, -6.26, -6.31, -6.36, -6.41, -6.45, -6.50, -6.55, -6.59,
        -6.64, -6.68, -6.72, -6.77, -6.81, -6.85, -6.89, -6.93, -6.98, -7.02,
        -7.06, -7.09, -7.13, -7.17, -7.21, -7.25, -7.28, -7.32, -7.36, -7.39,
        -7.43, -7.46, -7.50, -7.53, -7.57, -7.60, -7.63, -7.66, -7.70, -7.73,
        -7.76, -7.79, -7.82, -7.85, -7.88, -7.91, -7.94, -7.97, -7.99, -8.02,
        -8.05, -8.07, -8.10, -8.13, -8.15, -8.18, -8.20, -8.23, -8.25, -8.28,
        -8.30, -8.32, -8.34, -8.37, -8.39, -8.41, -8.43, -8.45, -8.47, -8.49,
        -8.51, -8.53, -8.55, -8.57, -8.59, -8.61, -8.62, -8.64, -8.66, -8.67,
        -8.69, -8.70, -8.72, -8.74, -8.75, -8.76, -8.78, -8.79, -8.81, -8.82,
        -8.83, -8.84, -8.86, -8.87, -8.88, -8.89, -8.90, -8.92, -8.93, -8.94,
        -8.95, -8.96, -8.97, -8.98, -9.00, -9.01, -9.02, -9.04, -9.04, -9.04,
        -9.04, -9.04, -9.04, -9.04, -9.04, -9.04, -9.04, -9.04, -9.04, -9.04,
        -9.04, -9.04, -9.04, -9.04, -9.04, -9.04, -9.04, -9.04, -9.04, -9.04,
        -9.04, -9.04, -9.04, -9.04, -9.04, -9.04, -9.04, -9.04, -9.04, -9.04,
        -9.04, -9.04, -9.04, -9.04, -9.04, -9.04, -9.04, -9.04, -9.04, -9.04,
        -9.04, -9.04, -9.04, -9.04, -9.04, -9.04, -9.04, -9.04, -9.04, -9.04,
        -9.04, -9.04
    ])
    return corr[tot]
