__all__ = ['add_distributed_noise', 'prepare_data']

import logging
import pandas as pd
import numpy as np

from sklearn import preprocessing as skp
from sklearn import pipeline as skpl

_log = logging.getLogger(__name__)


def add_distributed_noise(ts, noise_std=0.05, noise_mean=0.0, rng=None, random_seed=None):
    """Add white noise to a time series 1D array.

    Noise amplitude is scaled by the signal's own standard deviation, so
    noise_std is relative to the signal: 0.05 means roughly 5% of one std.

    Args:
        ts (1D numpy array): time series data.
        noise_std (float): noise amplitude relative to signal std. Default 0.05.
        noise_mean (float): noise mean (usually 0). Default 0.0.
        rng (numpy.random.Generator, optional): random generator. Created if None.
        random_seed (int, optional): seed used when rng is None.

    Returns:
        1D numpy array: same shape as ts with noise added.
    """
    if rng is None:
        rng = np.random.default_rng(seed=random_seed)

    ts_std = ts.std()
    noise = ts_std * (noise_std * rng.standard_normal(len(ts)) + noise_mean)
    return ts + noise


def prepare_data(data, lags_in, cols_in, steps_in, cols_out, steps_out,
                 augmentation_factor=0, noise_std=0.05, noise_mean=0.0,
                 scaler_in=None, scaler_out=None, verbose=1, random_seed: int = 0):
    """Prepare time series data as supervised learning windows.

    Args:
        data (DataFrame): input DataFrame with a contiguous datetime or integer index.
        lags_in (list or None): extra lag offsets to include as input features.
            If None, only the current timestep (lag 0) is used.
        cols_in (list): column names to use as input features.
        steps_in (int): number of input timesteps per window.
        cols_out (list): column names to predict.
        steps_out (int): number of output timesteps per window.
        augmentation_factor (int): number of additional noised copies to append.
            0 means no augmentation. Each copy uses a single coherent noise
            realization applied consistently across all features and timesteps,
            ensuring x and y remain aligned.
        noise_std (float): noise amplitude relative to each feature's std. Default 0.05.
        noise_mean (float): noise mean. Default 0.0.
        scaler_in (sklearn transformer, optional): pre-fit scaler for input columns.
            If None a StandardScaler+MinMaxScaler pipeline is fit on data.
            Pass a scaler fit only on training data to avoid leakage when data
            contains test-period rows.
        scaler_out (sklearn transformer, optional): same as scaler_in for outputs.
        verbose (int): verbosity. Default 1.
        random_seed (int): seed for the noise RNG used during augmentation. Default 0.

    Raises:
        ValueError: if data does not contain enough rows for one sample.

    Returns:
        X (ndarray, shape (n_total, steps_in, n_in_features)): input windows.
            n_total = samples * (augmentation_factor + 1).
        y (ndarray, shape (n_total, steps_out, n_out_features)): output windows.
        index (list): the original data index values.
        scaler_in: fitted input scaler (may be the one passed in).
        scaler_out: fitted output scaler (may be the one passed in).
        samples (int): number of original (non-augmented) windows.
    """
    if augmentation_factor is not None and augmentation_factor < 0:
        raise ValueError(
            'augmentation_factor must be >= 0, got {0}.'.format(
                augmentation_factor))
    if augmentation_factor is not None and augmentation_factor > 10_000:
        raise ValueError(
            'augmentation_factor={0} is suspiciously large (> 10,000). '
            'Pass a smaller value to avoid memory exhaustion.'.format(
                augmentation_factor))

    lags = sorted(set([0] + (lags_in or [])))
    steps_lag = max(lags)

    data_in = data[cols_in].values
    data_out = data[cols_out].values

    if verbose > 0:
        _log.info('data_in.shape  = %s', data_in.shape)
        _log.info('data_out.shape = %s', data_out.shape)

    samples = data_in.shape[0] - (steps_in + steps_out + steps_lag) + 1
    if samples < 1:
        raise ValueError('Not enough data to produce at least 1 sample.')

    if verbose > 0:
        _log.info('samples = %d', samples)

    index = list(data.index)

    # Fit or reuse scalers. Pass pre-fit scalers to avoid leakage when data
    # contains rows from the test period.
    if scaler_in is None:
        scaler_in = skpl.Pipeline([
            ('std', skp.StandardScaler()),
            ('minmax', skp.MinMaxScaler(feature_range=(0, 1)))])
        data_in_scaled = scaler_in.fit_transform(data_in)
    else:
        data_in_scaled = scaler_in.transform(data_in)

    if scaler_out is None:
        scaler_out = skpl.Pipeline([
            ('std', skp.StandardScaler()),
            ('minmax', skp.MinMaxScaler(feature_range=(0, 1)))])
        data_out_scaled = scaler_out.fit_transform(data_out)
    else:
        data_out_scaled = scaler_out.transform(data_out)

    n_in_features = len(cols_in) * len(lags)
    n_out_features = len(cols_out)
    n_augmented = augmentation_factor if augmentation_factor is not None else 0
    n_total = samples * (n_augmented + 1)

    X = np.empty((n_total, steps_in, n_in_features), dtype=np.float64)
    y = np.empty((n_total, steps_out, n_out_features), dtype=np.float64)

    def _fill_windows(slot_offset, d_in, d_out):
        """Extract sliding windows from d_in / d_out into X and y."""
        for sample in range(samples):
            slot = slot_offset + sample
            feat_idx = 0
            for col in range(len(cols_in)):
                for lag in lags:
                    for step in range(steps_in):
                        X[slot, step, feat_idx] = d_in[sample + steps_lag + step - lag, col]
                    feat_idx += 1
            for col in range(len(cols_out)):
                for step in range(steps_out):
                    y[slot, step, col] = d_out[sample + steps_lag + steps_in + step, col]

    # Original (un-noised) windows
    _fill_windows(0, data_in_scaled, data_out_scaled)

    # Augmented copies: one coherent noise realization per copy. The same rng
    # is reused across columns so that x and y noise are correlated within a
    # copy, preventing the input and target from drifting in opposite directions.
    rng = np.random.default_rng(seed=random_seed)
    for aug in range(n_augmented):
        data_in_noised = np.column_stack([
            add_distributed_noise(
                data_in_scaled[:, c],
                noise_std=noise_std,
                noise_mean=noise_mean,
                rng=rng)
            for c in range(data_in_scaled.shape[1])
        ])
        data_out_noised = np.column_stack([
            add_distributed_noise(
                data_out_scaled[:, c],
                noise_std=noise_std,
                noise_mean=noise_mean,
                rng=rng)
            for c in range(data_out_scaled.shape[1])
        ])
        _fill_windows(samples * (aug + 1), data_in_noised, data_out_noised)

    if verbose > 0:
        _log.info('X.shape = %s', X.shape)
        _log.info('y.shape = %s', y.shape)

    return X, y, index, scaler_in, scaler_out, samples
