"""Architecture-independent collider-object feature construction."""

from __future__ import annotations

from collections.abc import Mapping, Sequence
from numbers import Number
from typing import Any

import numpy as np
import torch

NODE_FEATURE_NAMES = (
    "pt",
    "eta",
    "phi",
    "energy",
    "btag",
    "charge",
    "node_type",
)

_CALCULATED_ENERGY = "CALC_E"
_NODE_TYPE = "NODE_TYPE"


def _as_feature_tensor(values: list[Any], dtype: torch.dtype) -> torch.Tensor:
    return torch.as_tensor(np.asarray(values, dtype=np.float32), dtype=dtype)


def build_node_features(
    event: Mapping[str, Any],
    feature_branches: Sequence[Any],
    object_types: Sequence[str],
    scales: Sequence[Number] | torch.Tensor,
    *,
    dtype: torch.dtype = torch.float32,
) -> tuple[torch.Tensor, list[int]]:
    """Build dense node features from one collider event.

    ``feature_branches`` follows the active legacy schema: each item describes
    one output column and contains one branch name or constant per object
    type. ``object_types`` contains ``"vector"`` or ``"single"`` for each
    object type. Rows are concatenated in object-type order.

    The returned tensor has columns in :data:`NODE_FEATURE_NAMES` order for
    the active schema, and the second result records rows contributed by each
    object type. Inputs are read but never modified.
    """
    if not feature_branches:
        return torch.empty((0, 0), dtype=dtype), []
    first_column = feature_branches[0]
    if len(first_column) != len(object_types):
        raise ValueError("one branch specification is required per object type")

    lengths: list[int] = []
    for branch, object_type in zip(first_column, object_types):
        if object_type == "single":
            lengths.append(1)
        elif object_type == "vector":
            lengths.append(len(event[branch]))
        else:
            raise ValueError(f"unknown object type: {object_type!r}")

    columns: list[torch.Tensor] = []
    for column_index, specification in enumerate(feature_branches):
        if specification == _CALCULATED_ENERGY:
            columns.append(columns[0] * torch.cosh(columns[1]))
            continue
        if specification == _NODE_TYPE:
            values = [
                object_type_index
                for object_type_index, length in enumerate(lengths)
                for _ in range(length)
            ]
            columns.append(torch.tensor(values, dtype=dtype))
            continue
        if len(specification) != len(object_types):
            raise ValueError(
                f"feature column {column_index} has the wrong number of branches"
            )

        values: list[Any] = []
        for object_type_index, (length, branch, object_type) in enumerate(
            zip(lengths, specification, object_types)
        ):
            if isinstance(branch, Number):
                values.extend([branch] * length)
            elif branch == _CALCULATED_ENERGY:
                start = sum(lengths[:object_type_index])
                stop = start + length
                values.extend(
                    (
                        columns[0][start:stop] * torch.cosh(columns[1][start:stop])
                    ).tolist()
                )
            elif object_type == "single":
                values.append(event[branch])
            else:
                values.extend(event[branch])
        columns.append(_as_feature_tensor(values, dtype))

    features = torch.stack(columns, dim=1)
    scale_tensor = torch.as_tensor(scales, dtype=dtype)
    if scale_tensor.ndim != 1 or scale_tensor.numel() != features.shape[1]:
        raise ValueError("scales must contain one value per feature column")
    return features * scale_tensor, lengths


def build_object_features(
    event: Mapping[str, Any],
    feature_branches: Sequence[Any],
    object_types: Sequence[str],
    scales: Sequence[Number] | torch.Tensor,
    *,
    dtype: torch.dtype = torch.float32,
) -> tuple[torch.Tensor, list[int]]:
    """Alias for :func:`build_node_features` using domain-neutral wording."""
    return build_node_features(
        event, feature_branches, object_types, scales, dtype=dtype
    )