"""Preprocessing utilities for time series data."""

from typing import Literal, Optional, Union, Tuple

import numpy as np
import pandas as pd


def fill_missing_dates(
    series: pd.Series,
    frequency: Literal["D", "W", "M"] = "D",
    fill_method: Literal["linear", "seasonal", "forward", "zero"] = "linear",
    seasonal_period: Optional[int] = None,
) -> pd.Series:
    """
    Fill missing dates in a time series.

    Parameters
    ----------
    series : pd.Series
        Time series with DatetimeIndex. May have gaps.
    frequency : {'D', 'W', 'M'}, default 'D'
        Frequency of the time series.
    fill_method : {'linear', 'seasonal', 'forward', 'zero'}, default 'linear'
        Method for filling missing values:
        - 'linear': Linear interpolation
        - 'seasonal': Use same period from previous cycle
        - 'forward': Forward fill (last observation carried forward)
        - 'zero': Fill with zeros
    seasonal_period : int, optional
        Period for seasonal interpolation. Defaults based on frequency.

    Returns
    -------
    filled_series : pd.Series
        Complete time series with filled values.

    Examples
    --------
    >>> dates = pd.date_range('2024-01-01', periods=5, freq='D')
    >>> dates = dates.delete(2)  # Remove middle date
    >>> series = pd.Series([100, 110, 130, 140], index=dates)
    >>> filled = fill_missing_dates(series, frequency='D', fill_method='linear')
    """
    if not isinstance(series.index, pd.DatetimeIndex):
        raise ValueError("Series must have DatetimeIndex")

    # Generate complete date range
    freq_map = {"D": "D", "W": "W-MON", "M": "MS"}
    freq = freq_map.get(frequency, "D")

    start_date = series.index.min()
    end_date = series.index.max()

    complete_dates = pd.date_range(start=start_date, end=end_date, freq=freq)

    # Reindex to complete dates (creates NaN for missing)
    result = series.reindex(complete_dates)

    # Set default seasonal period
    if seasonal_period is None:
        seasonal_period = {"D": 7, "W": 52, "M": 12}.get(frequency, 7)

    # Fill missing values based on method
    if fill_method == "linear":
        result = _linear_interpolate(result)
    elif fill_method == "seasonal":
        result = _seasonal_interpolate(result, seasonal_period)
    elif fill_method == "forward":
        result = result.ffill().bfill()  # Forward fill, then backward fill edges
    elif fill_method == "zero":
        result = result.fillna(0)
    else:
        raise ValueError(f"Unknown fill_method: {fill_method}")

    return result


def _linear_interpolate(values: pd.Series) -> pd.Series:
    """
    Fill NaN values using linear interpolation.

    Parameters
    ----------
    values : pd.Series
        Series with potential NaN values.

    Returns
    -------
    interpolated : pd.Series
        Series with NaN values filled via linear interpolation.
    """
    result = values.copy()

    # Use pandas interpolation for interior values
    result = result.interpolate(method="linear")

    # Handle edges (extrapolate using nearest value)
    result = result.bfill().ffill()

    return result


def _seasonal_interpolate(values: pd.Series, period: int) -> pd.Series:
    """
    Fill NaN values using seasonal interpolation.

    Uses the value from the same period in the previous cycle.

    Parameters
    ----------
    values : pd.Series
        Series with potential NaN values.
    period : int
        Seasonal period.

    Returns
    -------
    interpolated : pd.Series
        Series with NaN values filled seasonally.
    """
    result = values.copy()
    n = len(result)

    for i in range(n):
        if pd.isna(result.iloc[i]):
            # Try looking back one period
            if i >= period and not pd.isna(result.iloc[i - period]):
                result.iloc[i] = result.iloc[i - period]
            # Try looking forward one period
            elif i + period < n and not pd.isna(result.iloc[i + period]):
                result.iloc[i] = result.iloc[i + period]

    # Fall back to linear interpolation for remaining NaNs
    if result.isna().any():
        result = _linear_interpolate(result)

    return result


def detect_outliers(
    values: Union[pd.Series, np.ndarray],
    method: Literal["iqr", "zscore", "mad"] = "iqr",
    threshold: float = 3.0,
) -> np.ndarray:
    """
    Detect outliers in a time series.

    Parameters
    ----------
    values : pd.Series or np.ndarray
        Values to check for outliers.
    method : {'iqr', 'zscore', 'mad'}, default 'iqr'
        Detection method:
        - 'iqr': Interquartile range (threshold is multiplier, typically 1.5 or 3)
        - 'zscore': Z-score (threshold is number of std devs, typically 3)
        - 'mad': Median Absolute Deviation (threshold is modified z-score, typically 3.5)
    threshold : float, default 3.0
        Threshold for outlier detection (interpretation depends on method).

    Returns
    -------
    is_outlier : np.ndarray
        Boolean array where True indicates an outlier.

    Examples
    --------
    >>> values = pd.Series([100, 110, 105, 1000, 108, 112])
    >>> outliers = detect_outliers(values, method='iqr', threshold=1.5)
    >>> print(outliers)
    [False False False  True False False]
    """
    values = np.asarray(values)
    n = len(values)
    is_outlier = np.zeros(n, dtype=bool)

    # Remove NaN for calculations
    valid_mask = ~np.isnan(values)
    valid_values = values[valid_mask]

    if len(valid_values) < 3:
        return is_outlier

    if method == "iqr":
        q1 = np.percentile(valid_values, 25)
        q3 = np.percentile(valid_values, 75)
        iqr = q3 - q1
        lower_bound = q1 - threshold * iqr
        upper_bound = q3 + threshold * iqr

        is_outlier = (values < lower_bound) | (values > upper_bound)

    elif method == "zscore":
        mean_val = np.nanmean(values)
        std_val = np.nanstd(values)

        if std_val > 0:
            z_scores = np.abs(values - mean_val) / std_val
            is_outlier = z_scores > threshold

    elif method == "mad":
        # Median Absolute Deviation - robust to outliers
        median_val = np.nanmedian(values)
        abs_deviations = np.abs(values - median_val)
        mad = np.nanmedian(abs_deviations)

        if mad > 0:
            # Modified z-score using MAD
            modified_z = 0.6745 * (values - median_val) / mad
            is_outlier = np.abs(modified_z) > threshold

    else:
        raise ValueError(f"Unknown method: {method}")

    # NaN values are not outliers
    is_outlier = is_outlier & valid_mask

    return is_outlier


def handle_outliers(
    series: pd.Series,
    method: Literal["iqr", "zscore", "mad"] = "iqr",
    threshold: float = 3.0,
    treatment: Literal["remove", "cap", "interpolate", "flag"] = "cap",
) -> Union[pd.Series, Tuple[pd.Series, pd.Series]]:
    """
    Detect and handle outliers in a time series.

    Parameters
    ----------
    series : pd.Series
        Time series to process.
    method : {'iqr', 'zscore', 'mad'}, default 'iqr'
        Detection method.
    threshold : float, default 3.0
        Threshold for outlier detection.
    treatment : {'remove', 'cap', 'interpolate', 'flag'}, default 'cap'
        How to handle detected outliers:
        - 'remove': Remove outliers (set to NaN)
        - 'cap': Cap outliers at threshold boundaries (winsorize)
        - 'interpolate': Replace outliers with interpolated values
        - 'flag': Return original series with outlier mask

    Returns
    -------
    result : pd.Series
        Processed series (or original if treatment='flag').
    outlier_mask : pd.Series
        Boolean mask of outliers (only if treatment='flag').

    Examples
    --------
    >>> series = pd.Series([100, 110, 105, 1000, 108, 112])
    >>> cleaned = handle_outliers(series, treatment='cap')
    """
    is_outlier = detect_outliers(series.values, method=method, threshold=threshold)
    outlier_mask = pd.Series(is_outlier, index=series.index)

    if treatment == "flag":
        return series.copy(), outlier_mask

    result = series.copy()

    if treatment == "remove":
        result[is_outlier] = np.nan

    elif treatment == "cap":
        # Calculate bounds for capping
        valid_values = series[~is_outlier].dropna()

        if method == "iqr":
            q1 = valid_values.quantile(0.25)
            q3 = valid_values.quantile(0.75)
            iqr = q3 - q1
            lower_bound = q1 - threshold * iqr
            upper_bound = q3 + threshold * iqr
        else:
            # For zscore and mad, use percentiles as bounds
            lower_bound = valid_values.quantile(0.01)
            upper_bound = valid_values.quantile(0.99)

        result = result.clip(lower=lower_bound, upper=upper_bound)

    elif treatment == "interpolate":
        result[is_outlier] = np.nan
        result = result.interpolate(method="linear")
        result = result.bfill().ffill()

    else:
        raise ValueError(f"Unknown treatment: {treatment}")

    return result


def log_transform(
    values: Union[pd.Series, np.ndarray],
    offset: float = 1.0,
) -> Union[pd.Series, np.ndarray]:
    """
    Apply log transformation to stabilize variance.

    Uses log(x + offset) to handle zeros.

    Parameters
    ----------
    values : pd.Series or np.ndarray
        Values to transform.
    offset : float, default 1.0
        Offset to add before taking log (handles zeros).

    Returns
    -------
    transformed : pd.Series or np.ndarray
        Log-transformed values.

    Notes
    -----
    For multiplicative decomposition, transform the series to make it
    additive: log(Y) = log(T) + log(S) + log(R)
    """
    is_series = isinstance(values, pd.Series)

    arr = np.asarray(values)

    # Check for negative values
    if (arr < -offset).any():
        raise ValueError(
            f"Values must be >= {-offset} for log transform with offset={offset}"
        )

    result = np.log(arr + offset)

    if is_series:
        return pd.Series(result, index=values.index, name=values.name)
    return result


def inverse_log_transform(
    values: Union[pd.Series, np.ndarray],
    offset: float = 1.0,
) -> Union[pd.Series, np.ndarray]:
    """
    Inverse log transformation.

    Reverses log_transform: exp(x) - offset

    Parameters
    ----------
    values : pd.Series or np.ndarray
        Log-transformed values.
    offset : float, default 1.0
        Offset that was used in log_transform.

    Returns
    -------
    original : pd.Series or np.ndarray
        Original-scale values.
    """
    is_series = isinstance(values, pd.Series)

    arr = np.asarray(values)
    result = np.exp(arr) - offset

    # Ensure non-negative (numerical precision issues)
    result = np.maximum(result, 0)

    if is_series:
        return pd.Series(result, index=values.index, name=values.name)
    return result


def difference(
    series: pd.Series,
    order: int = 1,
    seasonal_period: Optional[int] = None,
) -> pd.Series:
    """
    Apply differencing to achieve stationarity.

    Parameters
    ----------
    series : pd.Series
        Time series to difference.
    order : int, default 1
        Order of differencing (number of times to difference).
    seasonal_period : int, optional
        If provided, apply seasonal differencing (Y_t - Y_{t-m}).

    Returns
    -------
    differenced : pd.Series
        Differenced series.

    Notes
    -----
    Regular differencing: Y'_t = Y_t - Y_{t-1}
    Seasonal differencing: Y'_t = Y_t - Y_{t-m}
    """
    result = series.copy()

    # Apply seasonal differencing first
    if seasonal_period is not None:
        result = result - result.shift(seasonal_period)

    # Apply regular differencing
    for _ in range(order):
        result = result.diff()

    return result.dropna()


def inverse_difference(
    differenced: pd.Series,
    original: pd.Series,
    order: int = 1,
    seasonal_period: Optional[int] = None,
) -> pd.Series:
    """
    Reverse differencing operation.

    Parameters
    ----------
    differenced : pd.Series
        Differenced series (forecasts).
    original : pd.Series
        Original series (for getting initial values).
    order : int, default 1
        Order of differencing that was applied.
    seasonal_period : int, optional
        Seasonal period if seasonal differencing was applied.

    Returns
    -------
    restored : pd.Series
        Series restored to original scale.
    """
    result = differenced.copy()

    # Reverse regular differencing
    for _ in range(order):
        # Get the last value before the forecast period
        last_value = original.iloc[-1]
        result = result.cumsum() + last_value

    # Reverse seasonal differencing
    if seasonal_period is not None:
        # Need last m values from original
        for i in range(len(result)):
            if i < seasonal_period:
                result.iloc[i] = (
                    result.iloc[i] + original.iloc[-(seasonal_period - i)]
                )
            else:
                result.iloc[i] = result.iloc[i] + result.iloc[i - seasonal_period]

    return result


def normalize(
    series: pd.Series,
    method: Literal["minmax", "zscore"] = "minmax",
) -> Tuple[pd.Series, dict]:
    """
    Normalize a time series.

    Parameters
    ----------
    series : pd.Series
        Series to normalize.
    method : {'minmax', 'zscore'}, default 'minmax'
        Normalization method.

    Returns
    -------
    normalized : pd.Series
        Normalized series.
    params : dict
        Parameters needed for denormalization.
    """
    if method == "minmax":
        min_val = series.min()
        max_val = series.max()
        range_val = max_val - min_val

        if range_val == 0:
            normalized = pd.Series(0.5, index=series.index)
        else:
            normalized = (series - min_val) / range_val

        params = {"method": "minmax", "min": min_val, "max": max_val}

    elif method == "zscore":
        mean_val = series.mean()
        std_val = series.std()

        if std_val == 0:
            normalized = pd.Series(0, index=series.index)
        else:
            normalized = (series - mean_val) / std_val

        params = {"method": "zscore", "mean": mean_val, "std": std_val}

    else:
        raise ValueError(f"Unknown method: {method}")

    return normalized, params


def denormalize(series: pd.Series, params: dict) -> pd.Series:
    """
    Reverse normalization.

    Parameters
    ----------
    series : pd.Series
        Normalized series.
    params : dict
        Parameters from normalize().

    Returns
    -------
    original : pd.Series
        Series in original scale.
    """
    method = params["method"]

    if method == "minmax":
        range_val = params["max"] - params["min"]
        return series * range_val + params["min"]

    elif method == "zscore":
        return series * params["std"] + params["mean"]

    else:
        raise ValueError(f"Unknown method: {method}")
