"""
Parameter optimization module.

This module provides grid search functionality for finding optimal
strategy parameters through systematic backtesting.
"""

from dataclasses import dataclass
from itertools import product
from pathlib import Path
from typing import Any

import pandas as pd

from jbqlab.engine import run_backtest
from jbqlab.types import BacktestResult


@dataclass
class GridSearchResult:
    """Results from a parameter grid search.

    Attributes:
        results: List of (params, BacktestResult) tuples for each combination.
        best_params: Parameters that achieved the best metric value.
        best_result: BacktestResult for the best parameters.
        summary_df: DataFrame summarizing all runs.
        optimize_metric: The metric that was optimized.
    """

    results: list[tuple[dict[str, Any], BacktestResult]]
    best_params: dict[str, Any]
    best_result: BacktestResult
    summary_df: pd.DataFrame
    optimize_metric: str

    def top_n(self, n: int = 5) -> pd.DataFrame:
        """Get top N parameter combinations by optimized metric.

        Args:
            n: Number of top results to return.

        Returns:
            DataFrame with top N parameter combinations and their metrics.
        """
        return self.summary_df.head(n)


def grid_search(
    data: str | Path | pd.DataFrame,
    strategy: str,
    param_grid: dict[str, list[Any]],
    optimize_metric: str = "sharpe",
    higher_is_better: bool = True,
    transaction_cost: float = 0.001,
    slippage: float = 0.0005,
    initial_capital: float = 100_000.0,
    verbose: bool = True,
) -> GridSearchResult:
    """Run grid search over strategy parameters.

    Systematically tests all combinations of parameters and returns
    the combination that optimizes the specified metric.

    Args:
        data: Path to CSV file or DataFrame with price data.
        strategy: Name of the strategy to optimize.
        param_grid: Dictionary mapping parameter names to lists of values to try.
            Example: {"fast": [5, 10, 15], "slow": [20, 30, 40]}
        optimize_metric: Metric to optimize (default: "sharpe").
            Options: total_return, cagr, sharpe, sortino, calmar, volatility, max_drawdown
        higher_is_better: Whether higher metric values are better (default: True).
            Set to False for metrics like max_drawdown or volatility.
        transaction_cost: Transaction cost for all backtests.
        slippage: Slippage for all backtests.
        initial_capital: Initial capital for all backtests.
        verbose: Whether to print progress (default: True).

    Returns:
        GridSearchResult containing all results and the best parameters.

    Example:
        >>> result = grid_search(
        ...     "data.csv",
        ...     strategy="sma_crossover",
        ...     param_grid={"fast": [5, 10, 15], "slow": [20, 30, 40]},
        ...     optimize_metric="sharpe",
        ... )
        >>> print(result.best_params)
        {'fast': 10, 'slow': 30}
    """
    # Generate all parameter combinations
    param_names = list(param_grid.keys())
    param_values = list(param_grid.values())
    combinations = list(product(*param_values))

    total_combinations = len(combinations)

    if verbose:
        print(f"🔍 Grid search: {total_combinations} parameter combinations")
        print(f"   Strategy: {strategy}")
        print(f"   Optimizing: {optimize_metric} ({'higher' if higher_is_better else 'lower'} is better)")
        print()

    results: list[tuple[dict[str, Any], BacktestResult]] = []

    for i, combo in enumerate(combinations):
        params = dict(zip(param_names, combo, strict=False))

        if verbose:
            print(f"   [{i + 1}/{total_combinations}] Testing {params}...", end=" ")

        try:
            result = run_backtest(
                data=data,
                strategy=strategy,
                transaction_cost=transaction_cost,
                slippage=slippage,
                initial_capital=initial_capital,
                **params,
            )
            results.append((params, result))

            if verbose:
                metric_value = result.metrics.get(optimize_metric, 0)
                print(f"{optimize_metric}={metric_value:.4f}")

        except Exception as e:
            if verbose:
                print(f"FAILED: {e}")
            continue

    if not results:
        raise ValueError("All parameter combinations failed. Check your parameter grid.")

    # Build summary DataFrame
    summary_data = []
    for params, result in results:
        row = {**params, **result.metrics}
        summary_data.append(row)

    summary_df = pd.DataFrame(summary_data)

    # Sort by optimize metric
    summary_df = summary_df.sort_values(
        optimize_metric, ascending=not higher_is_better
    ).reset_index(drop=True)

    # Find best
    best_idx = 0
    best_params = {name: summary_df.iloc[best_idx][name] for name in param_names}
    best_result = next(r for p, r in results if p == best_params)

    if verbose:
        print()
        print(f"✅ Best parameters: {best_params}")
        print(f"   {optimize_metric}: {best_result.metrics[optimize_metric]:.4f}")

    return GridSearchResult(
        results=results,
        best_params=best_params,
        best_result=best_result,
        summary_df=summary_df,
        optimize_metric=optimize_metric,
    )


def walk_forward_optimization(
    data: str | Path | pd.DataFrame,
    strategy: str,
    param_grid: dict[str, list[Any]],
    train_period: int = 252,
    test_period: int = 63,
    optimize_metric: str = "sharpe",
    **kwargs: Any,
) -> pd.DataFrame:
    """Run walk-forward optimization.

    Splits data into rolling train/test windows, optimizes parameters
    on each training window, and tests on the subsequent test window.

    Args:
        data: Path to CSV file or DataFrame with price data.
        strategy: Name of the strategy.
        param_grid: Parameter grid for optimization.
        train_period: Number of days for training window (default: 252 = 1 year).
        test_period: Number of days for testing window (default: 63 = 1 quarter).
        optimize_metric: Metric to optimize during training.
        **kwargs: Additional arguments passed to run_backtest.

    Returns:
        DataFrame with results from each walk-forward window.
    """
    # Load data if path
    if isinstance(data, (str, Path)):
        from jbqlab.data import load_csv
        df = load_csv(data)
    else:
        df = data.copy()
        if "date" in df.columns:
            df = df.set_index("date")

    total_days = len(df)
    window_results = []
    window_num = 0

    start_idx = 0

    while start_idx + train_period + test_period <= total_days:
        window_num += 1

        # Split into train/test
        train_end = start_idx + train_period
        test_end = train_end + test_period

        train_data = df.iloc[start_idx:train_end].copy()
        test_data = df.iloc[train_end:test_end].copy()

        # Optimize on training data
        grid_result = grid_search(
            data=train_data,
            strategy=strategy,
            param_grid=param_grid,
            optimize_metric=optimize_metric,
            verbose=False,
            **kwargs,
        )

        # Test on test data with best params
        test_result = run_backtest(
            data=test_data,
            strategy=strategy,
            **grid_result.best_params,
            **kwargs,
        )

        window_results.append({
            "window": window_num,
            "train_start": df.index[start_idx],
            "train_end": df.index[train_end - 1],
            "test_start": df.index[train_end],
            "test_end": df.index[test_end - 1],
            "best_params": str(grid_result.best_params),
            "train_sharpe": grid_result.best_result.metrics.get("sharpe", 0),
            "test_sharpe": test_result.metrics.get("sharpe", 0),
            "test_return": test_result.metrics.get("total_return", 0),
        })

        # Move to next window
        start_idx += test_period

    return pd.DataFrame(window_results)
