"""Stateful online feature construction shared by SR Stage 2 and Stage 3."""
from __future__ import annotations

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

import numpy as np

from .data_contracts import actions_to_alpha
from .feature_builder import build_feature_matrix, compute_dimensionless, compute_features


@dataclass
class FeatureState:
    """History required for one-step feature construction.

    The first step uses current sensors/forces/target as observation history,
    making lag features equal current features and their derivatives zero.
    Missing action history is initialized to zero alpha unless explicit
    previous actions are supplied.
    """

    previous_sensors: Optional[np.ndarray] = None
    previous_forces: Optional[np.ndarray] = None
    previous_target_force: Optional[np.ndarray] = None
    previous_action: Optional[np.ndarray] = None
    previous_action2: Optional[np.ndarray] = None
    step_index: int = 0

    def reset(self) -> None:
        self.previous_sensors = None
        self.previous_forces = None
        self.previous_target_force = None
        self.previous_action = None
        self.previous_action2 = None
        self.step_index = 0


class StatefulFeatureBuilder:
    """Generate online features with exactly the batch ``compute_features`` semantics."""

    def __init__(
        self,
        *,
        mu: float,
        dt_c: float,
        u0: float = 0.01,
        d: float = 20.0,
        rho: float = 1.0,
        action_mode: str = "physical",
        action_scale: float = 1.0,
        action_bias: Sequence[float] = (0.0, 0.0, 0.0),
        include_mu: bool = True,
        include_cos_sin: bool = True,
        require_target: bool = False,
        state: Optional[FeatureState] = None,
    ) -> None:
        if dt_c <= 0:
            raise ValueError("dt_c must be positive")
        self.mu = float(mu)
        self.dt_c = float(dt_c)
        self.u0 = float(u0)
        self.d = float(d)
        self.rho = float(rho)
        self.action_mode = action_mode
        self.action_scale = float(action_scale)
        self.action_bias = tuple(action_bias)
        self.include_mu = bool(include_mu)
        self.include_cos_sin = bool(include_cos_sin)
        self.require_target = bool(require_target)
        self.state = state if state is not None else FeatureState()

    @staticmethod
    def _row(name: str, value: np.ndarray, columns: int) -> np.ndarray:
        row = np.asarray(value, dtype=np.float64)
        if row.shape != (columns,):
            raise ValueError("{} must have shape ({},), got {}".format(name, columns, row.shape))
        return row

    def _alpha(self, action: np.ndarray) -> np.ndarray:
        row = self._row("action", action, 3).reshape(1, 3)
        return actions_to_alpha(
            row, action_mode=self.action_mode, u0=self.u0,
            action_scale=self.action_scale, action_bias=self.action_bias,
        )[0]

    def step(
        self,
        sensors: np.ndarray,
        forces: np.ndarray,
        *,
        previous_action: Optional[np.ndarray] = None,
        previous_action2: Optional[np.ndarray] = None,
        target_force: Optional[np.ndarray] = None,
        previous_target_force: Optional[np.ndarray] = None,
        feat_keys: Optional[Sequence[str]] = None,
        add_bias: bool = False,
        strict: bool = True,
    ) -> Tuple[Dict[str, np.ndarray], Optional[np.ndarray]]:
        """Consume current post-state and return scalar feature dict/vector.

        ``previous_action`` means the action immediately preceding the current
        state (the historical ``a_lag1``). If omitted, stored history is used.
        The current action is not consumed here; call ``observe_action`` after
        selecting/applying it, or pass explicit history on every call.
        """
        s = self._row("sensors", sensors, 6)
        f = self._row("forces", forces, 6)
        if self.require_target and target_force is None:
            raise ValueError("target_force is required for illusion features")
        tf = None if target_force is None else self._row("target_force", target_force, 2)

        p_action = self.state.previous_action if previous_action is None else self._alpha(previous_action)
        p_action2 = self.state.previous_action2 if previous_action2 is None else self._alpha(previous_action2)
        a1 = np.zeros(3) if p_action is None else np.asarray(p_action, dtype=np.float64)
        a2 = np.zeros(3) if p_action2 is None else np.asarray(p_action2, dtype=np.float64)

        ps = self.state.previous_sensors
        pf = self.state.previous_forces
        if ps is None or pf is None:
            ps, pf = s, f
        ptf = self.state.previous_target_force if previous_target_force is None else self._row(
            "previous_target_force", previous_target_force, 2
        )
        if tf is not None and ptf is None:
            ptf = tf

        sensors_raw = np.vstack([ps, s])
        forces_raw = np.vstack([pf, f])
        dim = compute_dimensionless(s.reshape(1, 6), f.reshape(1, 6), u0=self.u0, d=self.d, rho=self.rho)
        target_rows = None if tf is None else np.vstack([ptf, tf])
        feature_arrays = compute_features(
            dim, a1.reshape(1, 3), a2.reshape(1, 3), self.mu,
            alpha_mode=True, include_mu=self.include_mu,
            include_cos_sin=self.include_cos_sin, dt_c=self.dt_c, u0=self.u0,
            target_forces=None if target_rows is None else target_rows[-1:], rho=self.rho,
            sensors_raw=sensors_raw, forces_raw=forces_raw,
        )

        # compute_features needs both target rows to derive target-error history.
        if target_rows is not None:
            dim_pair = compute_dimensionless(sensors_raw, forces_raw, u0=self.u0, d=self.d, rho=self.rho)
            pair = compute_features(
                dim_pair, np.vstack([a2, a1]), np.vstack([np.zeros(3), a2]), self.mu,
                alpha_mode=True, include_mu=self.include_mu,
                include_cos_sin=self.include_cos_sin, dt_c=self.dt_c, u0=self.u0,
                target_forces=target_rows, rho=self.rho,
                sensors_raw=sensors_raw, forces_raw=forces_raw,
            )
            for key in ("Cd_err_lag1", "Cl_err_lag1", "dCd_err_dt", "dCl_err_dt"):
                feature_arrays[key] = pair[key][-1:]

        scalar_features = {key: np.asarray(value)[-1] for key, value in feature_arrays.items()}
        vector = None
        if feat_keys is not None:
            vector = build_feature_matrix(feature_arrays, list(feat_keys), add_bias=add_bias, strict=strict)[-1]

        self.state.previous_sensors = s.copy()
        self.state.previous_forces = f.copy()
        self.state.previous_target_force = None if tf is None else tf.copy()
        self.state.step_index += 1
        return scalar_features, vector

    def observe_action(self, action: np.ndarray) -> None:
        """Advance stored action history after an action is selected/applied."""
        alpha = self._alpha(action)
        self.state.previous_action2 = None if self.state.previous_action is None else self.state.previous_action.copy()
        self.state.previous_action = alpha

    def reset(self) -> None:
        self.state.reset()
