#!/usr/bin/env python3
import numba as nb
import numpy as np
import awkward1 as ak1
import uproot

from km3io.definitions import reconstruction as krec
from km3io.definitions import trigger as ktrg
from km3io.definitions import fitparameters as kfit

# 110 MB based on the size of the largest basket found so far in km3net
BASKET_CACHE_SIZE = 110 * 1024**2
BASKET_CACHE = uproot.cache.ThreadSafeArrayCache(BASKET_CACHE_SIZE)


class cached_property:
    """A simple cache decorator for properties."""
    def __init__(self, function):
        self.function = function

    def __get__(self, obj, cls):
        if obj is None:
            return self
        prop = obj.__dict__[self.function.__name__] = self.function(obj)
        return prop


def unfold_indices(obj, indices):
    """Unfolds an index chain and returns the corresponding item"""
    original_obj = obj
    for depth, idx in enumerate(indices):
        try:
            obj = obj[idx]
        except IndexError:
            raise IndexError(
                "IndexError while accessing an item from '{}' at depth {} ({}) "
                "using the index chain {}".format(repr(original_obj), depth,
                                                  idx, indices))
    return obj


def to_num(value):
    """Convert a value to a numerical one if possible"""
    for converter in (int, float):
        try:
            return converter(value)
        except (ValueError, TypeError):
            pass
    return value


@nb.jit(nopython=True)
def unique(array, dtype=np.int64):
    """Return the unique elements of an array with a given dtype.

    The performance is better for pre-sorted input arrays.

    """
    n = len(array)
    out = np.empty(n, dtype)
    last = array[0]
    entry_idx = 0
    out[entry_idx] = last
    for i in range(1, n):
        current = array[i]
        if current == last:  # shortcut for sorted arrays
            continue
        already_present = False
        for j in range(entry_idx + 1):
            if current == out[j]:
                already_present = True
                break
        if not already_present:
            entry_idx += 1
            out[entry_idx] = current
        last = current
    return out[:entry_idx + 1]


@nb.jit(nopython=True)
def uniquecount(array, dtype=np.int64):
    """Count the number of unique elements in a jagged Awkward1 array."""
    n = len(array)
    out = np.empty(n, dtype)
    for i in range(n):
        sub_array = array[i]
        if len(sub_array) == 0:
            out[i] = 0
        else:
            out[i] = len(unique(sub_array))
    return out


def w2list_genhen_keys():
    """names of the w2list parameters as defined in the official
    KM3NeT-Dataformat for genhen.

    Returns
    -------
    dict_keys
        genhen w2list keys.
    """
    return krec.w2list_genhen.keys()


def w2list_gseagen_keys():
    """names of the w2list parameters as defined in the official
    KM3NeT-Dataformat for gseagen.

    Returns
    -------
    dict_keys
        gseagen w2list keys.
    """
    return krec.w2list_gseagen.keys()


def get_w2list_param(events, generator, param):
    """get all the values of a specific parameter from the w2list
    in offline neutrino files.

    Parameters
    ----------
    events : class km3io.offline.OfflineBranch
        events class in offline neutrino files.
    generator : str
        the name of the software generating neutrinos, it is either
        'genhen' or 'gseagen'. 
    param : str
        the name of the parameters found in w2list as defined in the
        KM3NeT-Dataformat for both genhen and gseagen.

    Returns
    -------
    awkward array
        array of the values of interest.
    """
    if generator == "gseagen":
        return events.w2list[:, krec.w2list_gseagen[param]]
    if generator == "genhen":
        return events.w2list[:, krec.w2list_genhen[param]]


def rec_types():
    """name of the reconstruction type as defined in the official
    KM3NeT-Dataformat.

    Returns
    -------
    dict_keys
        reconstruction types.
    """
    return krec.keys()


def fitinf(fitparam, tracks):
    """Access fit parameters in tracks.fitinf.

    Parameters
    ----------
    fitparam : str
        the fit parameter name according to fitparameters defined in
        KM3NeT-Dataformat.
    tracks : class km3io.offline.OfflineBranch
        the tracks class. both full tracks branch or a slice of the
        tracks branch (example tracks[:, 0]) work.

    Returns
    -------
    awkward array
        awkward array of the values of the fit parameter requested.
    """
    fit = tracks.fitinf
    index = kfit[fitparam]
    try:
        params = fit[count_nested(fit, axis=2) > index]
        return ak1.Array([i[:, index] for i in params])
    except ValueError:
        # This is the case for tracks[:, 0] or any other selection.
        params = fit[count_nested(fit, axis=1) > index]
        return params[:, index]


def fitparams():
    """name of the fit parameters as defined in the official
    KM3NeT-Dataformat.

    Returns
    -------
    dict_keys
        fit parameters keys.
    """
    return kfit.keys()


def count_nested(Array, axis=0):
    """count elements in a nested awkward Array.

    Parameters
    ----------
    Array : Awkward1 Array
        Array of data. Example tracks.fitinf or tracks.rec_stages.
    axis : int, optional
        axis = 0: to count elements in the outmost level of nesting.
        axis = 1: to count elements in the first level of nesting.
        axis = 2: to count elements in the second level of nesting.

    Returns
    -------
    awkward1 Array or int
        counts of elements found in a nested awkward1 Array.
    """
    if axis == 0:
        return ak1.num(Array, axis=0)
    if axis == 1:
        return ak1.num(Array, axis=1)
    if axis == 2:
        return ak1.count(Array, axis=2)


@nb.jit(nopython=True)
def _find(rec_stages, stages, builder):
    """construct an awkward1 array with the same structure as tracks.rec_stages.
    When stages are found, the Array is filled with value 1, otherwise it is filled
    with value 0.

    Parameters
    ----------
    rec_stages : awkward1 Array
        tracks.rec_stages .
    stages : awkward1 Array
        reconstruction stages of interest.
    builder : awkward1.highlevel.ArrayBuilder
        awkward1 Array builder.
    """
    for s in rec_stages:
        builder.begin_list()
        for i in s:
            num_stages = len(i)
            if num_stages == len(stages):
                found = 0
                for j in range(num_stages):
                    if i[j] == stages[j]:
                        found += 1
                if found == num_stages:
                    builder.append(1)
                else:
                    builder.append(0)
            else:
                builder.append(0)
        builder.end_list()


def mask(rec_stages, stages):
    """create a mask on tracks.rec_stages .

    Parameters
    ----------
    rec_stages : awkward1 Array
        tracks.rec_stages .
    stages : list
        reconstruction stages of interest.

    Returns
    -------
    awkward1 Array
        an awkward1 Array mask where True corresponds to the positions
        where stages were found. False otherwise.
    """
    builder = ak1.ArrayBuilder()
    _find(rec_stages, ak1.Array(stages), builder)
    return builder.snapshot() == 1


def best_track(tracks, strategy="default", rec_type=None):
    """best track selection based on different strategies

    Parameters
    ----------
    tracks : class km3io.offline.OfflineBranch
        a subset of reconstructed tracks where `events.n_tracks > 0` is always true. 
    strategy : str
        the trategy desired to select the best tracks. It is either: 
            - "first" : to select the first tracks.
            - "default": to select the best tracks (the first ones) corresponding to a specific
            reconstruction algorithm (JGandalf, Jshowerfit, etc). This requires rec_type input.
            Example: best_track(my_tracks, strategy="default", rec_type="JPP_RECONSTRUCTION_TYPE").
    rec_type : str, optional
        reconstruction type as defined in the official KM3NeT-Dataformat.

    Returns
    -------
    class km3io.offline.OfflineBranch
        tracks class with the desired "best tracks" selection.

    Raises
    ------
    ValueError
        ValueError raised when:
            - an invalid strategy is requested.
            - a subset of events with empty tracks is used.
    """
    options = ['first', 'default']
    if strategy not in options:
        raise ValueError("{} not in {}".format(strategy, options))

    n_events = 1 if tracks.is_single else len(tracks)

    if n_events > 1 and any(count_nested(tracks.lik, axis=1) == 0):
        raise ValueError(
            "'events' should not contain empty tracks. Consider applying the mask: events.n_tracks>0"
        )

    if strategy == "first":
        if n_events == 1:
            out = tracks[0]
        else:
            out = tracks[:, 0]

    if strategy == "default" and rec_type is None:
        raise ValueError(
            "rec_type must be provided when the default strategy is used.")

    if strategy == "default" and rec_type is not None:
        if n_events == 1:
            rec_types = tracks[tracks.rec_type == krec[rec_type]]
            len_stages = count_nested(rec_types.rec_stages, axis=1)
            longest = rec_types[len_stages == ak1.max(len_stages, axis=0)]
            out = longest[longest.lik == ak1.max(longest.lik, axis=0)]
        else:
            rec_types = tracks[tracks.rec_type == krec[rec_type]]
            len_stages = count_nested(rec_types.rec_stages, axis=2)
            longest = rec_types[len_stages == ak1.max(len_stages, axis=1)]
            out = longest[longest.lik == ak1.max(longest.lik, axis=1)]

    return out


def get_multiplicity(tracks, rec_stages):
    """tracks selection based on specific reconstruction stages (for multiplicity
    calculations).

    Parameters
    ----------
    tracks : class km3io.offline.OfflineBranch
        tracks or a subste of tracks. 
    rec_stages : list
        the reconstruction stages of interest. Examle: [1, 2, 3, 4, 5].

    Returns
    -------
    class km3io.offline.OfflineBranch
        tracks branch with the desired reconstruction stages only.
    """
    return tracks[mask(tracks.rec_stages, rec_stages)]
