from __future__ import annotations

__all__ = (
    "index_select",
    "lyndon_words",
    "compress",
    "unravel_signature",
    "flatten",
    "term_at",
)

from collections import defaultdict
from functools import partial
from typing import cast

import jax
import jax.numpy as jnp


@jax.jit
def index_select(input: jax.Array, indices: jax.Array) -> jax.Array:
    """
    Select entries in m-level tensor based on given indices
    This function will help compressing log-signatures

    Args:
        input: size (dim, dim, ..., dim)
        indices: size (dim, n)
    Return:
        A 1D jnp.ndarray
    """

    dim = input.shape[0]
    ndim = input.ndim
    n = indices.shape[1]
    assert n <= ndim
    strides = jnp.array([dim**i for i in range(n)])
    # flatten matrix A in C-style
    flattened = input.ravel()
    strides = jnp.array([dim**i for i in range(n)])
    # flatten matrix A in Fortran-style
    flattened = input.ravel("F")

    def _select(index: jax.Array) -> jax.Array:
        """index is a `jnp.ndarray` int"""

        # this is the way to compute the position of
        # (C-style) raveled arrays
        position = jnp.sum(index * strides)
        return flattened[position]

    return jax.vmap(_select)(indices)


def lyndon_words(depth: int, dim: int) -> list[jax.Array]:
    """Generate Lyndon words of length `depth` over an `dim`-symbol alphabet
    Example in Python: https://gist.github.com/dvberkel/1950267

    Args:
        depth: int
        dim: int
    """
    list_of_words = defaultdict(list)
    word = [-1]
    while word:
        word[-1] += 1
        m = len(word)
        # yield word
        list_of_words[m - 1].append(jnp.array(word))
        while len(word) < depth:
            word.append(word[-m])
        while word and word[-1] == dim - 1:
            word.pop()

    return [jnp.stack(list_of_words[i]) for i in range(depth)]


def compress(
    input: list[jax.Array],
    indices: list[jax.Array],
) -> list[jax.Array]:
    """
    Compress expanded log-signatures using Lydon words

    Args:
        input: List of `jnp.ndarray`
        indices: generated by Lyndon words

    Returns:
        A list of compressed `jnp.ndarray`
    """

    return [index_select(term, index) for term, index in zip(input, indices)]


@jax.jit
def flatten(signature: list[jax.Array]) -> jax.Array:
    flattened_terms = tuple(map(jnp.ravel, signature))
    return jnp.concatenate(flattened_terms)


@partial(jax.jit, static_argnums=[0, 1])
def _get_depth(dim: int, depth: int) -> tuple[jax.Array, jax.Array]:
    offset = jax.lax.integer_pow(dim, depth)
    start = dim * (1 - offset) // (1 - dim)

    return offset, start


@partial(jax.jit, static_argnums=[1, 2, 3, 4])
def _term_at(
    flattened_signature: jax.Array,
    dim: int,
    term_i: int,
    start: int,
    offset: int,
) -> jax.Array:
    return jax.lax.dynamic_slice(
        flattened_signature,
        (start,),
        (offset,),
    ).reshape((term_i + 1) * (dim,))


def term_at(
    flattened_signature: jax.Array,
    dim: int,
    term_i: int,
) -> jax.Array:
    start, prev_offset = _get_depth(dim, term_i)

    return cast(
        jax.Array, _term_at(flattened_signature, dim, term_i, start, prev_offset * dim)
    )


def unravel_signature(signature: jax.Array, dim: int, depth: int) -> list[jax.Array]:
    unraveled: list[jax.Array] = []
    start, offset = 0, dim

    for term_i in range(depth):
        unraveled.append(_term_at(signature, dim, term_i, start, offset))
        start += offset
        offset *= dim

    return unraveled
