"""Explicit measured-context preparation for stateful trajectory fitting."""

from __future__ import annotations

import numpy as np

from .configuration import ActionCodec2Config
from .physical_codec import _scale, _SO3Runtime
from .trajectory import ActiveTrajectory, TrajectoryLike, prepare_trajectory, trajectory_boundary


def source_differences(
    trajectory: TrajectoryLike, config: ActionCodec2Config
) -> dict[str, tuple[tuple[int, ...], np.ndarray, np.ndarray]]:
    """Measure physical first/second differences with the declared group law."""
    values, dims = prepare_trajectory(trajectory, dimension=config.physical.input_dim, index=0)
    boundary = trajectory_boundary(trajectory, config, dims)
    output = {}
    for name, component in config.physical.components.items():
        present = tuple(d for d in component.dims if d in dims)
        if not present or component.primitive_order == 0:
            continue
        local = [dims.index(d) for d in present]
        indices = [component.dims.index(d) for d in present]
        scale = np.take(_scale(component), indices)
        raw = values[:, local].astype(np.float64) * scale
        previous = None if boundary is None else boundary[list(present)].astype(np.float64) * scale
        if component.space == "additive":
            if component.input_order == 0:
                first = np.diff(raw, axis=0, prepend=np.zeros((1, len(present))))
            elif component.input_order == 1:
                first = raw
            else:
                if previous is None:
                    raise ValueError("order-2 input statistics require previous_first_order")
                first = np.cumsum(raw, axis=0) + previous
            second = (
                np.diff(first, axis=0)
                if previous is None
                else np.diff(first, axis=0, prepend=previous[None])
            )
        else:
            if len(present) != 3:
                raise ValueError("differential statistics require complete SO(3) components")
            runtime = _SO3Runtime(component, (np.zeros(1),) * 3)
            quaternion = runtime._from_representation(raw)
            if component.input_order == 0:
                anchors = np.concatenate((np.array([[0.0, 0.0, 0.0, 1.0]]), quaternion[:-1]))
                increments = runtime._residual(quaternion, anchors)
            else:
                increments = quaternion
            first = runtime._to_representation(increments)
            if previous is None:
                changes = runtime._residual(increments[1:], increments[:-1])
            else:
                anchors = np.concatenate(
                    (runtime._from_representation(previous[None]), increments[:-1])
                )
                changes = runtime._residual(increments, anchors)
            second = runtime._to_representation(changes)
        output[name] = present, first, second
    return output


def bootstrap_second_order_trajectory(
    trajectory: ActiveTrajectory, config: ActionCodec2Config
) -> ActiveTrajectory:
    """Use two observed positions as context and encode the remaining suffix."""
    values, dims = prepare_trajectory(trajectory, dimension=config.physical.input_dim, index=0)
    if len(values) < 3:
        raise ValueError(
            "second-order bootstrap requires two context frames and at least one target"
        )
    if trajectory.previous_first_order is not None:
        raise ValueError("trajectory already has previous_first_order; do not bootstrap it again")
    result = values[2:].copy()
    previous = np.zeros(len(dims), dtype=np.float32)
    for component in config.physical.components.values():
        local = [dims.index(d) for d in component.dims if d in dims]
        if not local:
            continue
        if component.input_order != 0:
            raise ValueError("measured-context bootstrap requires input_order=0")
        if component.primitive_order == 0:
            continue
        if component.space == "additive":
            result[:, local] -= values[1, local]
            previous[local] = values[1, local] - values[0, local]
        else:
            if len(local) != 3:
                raise ValueError("bootstrap requires complete SO(3) components")
            runtime = _SO3Runtime(component, (np.zeros(1),) * 3)
            quaternion = runtime._from_representation(values[:, local] * _scale(component))
            increment = runtime._residual(quaternion[1:2], quaternion[:1])

            from ._so3 import conjugate_quaternion, multiply_quaternions

            anchor = conjugate_quaternion(quaternion[1:2])
            local_path = (
                multiply_quaternions(quaternion[2:], anchor)
                if runtime.left
                else multiply_quaternions(anchor, quaternion[2:])
            )
            result[:, local] = runtime._to_representation(local_path) / _scale(component)
            previous[local] = runtime._to_representation(increment)[0] / _scale(component)
    return ActiveTrajectory(result, dims, previous_first_order=previous)