"""Elastic Net / LASSO signal selection pipeline functions.

Pipeline functions:
    signal_selector - Automatic signal selection via regularized regression
"""

from __future__ import annotations

from typing import Any, Callable

from horizon._horizon import (
    elastic_net_fit,
    elastic_net_predict,
)
from horizon.context import Context


def signal_selector(
    signals: list[str],
    target: str = "return",
    window: int = 100,
    refit_interval: int = 50,
    alpha: float = 0.1,
    l1_ratio: float = 0.7,
) -> Callable[[Context], dict[str, Any]]:
    """Create a signal selection pipeline using elastic net regression.

    Automatically selects the most predictive signals from a set of
    candidates using L1/L2 regularization.

    Args:
        signals: List of signal parameter names to use as features.
        target: Target parameter name (default: "return").
        window: Rolling window for fitting.
        refit_interval: How often to refit the model (in ticks).
        alpha: Regularization strength.
        l1_ratio: L1 vs L2 mix (1.0 = LASSO, 0.0 = Ridge).

    Returns:
        Pipeline function: (Context) -> dict with keys:
            prediction, selected_signals, coefficients, r_squared
    """
    from horizon._horizon import auth_require_pro
    auth_require_pro()

    histories: dict[str, list[list[float]]] = {}
    target_histories: dict[str, list[float]] = {}
    tick_counts: dict[str, int] = {}
    last_results: dict[str, Any] = {}

    def _signal_selector(ctx: Context) -> dict[str, Any]:
        market_id = ctx.market.id if ctx.market else "__default__"

        if market_id not in histories:
            histories[market_id] = []
            target_histories[market_id] = []
            tick_counts[market_id] = 0
            last_results[market_id] = None

        tick_counts[market_id] += 1

        # Collect current signal values
        row = []
        for sig_name in signals:
            val = ctx.params.get(sig_name, 0.0)
            row.append(float(val) if isinstance(val, (int, float)) else 0.0)

        histories[market_id].append(row)

        # Collect target
        target_val = ctx.params.get(target, 0.0)
        target_histories[market_id].append(
            float(target_val) if isinstance(target_val, (int, float)) else 0.0
        )

        # Trim to window
        if len(histories[market_id]) > window:
            histories[market_id] = histories[market_id][-window:]
            target_histories[market_id] = target_histories[market_id][-window:]

        result = {
            "prediction": 0.0,
            "selected_signals": [],
            "coefficients": {},
            "r_squared": 0.0,
        }

        # Refit periodically
        if (
            len(histories[market_id]) >= 10
            and tick_counts[market_id] % refit_interval == 0
        ):
            try:
                x = histories[market_id]
                y = target_histories[market_id]

                fit = elastic_net_fit(x, y, alpha, l1_ratio)
                last_results[market_id] = fit

                # Predict with current features
                preds = elastic_net_predict([row], fit)
                result["prediction"] = preds[0] if preds else 0.0
                result["r_squared"] = fit.r_squared

                # Selected signals (nonzero coefficients)
                selected = []
                coeffs = {}
                for i, coeff in enumerate(fit.coefficients):
                    if abs(coeff) > 1e-10 and i < len(signals):
                        selected.append(signals[i])
                        coeffs[signals[i]] = round(coeff, 6)

                result["selected_signals"] = selected
                result["coefficients"] = coeffs

            except (ValueError, RuntimeError):
                pass

        elif last_results[market_id] is not None:
            # Use cached model for prediction
            try:
                fit = last_results[market_id]
                preds = elastic_net_predict([row], fit)
                result["prediction"] = preds[0] if preds else 0.0
                result["r_squared"] = fit.r_squared

                selected = []
                coeffs = {}
                for i, coeff in enumerate(fit.coefficients):
                    if abs(coeff) > 1e-10 and i < len(signals):
                        selected.append(signals[i])
                        coeffs[signals[i]] = round(coeff, 6)

                result["selected_signals"] = selected
                result["coefficients"] = coeffs
            except (ValueError, RuntimeError):
                pass

        return result

    _signal_selector.__name__ = "signal_selector"
    return _signal_selector
