"""Validated data contracts and explicit Stage 2/3 trajectory alignment."""
from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Dict, Mapping, Optional, Sequence, Tuple

import numpy as np

from .feature_builder import build_feature_matrix, compute_dimensionless, compute_features

HISTORICAL_ALIGNMENT = "historical_post_state_same_action"
CAUSAL_ALIGNMENT = "causal_post_state_to_next_action"
ALIGNMENTS = (HISTORICAL_ALIGNMENT, CAUSAL_ALIGNMENT)
OFFLINE_SPLIT_NAMES = ("train", "validation", "blind")


def contiguous_block_slices(
    n_samples: int,
    fractions: Sequence[float] = (0.6, 0.2, 0.2),
) -> Dict[str, slice]:
    """Return ordered, non-overlapping train/validation/blind blocks."""
    if n_samples < 3:
        raise ValueError("at least three samples are required for offline splits")
    values = np.asarray(fractions, dtype=np.float64)
    if values.shape != (3,) or np.any(values <= 0) or not np.isclose(values.sum(), 1.0):
        raise ValueError("split fractions must be three positive values summing to one")
    train_end = max(1, min(n_samples - 2, int(np.floor(n_samples * values[0]))))
    validation_count = max(1, int(np.floor(n_samples * values[1])))
    validation_end = min(n_samples - 1, train_end + validation_count)
    return {
        "train": slice(0, train_end),
        "validation": slice(train_end, validation_end),
        "blind": slice(validation_end, n_samples),
    }


@dataclass(frozen=True)
class BatchDataset:
    """Aligned symbolic-regression dataset and its provenance."""

    X: np.ndarray
    Y: np.ndarray
    kept_indices: np.ndarray
    dropped_indices: np.ndarray
    metadata: Dict[str, Any]
    features: Dict[str, np.ndarray]

    def __iter__(self):
        """Allow the convenient ``X, Y = dataset`` compatibility idiom."""
        yield self.X
        yield self.Y


def _matrix(name: str, value: np.ndarray, columns: int, rows: Optional[int] = None) -> np.ndarray:
    array = np.asarray(value, dtype=np.float64)
    if array.ndim != 2 or array.shape[1] != columns:
        raise ValueError("{} must have shape (T, {}), got {}".format(name, columns, array.shape))
    if rows is not None and array.shape[0] != rows:
        raise ValueError("{} has {} rows; expected {}".format(name, array.shape[0], rows))
    if not np.all(np.isfinite(array)):
        raise ValueError("{} contains non-finite values".format(name))
    return array


def validate_controlled_trajectory(
    sensors: np.ndarray,
    forces: np.ndarray,
    actions: np.ndarray,
    target_forces: Optional[np.ndarray] = None,
    *,
    illusion: bool = False,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]:
    """Validate the legacy ``controlled.npz`` array schema.

    The legacy contract is sensors ``(T,6)``, forces ``(T,6)``, actions
    ``(T,3)``, and optional illusion target forces ``(T,2)``.
    """
    s = _matrix("sensors", sensors, 6)
    f = _matrix("forces", forces, 6, s.shape[0])
    a = _matrix("actions", actions, 3, s.shape[0])
    if illusion and target_forces is None:
        raise ValueError("illusion trajectories require target_forces with shape (T, 2)")
    tf = None if target_forces is None else _matrix("target_forces", target_forces, 2, s.shape[0])
    return s, f, a, tf


def actions_to_alpha(
    actions: np.ndarray,
    *,
    action_mode: str,
    u0: float,
    action_scale: float = 1.0,
    action_bias: Sequence[float] = (0.0, 0.0, 0.0),
) -> np.ndarray:
    """Convert normalized, physical-omega, or alpha actions to alpha."""
    a = np.asarray(actions, dtype=np.float64)
    if action_mode == "alpha":
        return a.copy()
    if action_mode == "physical":
        return a / u0
    if action_mode == "normalized":
        bias = np.asarray(action_bias, dtype=np.float64)
        if bias.shape != (3,):
            raise ValueError("action_bias must have shape (3,)")
        return a * float(action_scale) + bias
    raise ValueError("action_mode must be 'normalized', 'physical', or 'alpha'")


def build_batch_dataset(
    sensors: np.ndarray,
    forces: np.ndarray,
    actions: np.ndarray,
    feat_keys: Sequence[str],
    *,
    mu: float,
    dt_c: float,
    alignment: str = HISTORICAL_ALIGNMENT,
    action_mode: str = "physical",
    action_scale: float = 1.0,
    action_bias: Sequence[float] = (0.0, 0.0, 0.0),
    target_forces: Optional[np.ndarray] = None,
    illusion: bool = False,
    u0: float = 0.01,
    d: float = 20.0,
    rho: float = 1.0,
    include_mu: Optional[bool] = None,
    include_cos_sin: bool = False,
    add_bias: bool = False,
    warmup: int = 2,
    sample_segment: str = "all",
    transient_count: int = 50,
    strict: bool = True,
) -> BatchDataset:
    """Build an aligned Stage 2/3 dataset from a controlled trajectory.

    ``historical_post_state_same_action`` reproduces the legacy pairing
    ``post_state[i] -> action[i]`` with action history ``i-1, i-2``.
    ``causal_post_state_to_next_action`` pairs ``post_state[i] -> action[i+1]``
    and therefore uses the causally available history ``i, i-1``.
    """
    if alignment not in ALIGNMENTS:
        raise ValueError("unknown alignment {!r}; expected one of {}".format(alignment, ALIGNMENTS))
    if dt_c <= 0:
        raise ValueError("dt_c must be positive")
    if warmup < 0:
        raise ValueError("warmup must be non-negative")
    if sample_segment not in {"all", "transient", "steady"}:
        raise ValueError("sample_segment must be 'all', 'transient', or 'steady'")
    if transient_count < 0:
        raise ValueError("transient_count must be non-negative")

    s, f, a_raw, tf = validate_controlled_trajectory(
        sensors, forces, actions, target_forces, illusion=illusion
    )
    alpha = actions_to_alpha(
        a_raw, action_mode=action_mode, u0=u0,
        action_scale=action_scale, action_bias=action_bias,
    )
    T = s.shape[0]
    a1 = np.zeros((T, 3), dtype=np.float64)
    a2 = np.zeros((T, 3), dtype=np.float64)
    if alignment == HISTORICAL_ALIGNMENT:
        a1[1:] = alpha[:-1]
        a2[2:] = alpha[:-2]
        state_indices = np.arange(T, dtype=np.int64)
        target_indices = state_indices.copy()
    else:
        a1[:] = alpha
        a2[1:] = alpha[:-1]
        state_indices = np.arange(max(T - 1, 0), dtype=np.int64)
        target_indices = state_indices + 1

    dim = compute_dimensionless(s, f, u0=u0, d=d, rho=rho)
    features = compute_features(
        dim, a1, a2, mu, alpha_mode=True,
        include_mu=(any(key.startswith("mu") for key in feat_keys) if include_mu is None else include_mu),
        include_cos_sin=include_cos_sin, dt_c=dt_c, u0=u0,
        target_forces=tf, rho=rho, sensors_raw=s, forces_raw=f,
    )
    full_X = build_feature_matrix(features, list(feat_keys), add_bias=add_bias, strict=strict)
    valid = state_indices >= warmup
    if sample_segment == "transient":
        valid &= state_indices < transient_count
    elif sample_segment == "steady":
        valid &= state_indices >= transient_count
    candidate = state_indices[valid]
    target_candidate = target_indices[valid]
    X = full_X[candidate]
    Y = alpha[target_candidate]
    kept = candidate
    dropped = np.setdiff1d(np.arange(T, dtype=np.int64), kept, assume_unique=True)
    metadata = {
        "alignment": alignment,
        "action_mode": action_mode,
        "feature_keys": list(feat_keys),
        "add_bias": bool(add_bias),
        "strict": bool(strict),
        "warmup": int(warmup),
        "sample_segment": sample_segment,
        "transient_count": int(transient_count),
        "sample_segments": {
            "transient": [int(warmup), int(min(transient_count, T))],
            "steady": [int(max(warmup, transient_count)), int(T)],
        },
        "state_indices": kept.copy(),
        "target_indices": target_candidate.copy(),
        "history_policy": "drop_warmup_zero_initialized",
        "illusion": bool(illusion),
        "dt_c": float(dt_c),
    }
    return BatchDataset(X, Y, kept, dropped, metadata, features)


def build_batch_dataset_from_mapping(
    data: Mapping[str, np.ndarray], feat_keys: Sequence[str], **kwargs: Any
) -> BatchDataset:
    """Build from a loaded legacy NPZ/mapping while preserving validation."""
    required = [key for key in ("sensors", "forces", "actions") if key not in data]
    if required:
        raise KeyError("controlled trajectory missing required array(s): {}".format(", ".join(required)))
    return build_batch_dataset(
        data["sensors"], data["forces"], data["actions"], feat_keys,
        target_forces=data["target_forces"] if "target_forces" in data else None,
        **kwargs
    )
