"""Signal processing and generation module for jbqlab.

This module provides utilities for:
- Signal smoothing and filtering
- Signal combination and ensemble
- Entry/exit timing refinement
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np
import pandas as pd

if TYPE_CHECKING:
    from collections.abc import Sequence

__all__ = [
    "smooth_signal",
    "ensemble_signals",
    "signal_lag",
    "signal_filter",
    "signal_confirm",
    "signal_to_trades",
    "calculate_signal_quality",
]


def smooth_signal(
    signal: pd.Series,
    method: str = "ema",
    window: int = 3,
    threshold: float = 0.5,
) -> pd.Series:
    """Smooth a binary signal to reduce whipsaws.

    Args:
        signal: Binary signal series (0 or 1).
        method: Smoothing method ('ema', 'sma', 'majority').
        window: Smoothing window size.
        threshold: Threshold for converting smoothed to binary.

    Returns:
        Smoothed binary signal.
    """
    if method == "ema":
        smoothed = signal.ewm(span=window, adjust=False).mean()
    elif method == "sma":
        smoothed = signal.rolling(window=window, min_periods=1).mean()
    elif method == "majority":
        # Take the majority vote over the window
        smoothed = signal.rolling(window=window, min_periods=1).apply(
            lambda x: 1 if x.mean() >= threshold else 0, raw=True
        )
        return smoothed.fillna(0).astype(int)
    else:
        raise ValueError(f"Invalid method: {method}")

    # Convert back to binary
    return (smoothed >= threshold).astype(int)


def ensemble_signals(
    signals: Sequence[pd.Series],
    method: str = "majority",
    weights: Sequence[float] | None = None,
    threshold: float = 0.5,
) -> pd.Series:
    """Combine multiple signals into an ensemble signal.

    Args:
        signals: List of signal series.
        method: Combination method ('majority', 'unanimous', 'weighted', 'any').
        weights: Weights for weighted method (must sum to 1).
        threshold: Threshold for majority/weighted methods.

    Returns:
        Combined ensemble signal.
    """
    if len(signals) == 0:
        raise ValueError("Must provide at least one signal")

    if len(signals) == 1:
        return signals[0].copy()

    # Align all signals
    df = pd.concat(signals, axis=1)
    df.columns = range(len(signals))
    df = df.fillna(0)

    if method == "majority":
        # At least threshold% must agree
        mean_signal = df.mean(axis=1)
        return (mean_signal >= threshold).astype(int)

    elif method == "unanimous":
        # All signals must agree (all 1s)
        return df.min(axis=1).astype(int)

    elif method == "any":
        # Any signal is 1
        return df.max(axis=1).astype(int)

    elif method == "weighted":
        if weights is None:
            weights = [1 / len(signals)] * len(signals)
        if len(weights) != len(signals):
            raise ValueError("Weights must match number of signals")

        weighted = sum(df.iloc[:, i] * w for i, w in enumerate(weights))
        return (weighted >= threshold).astype(int)

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


def signal_lag(
    signal: pd.Series,
    lag: int = 1,
) -> pd.Series:
    """Apply lag to a signal (delayed execution).

    Args:
        signal: Signal series.
        lag: Number of periods to lag (positive = delay).

    Returns:
        Lagged signal.
    """
    if lag == 0:
        return signal.copy()

    lagged = signal.shift(lag)
    return lagged.fillna(0).astype(int)


def signal_filter(
    signal: pd.Series,
    filter_condition: pd.Series,
    mode: str = "allow",
) -> pd.Series:
    """Filter a signal based on a condition.

    Args:
        signal: Signal series to filter.
        filter_condition: Boolean series (True = condition met).
        mode: 'allow' (only signal when condition True) or
              'block' (block signal when condition True).

    Returns:
        Filtered signal.
    """
    # Align series
    combined = pd.concat([signal, filter_condition], axis=1).fillna(False)
    sig = combined.iloc[:, 0]
    cond = combined.iloc[:, 1].astype(bool)

    if mode == "allow":
        # Only allow signal when condition is True
        return (sig & cond).astype(int)
    elif mode == "block":
        # Block signal when condition is True
        return (sig & ~cond).astype(int)
    else:
        raise ValueError(f"Invalid mode: {mode}")


def signal_confirm(
    signal: pd.Series,
    confirm_periods: int = 2,
) -> pd.Series:
    """Require signal to persist for n periods before confirming.

    Args:
        signal: Signal series.
        confirm_periods: Number of consecutive periods required.

    Returns:
        Confirmed signal (only 1 after n consecutive 1s).
    """
    if confirm_periods <= 1:
        return signal.copy()

    # Calculate consecutive 1s
    groups = (signal != signal.shift()).cumsum()
    consecutive = signal.groupby(groups).cumcount() + 1

    # Only confirm after n periods
    confirmed = (signal == 1) & (consecutive >= confirm_periods)
    return confirmed.astype(int)


def signal_to_trades(
    signal: pd.Series,
) -> pd.DataFrame:
    """Convert a signal series to a trade list.

    Args:
        signal: Signal series (1 = long, 0 = flat).

    Returns:
        DataFrame with entry_date, exit_date, duration columns.
    """
    trades = []
    in_trade = False
    entry_date = None

    for date, value in signal.items():
        if value == 1 and not in_trade:
            # Entry
            in_trade = True
            entry_date = date
        elif value == 0 and in_trade:
            # Exit
            in_trade = False
            trades.append(
                {
                    "entry_date": entry_date,
                    "exit_date": date,
                }
            )
            entry_date = None

    # Handle open trade at end
    if in_trade and entry_date is not None:
        trades.append(
            {
                "entry_date": entry_date,
                "exit_date": signal.index[-1],
            }
        )

    df = pd.DataFrame(trades)

    if len(df) > 0:
        df["duration"] = (df["exit_date"] - df["entry_date"]).dt.days

    return df


def calculate_signal_quality(
    signal: pd.Series,
    returns: pd.Series,
) -> dict:
    """Calculate quality metrics for a signal.

    Args:
        signal: Signal series.
        returns: Return series (aligned with signal).

    Returns:
        Dictionary with quality metrics.
    """
    # Align series
    combined = pd.concat([signal, returns], axis=1).dropna()
    if len(combined) < 10:
        return {"error": "Insufficient data"}

    sig = combined.iloc[:, 0]
    ret = combined.iloc[:, 1]

    # Strategy returns (signal * next day return)
    strategy_returns = sig.shift(1) * ret
    strategy_returns = strategy_returns.dropna()

    # Count trades
    changes = sig.diff().abs()
    n_entries = int((changes == 1).sum() / 2)

    # Time in market
    time_in_market = sig.mean()

    # Win rate when in position
    in_position_returns = ret[sig.shift(1) == 1].dropna()
    if len(in_position_returns) > 0:
        win_rate = (in_position_returns > 0).mean()
        avg_win = in_position_returns[in_position_returns > 0].mean() if (in_position_returns > 0).any() else 0
        avg_loss = in_position_returns[in_position_returns < 0].mean() if (in_position_returns < 0).any() else 0
    else:
        win_rate = 0
        avg_win = 0
        avg_loss = 0

    # Signal turnover
    turnover = changes.mean()

    # Information coefficient (correlation of signal with forward return)
    ic = sig.corr(ret.shift(-1))

    return {
        "n_entries": n_entries,
        "time_in_market": float(time_in_market),
        "win_rate": float(win_rate),
        "avg_win": float(avg_win) if not np.isnan(avg_win) else 0,
        "avg_loss": float(avg_loss) if not np.isnan(avg_loss) else 0,
        "turnover": float(turnover),
        "information_coefficient": float(ic) if not np.isnan(ic) else 0,
        "total_return": float((1 + strategy_returns).prod() - 1),
        "sharpe_approx": float(strategy_returns.mean() / strategy_returns.std() * np.sqrt(252)) if strategy_returns.std() > 0 else 0,
    }


def generate_crossover_signal(
    series: pd.Series,
    fast_window: int,
    slow_window: int,
) -> pd.Series:
    """Generate crossover signal from any series.

    Args:
        series: Price or indicator series.
        fast_window: Fast MA window.
        slow_window: Slow MA window.

    Returns:
        Signal series (1 when fast > slow).
    """
    fast_ma = series.rolling(window=fast_window, min_periods=fast_window).mean()
    slow_ma = series.rolling(window=slow_window, min_periods=slow_window).mean()

    signal = (fast_ma > slow_ma).astype(int)
    return signal.fillna(0).astype(int)


def generate_threshold_signal(
    series: pd.Series,
    upper_threshold: float,
    lower_threshold: float,
    mode: str = "breakout",
) -> pd.Series:
    """Generate signal based on threshold crossings.

    Args:
        series: Indicator series.
        upper_threshold: Upper threshold value.
        lower_threshold: Lower threshold value.
        mode: 'breakout' (long above upper) or 'mean_reversion' (long below lower).

    Returns:
        Signal series.
    """
    if mode == "breakout":
        # Long when above upper, flat when below lower
        signal = pd.Series(0, index=series.index)
        position = 0
        for i in range(len(series)):
            if position == 0 and series.iloc[i] > upper_threshold:
                position = 1
            elif position == 1 and series.iloc[i] < lower_threshold:
                position = 0
            signal.iloc[i] = position
        return signal

    elif mode == "mean_reversion":
        # Long when below lower, flat when above upper
        signal = pd.Series(0, index=series.index)
        position = 0
        for i in range(len(series)):
            if position == 0 and series.iloc[i] < lower_threshold:
                position = 1
            elif position == 1 and series.iloc[i] > upper_threshold:
                position = 0
            signal.iloc[i] = position
        return signal

    else:
        raise ValueError(f"Invalid mode: {mode}")
