"""FID adapter for the augmented state of the heterogeneous waves dataset.

The raw state is 392-D = concat(u_coarse[196], v_hat_coarse[196]). The image and the physical diagnostics read
the observed displacement field u only (the velocity is a model input, not the observable). This adapter slices the first 196
coordinates and otherwise delegates to ``HeterogeneousWavesAdapter`` (identical reconstruction,
clip rule, physical diagnostics) -- no numerical difference from the plain-u adapter once sliced.
"""
from __future__ import annotations

import numpy as np

from .adapters import register_adapter
from .heterogeneous_waves import HeterogeneousWavesAdapter

_D_U = 196


@register_adapter("heterogeneous_waves_augmented_coarse")
class HeterogeneousWavesAugmentedAdapter(HeterogeneousWavesAdapter):
    def to_fields(self, states_raw: np.ndarray) -> np.ndarray:
        x = np.asarray(states_raw, dtype=np.float64)
        if x.ndim == 1:
            x = x[None]
        if x.shape[1] not in (_D_U, 2 * _D_U):
            raise ValueError(f"expected {_D_U}-D or {2 * _D_U}-D (augmented) states, got D={x.shape[1]}")
        return super().to_fields(x[:, :_D_U])

    def _moments(self, v: np.ndarray) -> tuple[float, float]:
        v = np.asarray(v, np.float64)
        if v.shape[1] == 2 * _D_U:
            v = v[:, :_D_U]
        return super()._moments(v)

    def card(self) -> dict:
        c = super().card()
        c.update({"state": "392-D augmented state (196-D coarse displacement u + 196-D finite-difference "
                           "coarse velocity v_hat); this adapter reads only the first 196 (displacement) "
                           "coordinates -- identical reconstruction to heterogeneous_waves_coarse"})
        return c