# Copyright© 2025-2026 Gesellschaft zur Förderung der angewandten Forschung e.V.
# acting on behalf of its Fraunhofer Institut für Graphische Datenverarbeitung.
# Licensed under the EUPL. See LICENSE.txt.

from typing import Sequence
from math import prod, ceil

from .backend import ArrayNamespace, get_index_dtype, ArrayLike
from .dimension import Dimension
from .trainbase import TrainBase
from .trainshape import TrainShape
from .exactaddition import ExactAddition

def slice_operator[T: ArrayLike](xp: ArrayNamespace[T], dim: Dimension, step: int) -> TrainBase[T]:
    tmp = step
    idx = 0
    for i in range(len(dim)-1, -1, -1):
        tmp //= dim[i].base
        if tmp == 1:
            idx = i
            break

    if dim.size() / step < step:
        return sum_slice(xp, dim, step)
    elif prod(digit.base for digit in dim[idx:]) % step != 0:
        return factor_slice(xp, dim, step)
    else:
        return factor_slice_exact(xp, dim, step)

def sum_slice(xp: ArrayNamespace, dim: Dimension, step: int):
    idxs = xp.asarray([step], dtype=get_index_dtype(xp))[:,xp.newaxis]
    digits = dim.to_digits(idxs)[:,0]

    add = ExactAddition()

    train = val_cores(xp, dim, 0, 1.0)
    for i in range(1, ceil(dim.size()/step)):
        train = add(train, val_cores(xp, dim, i*step, 1.0))

    return train

def val_cores[T: ArrayLike](xp: ArrayNamespace[T], dim: Dimension, idx: int, val: float) -> TrainBase[T]:
    idxs = xp.asarray([idx], dtype=get_index_dtype(xp))[:,xp.newaxis]
    digits = dim.to_digits(idxs)[:,0]
    cores = []
    for i in range(len(dim)):
        core = xp.zeros((1, dim[i].base, 1))
        core[0,digits[i],0] = 1.0
        cores.append(core)
    cores[0] *= val
    shape = TrainShape(dim, [(d,) for d in dim])
    return TrainBase(shape, cores, copy_data=False)


def factor_slice[T: ArrayLike](xp: ArrayNamespace[T], dim: Dimension, step: int) -> TrainBase[T]:
    cores = []
    offs = xp.asarray([0])

    tmp = dim.size()
    for digit in dim:
        tmp //= digit.base

        noffs = len(offs)

        # generate indices
        left = xp.arange(noffs)[:,None]
        mid = xp.arange(digit.base)[None,:]
        right = (offs[:,None] - tmp*mid) % step

        # get unique columns indices
        roffs = xp.unique_all(right)[0]
        nright = xp.searchsorted(roffs, right)

        # generate matrix
        mat = xp.zeros((noffs, digit.base, roffs.size))
        mat[left, mid, nright] = 1.0
        cores.append(mat)
        offs = roffs
    cores[-1] = cores[-1][:,:,0][:,:,None]

    # generate vector
    shape = TrainShape(dim, [(d,) for d in dim])
    return TrainBase(shape, cores, copy_data=False)

def factor_slice_exact[T: ArrayLike](xp: ArrayNamespace[T], dim: Dimension, step: int) -> TrainBase[T]:
    tmp = step
    idx = 0
    for i in range(len(dim)-1, -1, -1):
        tmp //= dim[i].base
        if tmp == 1:
            idx = i
            break

    if prod(digit.base for digit in dim[idx:]) % step != 0:
        raise ValueError("period must be divisible by dimension size")

    idxs = xp.asarray([step], dtype=get_index_dtype(xp))[:,xp.newaxis]
    digits = dim.to_digits(idxs)[:,0]

    cores = []
    for i in range(len(dim)):
        core = xp.zeros((1, dim[i].base, 1))
        if i < idx:
            core[0,:,0] = 1.0
        else:
            core[0,digits[i],0] = 1.0
        cores.append(core)

    shape = TrainShape(dim, [(d,) for d in dim])
    return TrainBase(shape, cores, copy_data=False)
