"""Sequence creation and scaling utilities for recurrent neural networks (LSTM, GRU, etc.).

This module provides classes for transforming time-series data into
windowed sequences and scaling 3-D tensors — the two most common
preprocessing steps before training sequence-based models.

Example:
    >>> from wael_lib import SequenceCreator, SequenceScaler
    >>> seq = SequenceCreator(window=30, horizon=1)
    >>> X_seq, y_seq = seq.create_sequences(X_features, y_target)
    >>> scaler = SequenceScaler()
    >>> X_scaled = scaler.fit_transform(X_seq)
"""

from __future__ import annotations

from typing import Tuple, Union

import numpy as np
import pandas as pd
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.discriminant_analysis import StandardScaler
from sklearn.preprocessing import MinMaxScaler


class SequenceCreator:
    """Create sliding-window sequences from time-series data for RNNs.

    Transforms flat time-series arrays into ``(samples, timesteps, features)``
    tensors suitable for Keras / PyTorch LSTM, GRU, and Transformer models.

    Attributes:
        window (int): Number of past steps used as input per sequence.
        horizon (int): Number of steps ahead to predict.

    Example:
        >>> import numpy as np
        >>> seq = SequenceCreator(window=5, horizon=1)
        >>> X = np.arange(10).reshape(-1, 1)
        >>> y = np.arange(10) * 10
        >>> X_seq, y_seq = seq.create_sequences(X, y)
        >>> X_seq.shape   # (5, 5, 1)
        >>> y_seq.shape   # (5,)
    """

    def __init__(self, window: int = 10, horizon: int = 1) -> None:
        """Initialise the SequenceCreator.

        Args:
            window (int): Number of past time steps used as input for each
                sequence.  Must be a positive integer.  Defaults to ``10``.
            horizon (int): Number of steps ahead to predict.  Must be a
                positive integer.  Defaults to ``1``.

        Raises:
            ValueError: If *window* or *horizon* are not positive integers.

        Example:
            >>> seq = SequenceCreator(window=30, horizon=1)
        """
        if window <= 0 or horizon <= 0:
            raise ValueError("Both window and horizon must be positive integers.")
        self.window: int = window
        self.horizon: int = horizon

    def _to_np(self, arr: Union[np.ndarray, pd.DataFrame, pd.Series]) -> np.ndarray:
        """Convert a pandas or array-like object to a NumPy array.

        Args:
            arr (Union[np.ndarray, pd.DataFrame, pd.Series]): Input data.

        Returns:
            np.ndarray: The data as a NumPy array.
        """
        return arr.values if hasattr(arr, "values") else np.asarray(arr)

    def create_sequences(
        self,
        X: Union[np.ndarray, pd.DataFrame, pd.Series],
        y: Union[np.ndarray, pd.Series],
    ) -> Tuple[np.ndarray, np.ndarray]:
        """Create input/target sequence pairs using a sliding window.

        For each position *i* in ``[0, len(X) - window)``, the method
        creates:

        * ``X_seq[i] = X[i : i + window]``
        * ``y_seq[i] = y[i + window]``

        If *X* is 1-D, an extra feature dimension is added automatically
        so the output is always 3-D.

        Args:
            X (Union[np.ndarray, pd.DataFrame, pd.Series]): Input features.
                Shape ``(n_samples,)`` for univariate or
                ``(n_samples, n_features)`` for multivariate data.
            y (Union[np.ndarray, pd.Series]): Target values.
                Shape ``(n_samples,)``.

        Returns:
            Tuple[np.ndarray, np.ndarray]: ``(X_seq, y_seq)`` where:

            * **X_seq** – shape ``(n_sequences, window, n_features)``.
            * **y_seq** – shape ``(n_sequences,)``.

        Example:
            >>> X_seq, y_seq = seq.create_sequences(df[features], df["target"])
        """
        X = self._to_np(X)
        y = self._to_np(y)

        X_seq, y_seq = [], []
        for i in range(len(X) - self.window):
            X_seq.append(X[i : i + self.window])
            y_seq.append(y[i + self.window])

        X_seq = np.array(X_seq)
        y_seq = np.array(y_seq)

        if X_seq.ndim == 2:
            X_seq = X_seq[..., np.newaxis]

        return X_seq, y_seq

    def create_X_sequences(
        self,
        X: Union[np.ndarray, pd.DataFrame, pd.Series],
    ) -> np.ndarray:
        """Create input sequences for prediction (no target values needed).

        Unlike :meth:`create_sequences`, this method does **not** require
        target values, making it suitable for inference on new/unseen data.

        Args:
            X (Union[np.ndarray, pd.DataFrame, pd.Series]): Input features.
                Shape ``(n_samples,)`` for univariate or
                ``(n_samples, n_features)`` for multivariate data.

        Returns:
            np.ndarray: Sequences of shape
            ``(n_sequences, window, n_features)`` where
            ``n_sequences = len(X) - window``.

        Example:
            >>> X_seq = seq.create_X_sequences(new_data[features])
            >>> predictions = model.predict(X_seq)
        """
        X = self._to_np(X)

        X_seq = []
        for i in range(len(X) - self.window):
            X_seq.append(X[i : i + self.window])

        X_seq = np.array(X_seq)

        if X_seq.ndim == 2:
            X_seq = X_seq[..., np.newaxis]

        return X_seq


class SequenceScaler(BaseEstimator, TransformerMixin):
    """Scikit-learn compatible MinMaxScaler for 3-D sequence tensors.

    Reshapes the 3-D input ``(batch, timesteps, features)`` to 2-D, applies
    ``MinMaxScaler``, and reshapes back.  Drop-in replacement for
    ``MinMaxScaler`` inside an ``sklearn.pipeline.Pipeline``.

    Attributes:
        scaler (MinMaxScaler): Underlying scaler instance.

    Example:
        >>> from sklearn.pipeline import Pipeline
        >>> pipe = Pipeline([
        ...     ("scaler", SequenceScaler()),
        ...     ("model", KerasRegressor(model=build_lstm)),
        ... ])
        >>> pipe.fit(X_train_3d, y_train)
    """

    def __init__(self) -> None:
        """Initialise the SequenceScaler with a default MinMaxScaler.

        Example:
            >>> scaler = SequenceScaler()
        """
        self.scaler: MinMaxScaler = MinMaxScaler()

    def fit(
        self,
        X: np.ndarray,
        y: Union[np.ndarray, None] = None,
    ) -> SequenceScaler:
        """Learn min/max statistics from 3-D sequence data.

        Args:
            X (np.ndarray): Training data of shape
                ``(batch_size, timesteps, n_features)``.
            y (Union[np.ndarray, None]): Ignored.  Present for sklearn
                Pipeline compatibility.

        Returns:
            SequenceScaler: ``self`` (fitted instance, for method chaining).

        Raises:
            ValueError: If *X* is not 3-dimensional.

        Example:
            >>> scaler = SequenceScaler()
            >>> scaler.fit(X_train)
        """
        batch_size, timesteps, n_features = X.shape
        X_reshaped: np.ndarray = X.reshape(-1, n_features)
        self.scaler.fit(X_reshaped)
        return self

    def transform(self, X: np.ndarray) -> np.ndarray:
        """Scale 3-D sequence data to [0, 1] using fitted parameters.

        Args:
            X (np.ndarray): Data of shape
                ``(batch_size, timesteps, n_features)``.
                *n_features* must match the data used during :meth:`fit`.

        Returns:
            np.ndarray: Scaled data with the **same shape** as input.
            Each feature is independently scaled to ``[0, 1]``.

        Raises:
            sklearn.exceptions.NotFittedError: If called before :meth:`fit`.

        Example:
            >>> X_scaled = scaler.transform(X_test)
        """
        batch_size, timesteps, n_features = X.shape
        X_reshaped: np.ndarray = X.reshape(-1, n_features)
        X_scaled: np.ndarray = self.scaler.transform(X_reshaped)
        return X_scaled.reshape(batch_size, timesteps, n_features)


class MySequenceScaler(BaseEstimator, TransformerMixin):
    """Scikit-learn compatible StandardScaler for 3-D sequence tensors.

    Identical API to :class:`SequenceScaler` but uses ``StandardScaler``
    (zero-mean, unit-variance) instead of ``MinMaxScaler``.

    Attributes:
        scaler (StandardScaler): Underlying scaler instance.

    Example:
        >>> scaler = MySequenceScaler()
        >>> X_scaled = scaler.fit_transform(X_train_3d)
    """

    def __init__(self) -> None:
        """Initialise the MySequenceScaler with a default StandardScaler.

        Example:
            >>> scaler = MySequenceScaler()
        """
        self.scaler: StandardScaler = StandardScaler()

    def fit(
        self,
        X: np.ndarray,
        y: Union[np.ndarray, None] = None,
    ) -> MySequenceScaler:
        """Learn mean/std statistics from 3-D sequence data.

        Args:
            X (np.ndarray): Training data of shape
                ``(batch_size, timesteps, n_features)``.
            y (Union[np.ndarray, None]): Ignored.  Present for sklearn
                Pipeline compatibility.

        Returns:
            MySequenceScaler: ``self`` (fitted instance).

        Raises:
            ValueError: If *X* is not 3-dimensional.

        Example:
            >>> scaler = MySequenceScaler()
            >>> scaler.fit(X_train)
        """
        n_samples, n_timesteps, n_features = X.shape
        X_reshaped = X.reshape(-1, n_features)
        self.scaler.fit(X_reshaped)
        return self

    def transform(self, X: np.ndarray) -> np.ndarray:
        """Standardise 3-D sequence data (zero-mean, unit-variance).

        Args:
            X (np.ndarray): Data of shape
                ``(batch_size, timesteps, n_features)``.

        Returns:
            np.ndarray: Standardised data with the **same shape** as input.

        Raises:
            sklearn.exceptions.NotFittedError: If called before :meth:`fit`.

        Example:
            >>> X_scaled = scaler.transform(X_test)
        """
        n_samples, n_timesteps, n_features = X.shape
        X_reshaped = X.reshape(-1, n_features)
        X_scaled = self.scaler.transform(X_reshaped)
        return X_scaled.reshape(n_samples, n_timesteps, n_features)
