from typing import Any, Callable, List, Optional, Sequence, Tuple, TypeVar, cast

from ..config import registry
from ..model import Model
from ..types import ArrayXd, ListXd

ItemT = TypeVar("ItemT")
InT = Sequence[Sequence[ItemT]]
OutT = ListXd
InnerInT = Sequence[ItemT]
InnerOutT = ArrayXd


@registry.layers("with_flatten.v1")
def with_flatten(layer: Model[InnerInT[ItemT], InnerOutT]) -> Model[InT[ItemT], OutT]:
    return Model(f"with_flatten({layer.name})", forward, layers=[layer], init=init)


def forward(
    model: Model[InT, OutT], Xnest: InT, is_train: bool
) -> Tuple[OutT, Callable]:
    layer: Model[InnerInT, InnerOutT] = model.layers[0]
    Xflat = _flatten(Xnest)
    Yflat, backprop_layer = layer(Xflat, is_train)
    # Get the split points. We want n-1 splits for n items.
    arr = layer.ops.asarray1i([len(x) for x in Xnest[:-1]])
    splits = arr.cumsum()
    Ynest = layer.ops.xp.split(Yflat, splits, axis=0)

    def backprop(dYnest: OutT) -> InT:
        dYflat = model.ops.flatten(dYnest)  # type: ignore[arg-type, var-annotated]
        # type ignore necessary for older versions of Mypy/Pydantic
        dXflat = backprop_layer(dYflat)
        dXnest = layer.ops.xp.split(dXflat, splits, axis=-1)
        return dXnest

    return Ynest, backprop


def _flatten(nested: InT) -> InnerInT:
    flat: List = []
    for item in nested:
        flat.extend(item)
    return cast(InT, flat)


def init(
    model: Model[InT, OutT], X: Optional[InT] = None, Y: Optional[OutT] = None
) -> None:
    model.layers[0].initialize(
        _flatten(X) if X is not None else None,
        model.layers[0].ops.xp.hstack(Y) if Y is not None else None,
    )