# -*- coding: utf-8 -*-
import numpy as np
from numpy import concatenate as join, ndarray

from awkward import unflatten, Array

def shape2d(arr): return arr.shape[:2]
def cut(shp): return np.full(shp[0], shp[1], dtype=np.int64)
def flatten(arr): return arr.flatten()


class JaggedArray(Array):
    """
    A numba-jittable class that handles matrices with a variable 
    number of columns per row.

    The class is actually an interface to `awkward.Array`,
    with some additional features, specific to 2d jagged arrays.

    Parameters
    ----------
    data : Iterable
        A 2d numpy array or a list of such objects.
    behaviour
        Passed forward to `awkward.Array`, see it's documentation
        for the details.
    with_name
        Passed forward to `awkward.Array`, see it's documentation
        for the details.
    check_valid
        Passed forward to `awkward.Array`, see it's documentation
        for the details.
    cache
        Passed forward to `awkward.Array`, see it's documentation
        for the details.
    kernels
        Passed forward to `awkward.Array`, see it's documentation
        for the details.
    cuts : Iterable
        An iterable that tells how to unflatten an 1d array into a 
        2d jagged shape.

    Returns
    -------
    JaggedArray

    """

    def __init__(self, data, behavior=None, with_name=None,
                 check_valid=False, cache=None, kernels=None,
                 cuts=None) -> 'JaggedArray':
        if isinstance(data, np.ndarray):
            nD = len(data.shape)
            assert nD <= 2, "Only 2 dimensional arrays are supported!"
            if nD == 1:
                assert isinstance(cuts, np.ndarray)
                data = unflatten(data, cuts)
            else:
                assert nD == 2
        elif isinstance(data, list):
            if isinstance(data[0], ndarray):
                # data is a list of 2d numpy arrays
                assert all(map(lambda arr: len(arr.shape) == 2, data)), \
                    "Only 2 dimensional arrays are supported!"
                # NOTE - implementaion 1
                # > Through the back door, but this is probably the cleanest solution of all.
                # > It only requires to create one python list, without further operations on it.
                data = unflatten(join(list(map(flatten, data))),
                                 join(list(map(cut, map(shape2d, data)))))
                # NOTE - implementaion 2
                #from operator import add
                #from functools import reduce
                # > This also works, but requires data to jump back and forth just to
                # > have a merged list of lists. It also requires to add nested python lists,
                # > which is probably not the quickest operation in the computational world.
                #data = ak.from_iter(reduce(add, map(lambda arr : ak.Array(arr).to_list(), data)))
                # NOTE - implementaion 3
                # > This also works at creation, but fails later at some operation due to
                # > the specific layout generated by ak.concatenate
                #data = ak.concatenate(list(map(lambda arr : ak.Array(arr), data)))
        super().__init__(data, behavior=behavior, with_name=with_name,
                         check_valid=check_valid, cache=cache, kernels=kernels)
