"""Dataclass support for aligned residue-level fields."""

from __future__ import annotations

from abc import ABC, abstractmethod
from collections.abc import Iterable
from dataclasses import Field, dataclass, fields, replace
from typing import Any, Self

import numpy as np

from .esmfold2_misc import concat_objects, slice_any_object

Index = int | list[int] | slice | np.ndarray


def _is_sequential(field: Field[Any]) -> bool:
    return bool(field.metadata.get("sequence", False))


def _sequence_axis(field: Field[Any]) -> int:
    axis = int(field.metadata.get("sequence_dim", 0))
    if axis not in (0, 1):
        raise NotImplementedError("SequentialDataclass supports sequence_dim values zero and one.")
    return axis


def _slice_value(value: Any, index: Index, axis: int) -> Any:
    if axis == 0:
        return slice_any_object(value, index)
    sliced = [slice_any_object(track, index) for track in value]
    return value.__class__(sliced)


def _iter_sequence_lengths(value: Any, axis: int) -> Iterable[int]:
    if axis == 0:
        yield len(value)
    else:
        yield from (len(track) for track in value)


@dataclass(frozen=True)
class SequentialDataclass(ABC):
    """Keep dataclass fields aligned along a shared residue dimension.

    A subclass marks aligned fields with ``metadata={"sequence": True}``.
    ``sequence_dim`` may be zero for a direct sequence or one for a collection
    of aligned tracks. ``join_token`` is passed to the package concatenation
    helper when instances are joined.
    """

    def __post_init__(self) -> None:
        expected = len(self)
        for field in fields(self):
            if not _is_sequential(field) or field.name == "complex":
                continue
            value = getattr(self, field.name)
            if value is None:
                continue
            for actual in _iter_sequence_lengths(value, _sequence_axis(field)):
                if actual != expected:
                    raise ValueError(
                        f"Mismatch in sequence length for field: {field.name}. "
                        f"Expected {expected}, received {actual}"
                    )

    @abstractmethod
    def __len__(self) -> int:
        """Return the shared sequence length."""

        raise NotImplementedError

    def __getitem__(self, index: Index) -> Self:
        """Apply one sequence index to every aligned field."""

        normalized_index: Index = [index] if isinstance(index, int) else index
        updates: dict[str, Any] = {}
        for field in fields(self):
            if not _is_sequential(field):
                continue
            value = getattr(self, field.name)
            if value is not None:
                updates[field.name] = _slice_value(value, normalized_index, _sequence_axis(field))
        return replace(self, **updates)

    @classmethod
    def concat(cls, items: list[Self], **overrides: Any) -> Self:
        """Join aligned fields and retain non-sequential values from the first item."""

        if not items:
            raise ValueError("SequentialDataclass.concat requires at least one item.")

        updates: dict[str, Any] = {}
        for field in fields(cls):
            if not _is_sequential(field):
                continue
            first_value = getattr(items[0], field.name)
            if first_value is None:
                continue
            values = [getattr(item, field.name) for item in items]
            join_token = field.metadata.get("join_token")
            if _sequence_axis(field) == 0:
                updates[field.name] = concat_objects(values, join_token)
            else:
                tracks = [concat_objects(track, join_token) for track in zip(*values, strict=True)]
                updates[field.name] = first_value.__class__(tracks)

        updates.update(overrides)
        return replace(items[0], **updates)


__all__ = ["SequentialDataclass"]