"""Causal baseline predictors for the Wunder stepwise API."""

from __future__ import annotations

from typing import Optional

import numpy as np

from src.data.causal_features import build_causal_tabular_features


class PersistencePredictionModel:
    """Predict the current state as the next state."""

    def __init__(self, n_features: int = 32):
        self.n_features = n_features
        self.current_seq: Optional[int] = None
        self.last_state: Optional[np.ndarray] = None

    def predict(self, data_point):
        if self.current_seq != data_point.seq_ix:
            self.current_seq = data_point.seq_ix
            self.last_state = None
        self.last_state = data_point.state.astype(np.float32, copy=True)
        if not data_point.need_prediction:
            return None
        return self.last_state.astype(np.float32, copy=True)


class MomentumPredictionModel:
    """Predict next state with a first-difference momentum term."""

    def __init__(self, alpha: float = 1.0, n_features: int = 32):
        self.alpha = float(alpha)
        self.n_features = n_features
        self.current_seq: Optional[int] = None
        self.previous_state: Optional[np.ndarray] = None
        self.current_state: Optional[np.ndarray] = None

    def predict(self, data_point):
        if self.current_seq != data_point.seq_ix:
            self.current_seq = data_point.seq_ix
            self.previous_state = None
            self.current_state = None

        self.previous_state = self.current_state
        self.current_state = data_point.state.astype(np.float32, copy=True)
        if not data_point.need_prediction:
            return None
        if self.previous_state is None:
            return self.current_state.astype(np.float32, copy=True)
        pred = self.current_state + self.alpha * (self.current_state - self.previous_state)
        return pred.astype(np.float32)


class EWMAPredictionModel:
    """Predict with an exponentially weighted moving average state."""

    def __init__(self, alpha: float = 0.2, n_features: int = 32):
        if not 0.0 < alpha <= 1.0:
            raise ValueError("alpha must be in (0, 1]")
        self.alpha = float(alpha)
        self.n_features = n_features
        self.current_seq: Optional[int] = None
        self.ewma_state: Optional[np.ndarray] = None

    def predict(self, data_point):
        if self.current_seq != data_point.seq_ix:
            self.current_seq = data_point.seq_ix
            self.ewma_state = None
        state = data_point.state.astype(np.float32, copy=False)
        if self.ewma_state is None:
            self.ewma_state = state.copy()
        else:
            self.ewma_state = self.alpha * state + (1.0 - self.alpha) * self.ewma_state
        if not data_point.need_prediction:
            return None
        return self.ewma_state.astype(np.float32, copy=True)


class TabularStatefulPredictionModel:
    """Stateful wrapper for sklearn-style multi-output tabular regressors."""

    def __init__(self, estimator, n_features: int = 32, feature_schema: str = "compact_v1"):
        self.estimator = estimator
        self.n_features = n_features
        self.feature_schema = feature_schema
        self.current_seq: Optional[int] = None
        self.history: list[np.ndarray] = []

    def predict(self, data_point):
        if self.current_seq != data_point.seq_ix:
            self.current_seq = data_point.seq_ix
            self.history = []
        self.history.append(data_point.state.astype(np.float32, copy=True))
        if not data_point.need_prediction:
            return None
        features = build_causal_tabular_features(
            np.asarray(self.history, dtype=np.float32),
            schema=self.feature_schema,
        )
        pred = self.estimator.predict(features.reshape(1, -1))[0]
        return np.asarray(pred, dtype=np.float32)