"""
Parameter tuning for Kalman Pairs Trader.

Grid search over transition_covariance (and optionally observation_covariance)
to find optimal settings that maximize out-of-sample Sharpe while controlling turnover.

Design:
  - Train/validation split (chronological, no look-ahead)
  - Grid search over log-spaced transition_covariance values
  - Objective: maximize Sharpe_validation - λ * max_drawdown_penalty
  - Reports top-N parameter sets with full metrics
"""

from __future__ import annotations

import numpy as np
from typing import Optional
from .kalman_filter import KalmanPairsTrader
from .backtest import backtest_kalman_pairs


def grid_search_transition_cov(
    X_train: np.ndarray,
    Y_train: np.ndarray,
    X_val: np.ndarray,
    Y_val: np.ndarray,
    transition_cov_range: tuple[float, float, int] = (1e-6, 1e-1, 20),
    observation_covariance: float = 1e-2,
    z_entry: float = 2.0,
    z_exit: float = 0.5,
    max_drawdown_penalty: float = 0.5,
    trade_size_usd: float = 100.0,
    transaction_cost_bps: float = 2.5,
) -> list[dict]:
    """
    Grid search optimal transition_covariance.

    Strategy:
        1. Split data chronologically (train → validation).
        2. For each Q value, run Kalman backtest on validation set
           (with no pre-training — Kalman adapts online).
        3. Score = Sharpe − λ * max_drawdown.
        4. Return sorted results.

    Args:
        X_train, Y_train: Training price series (used for initialization only).
        X_val, Y_val: Validation price series (out-of-sample test).
        transition_cov_range: (min, max, num_steps) in log space.
        max_drawdown_penalty: Weight for drawdown penalty in scoring.

    Returns:
        List of dicts sorted by score (descending), each with:
            transition_cov, sharpe, sortino, max_drawdown, win_rate, total_trades, score.
    """
    q_min, q_max, n_steps = transition_cov_range
    q_values = np.logspace(np.log10(q_min), np.log10(q_max), n_steps)

    results = []
    for q in q_values:
        trader = KalmanPairsTrader(
            transition_covariance=float(q),
            observation_covariance=observation_covariance,
            z_entry=z_entry,
            z_exit=z_exit,
        )

        # Pre-warm on training data (online filtering, no position taking)
        for x, y in zip(X_train, Y_train):
            trader.kf.update(float(x), float(y))

        # Backtest on validation
        bt = backtest_kalman_pairs(
            X_val, Y_val, trader,
            trade_size_usd=trade_size_usd,
            transaction_cost_bps=transaction_cost_bps,
        )

        score = bt["sharpe"] - max_drawdown_penalty * bt["max_drawdown"]

        results.append({
            "transition_cov": float(q),
            "sharpe": bt["sharpe"],
            "sortino": bt["sortino"],
            "max_drawdown": bt["max_drawdown"],
            "win_rate": bt["win_rate"],
            "total_trades": bt["total_trades"],
            "pnl_pct": bt["pnl_pct"],
            "score": round(score, 4),
        })

    results.sort(key=lambda r: r["score"], reverse=True)
    return results


def find_optimal_params(
    X: np.ndarray,
    Y: np.ndarray,
    train_frac: float = 0.6,
    **grid_kwargs,
) -> dict:
    """
    One-shot: split data, run grid search, return best params.

    Returns:
        dict with: best_params, all_results, train_size, val_size.
    """
    n = len(X)
    split = int(n * train_frac)
    X_train, X_val = X[:split], X[split:]
    Y_train, Y_val = Y[:split], Y[split:]

    grid = grid_search_transition_cov(X_train, Y_train, X_val, Y_val, **grid_kwargs)

    return {
        "best_params": {
            "transition_covariance": grid[0]["transition_cov"] if grid else 1e-4,
        },
        "best_score": grid[0]["score"] if grid else 0.0,
        "all_results": grid,
        "train_size": len(X_train),
        "val_size": len(X_val),
    }