"""Monte Carlo simulation module for jbqlab.

This module provides Monte Carlo methods for:
- Return simulation
- Drawdown analysis
- Confidence intervals
- Risk assessment
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

import numpy as np
import pandas as pd

if TYPE_CHECKING:
    from collections.abc import Callable

__all__ = [
    "MonteCarloResult",
    "monte_carlo_simulation",
    "bootstrap_returns",
    "simulate_paths",
    "confidence_intervals",
    "probability_of_loss",
    "expected_shortfall_mc",
]


@dataclass
class MonteCarloResult:
    """Results from Monte Carlo simulation."""

    n_simulations: int
    time_horizon: int
    final_values: np.ndarray
    all_paths: np.ndarray | None
    mean_final: float
    median_final: float
    std_final: float
    percentile_5: float
    percentile_25: float
    percentile_75: float
    percentile_95: float
    prob_loss: float
    max_drawdown_mean: float
    max_drawdown_95th: float

    def summary(self) -> str:
        """Generate text summary of results."""
        lines = [
            "Monte Carlo Simulation Results",
            "=" * 40,
            f"Simulations: {self.n_simulations:,}",
            f"Time Horizon: {self.time_horizon} periods",
            "",
            "Final Portfolio Value Distribution:",
            f"  Mean: {self.mean_final:.2%}",
            f"  Median: {self.median_final:.2%}",
            f"  Std Dev: {self.std_final:.2%}",
            "",
            "Percentiles:",
            f"  5th: {self.percentile_5:.2%}",
            f"  25th: {self.percentile_25:.2%}",
            f"  75th: {self.percentile_75:.2%}",
            f"  95th: {self.percentile_95:.2%}",
            "",
            f"Probability of Loss: {self.prob_loss:.1%}",
            f"Mean Max Drawdown: {self.max_drawdown_mean:.2%}",
            f"95th Percentile Max DD: {self.max_drawdown_95th:.2%}",
        ]
        return "\n".join(lines)


def monte_carlo_simulation(
    returns: pd.Series,
    n_simulations: int = 10000,
    time_horizon: int = 252,
    initial_value: float = 1.0,
    random_seed: int | None = None,
    save_paths: bool = False,
) -> MonteCarloResult:
    """Run Monte Carlo simulation based on historical returns.

    Uses bootstrap sampling from historical returns to simulate
    possible future paths.

    Args:
        returns: Historical return series.
        n_simulations: Number of simulation paths (default 10000).
        time_horizon: Number of periods to simulate (default 252 = 1 year).
        initial_value: Starting portfolio value (default 1.0).
        random_seed: Random seed for reproducibility.
        save_paths: Whether to save all paths (memory intensive).

    Returns:
        MonteCarloResult with simulation statistics.
    """
    if random_seed is not None:
        np.random.seed(random_seed)

    returns_arr = returns.dropna().values
    n_returns = len(returns_arr)

    if n_returns < 10:
        raise ValueError("Need at least 10 historical returns for simulation")

    # Simulate paths
    final_values = np.zeros(n_simulations)
    max_drawdowns = np.zeros(n_simulations)
    all_paths = np.zeros((n_simulations, time_horizon + 1)) if save_paths else None

    for i in range(n_simulations):
        # Bootstrap sample returns
        sampled_indices = np.random.randint(0, n_returns, size=time_horizon)
        sampled_returns = returns_arr[sampled_indices]

        # Compute path
        path = initial_value * np.cumprod(1 + sampled_returns)
        path = np.insert(path, 0, initial_value)

        final_values[i] = path[-1]

        # Compute max drawdown
        running_max = np.maximum.accumulate(path)
        drawdown = (path - running_max) / running_max
        max_drawdowns[i] = drawdown.min()

        if save_paths:
            all_paths[i] = path

    # Compute statistics
    final_returns = (final_values / initial_value) - 1

    return MonteCarloResult(
        n_simulations=n_simulations,
        time_horizon=time_horizon,
        final_values=final_values,
        all_paths=all_paths,
        mean_final=float(np.mean(final_returns)),
        median_final=float(np.median(final_returns)),
        std_final=float(np.std(final_returns)),
        percentile_5=float(np.percentile(final_returns, 5)),
        percentile_25=float(np.percentile(final_returns, 25)),
        percentile_75=float(np.percentile(final_returns, 75)),
        percentile_95=float(np.percentile(final_returns, 95)),
        prob_loss=float(np.mean(final_returns < 0)),
        max_drawdown_mean=float(np.mean(max_drawdowns)),
        max_drawdown_95th=float(np.percentile(max_drawdowns, 5)),  # 5th = worst 95%
    )


def bootstrap_returns(
    returns: pd.Series,
    n_samples: int = 1000,
    block_size: int | None = None,
    random_seed: int | None = None,
) -> np.ndarray:
    """Bootstrap sample returns with optional block bootstrap.

    Args:
        returns: Historical return series.
        n_samples: Number of samples per bootstrap (default = len(returns)).
        block_size: Size of blocks for block bootstrap (preserves autocorrelation).
        random_seed: Random seed for reproducibility.

    Returns:
        Array of bootstrapped returns.
    """
    if random_seed is not None:
        np.random.seed(random_seed)

    returns_arr = returns.dropna().values
    n = len(returns_arr)

    if n_samples is None:
        n_samples = n

    if block_size is None:
        # Simple bootstrap
        indices = np.random.randint(0, n, size=n_samples)
        return returns_arr[indices]
    else:
        # Block bootstrap
        n_blocks = int(np.ceil(n_samples / block_size))
        block_starts = np.random.randint(0, n - block_size + 1, size=n_blocks)

        bootstrapped = []
        for start in block_starts:
            bootstrapped.extend(returns_arr[start : start + block_size])

        return np.array(bootstrapped[:n_samples])


def simulate_paths(
    mu: float,
    sigma: float,
    n_paths: int = 1000,
    n_steps: int = 252,
    initial_value: float = 1.0,
    dt: float = 1 / 252,
    random_seed: int | None = None,
) -> np.ndarray:
    """Simulate price paths using Geometric Brownian Motion.

    dS = mu * S * dt + sigma * S * dW

    Args:
        mu: Annual drift (expected return).
        sigma: Annual volatility.
        n_paths: Number of paths to simulate.
        n_steps: Number of time steps.
        initial_value: Starting value.
        dt: Time step (default 1/252 = 1 trading day).
        random_seed: Random seed.

    Returns:
        Array of shape (n_paths, n_steps + 1) with simulated paths.
    """
    if random_seed is not None:
        np.random.seed(random_seed)

    # Generate random increments
    dW = np.random.standard_normal((n_paths, n_steps))

    # GBM formula: S(t+dt) = S(t) * exp((mu - 0.5*sigma^2)*dt + sigma*sqrt(dt)*dW)
    drift = (mu - 0.5 * sigma**2) * dt
    diffusion = sigma * np.sqrt(dt) * dW

    log_returns = drift + diffusion
    log_paths = np.cumsum(log_returns, axis=1)

    # Prepend zeros for initial value
    log_paths = np.hstack([np.zeros((n_paths, 1)), log_paths])

    paths = initial_value * np.exp(log_paths)

    return paths


def confidence_intervals(
    statistic_samples: np.ndarray,
    confidence_levels: list[float] | None = None,
) -> dict[float, tuple[float, float]]:
    """Calculate confidence intervals from sample distribution.

    Args:
        statistic_samples: Array of sampled statistics.
        confidence_levels: List of confidence levels (default [0.90, 0.95, 0.99]).

    Returns:
        Dictionary mapping confidence level to (lower, upper) bounds.
    """
    if confidence_levels is None:
        confidence_levels = [0.90, 0.95, 0.99]

    result = {}
    for level in confidence_levels:
        alpha = 1 - level
        lower = np.percentile(statistic_samples, alpha / 2 * 100)
        upper = np.percentile(statistic_samples, (1 - alpha / 2) * 100)
        result[level] = (float(lower), float(upper))

    return result


def probability_of_loss(
    returns: pd.Series,
    threshold: float = 0.0,
    time_horizon: int = 252,
    n_simulations: int = 10000,
    random_seed: int | None = None,
) -> float:
    """Calculate probability of loss over a time horizon.

    Args:
        returns: Historical return series.
        threshold: Loss threshold (default 0 = any loss).
        time_horizon: Time horizon in periods.
        n_simulations: Number of simulations.
        random_seed: Random seed.

    Returns:
        Probability of experiencing a loss >= threshold.
    """
    result = monte_carlo_simulation(
        returns=returns,
        n_simulations=n_simulations,
        time_horizon=time_horizon,
        random_seed=random_seed,
        save_paths=False,
    )

    final_returns = (result.final_values / 1.0) - 1
    return float(np.mean(final_returns < threshold))


def expected_shortfall_mc(
    returns: pd.Series,
    confidence_level: float = 0.95,
    time_horizon: int = 1,
    n_simulations: int = 10000,
    random_seed: int | None = None,
) -> float:
    """Calculate Expected Shortfall (CVaR) using Monte Carlo.

    Args:
        returns: Historical return series.
        confidence_level: Confidence level (default 0.95).
        time_horizon: Time horizon in periods.
        n_simulations: Number of simulations.
        random_seed: Random seed.

    Returns:
        Expected Shortfall (average loss in worst cases).
    """
    result = monte_carlo_simulation(
        returns=returns,
        n_simulations=n_simulations,
        time_horizon=time_horizon,
        random_seed=random_seed,
        save_paths=False,
    )

    final_returns = (result.final_values / 1.0) - 1

    # VaR is the threshold
    var = np.percentile(final_returns, (1 - confidence_level) * 100)

    # ES is the mean of losses worse than VaR
    worst_returns = final_returns[final_returns <= var]

    return float(np.mean(worst_returns)) if len(worst_returns) > 0 else float(var)


def stress_test(
    returns: pd.Series,
    shock_scenarios: dict[str, float],
    n_simulations: int = 1000,
    time_horizon: int = 21,  # 1 month
    random_seed: int | None = None,
) -> dict[str, MonteCarloResult]:
    """Run stress tests with different shock scenarios.

    Args:
        returns: Historical return series.
        shock_scenarios: Dict mapping scenario name to return shock.
            e.g., {"market_crash": -0.20, "flash_crash": -0.10}
        n_simulations: Simulations per scenario.
        time_horizon: Horizon for each simulation.
        random_seed: Random seed.

    Returns:
        Dictionary mapping scenario name to MonteCarloResult.
    """
    results = {}

    for name, shock in shock_scenarios.items():
        # Apply shock to returns
        shocked_returns = returns + shock / time_horizon

        result = monte_carlo_simulation(
            returns=shocked_returns,
            n_simulations=n_simulations,
            time_horizon=time_horizon,
            random_seed=random_seed,
            save_paths=False,
        )
        results[name] = result

    return results


def parameter_sensitivity(
    base_returns: pd.Series,
    parameter_range: np.ndarray,
    parameter_transform: Callable[[pd.Series, float], pd.Series],
    metric_func: Callable[[pd.Series], float],
    n_simulations: int = 1000,
    time_horizon: int = 252,
) -> pd.DataFrame:
    """Analyze sensitivity of a metric to a parameter.

    Args:
        base_returns: Base historical return series.
        parameter_range: Array of parameter values to test.
        parameter_transform: Function(returns, param) -> transformed_returns.
        metric_func: Function(returns) -> metric_value.
        n_simulations: Simulations per parameter value.
        time_horizon: Time horizon for simulation.

    Returns:
        DataFrame with parameter values and corresponding metrics.
    """
    results = []

    for param in parameter_range:
        transformed = parameter_transform(base_returns, param)
        mc_result = monte_carlo_simulation(
            returns=transformed,
            n_simulations=n_simulations,
            time_horizon=time_horizon,
            save_paths=False,
        )

        # Simulate final returns and compute metric
        final_returns = pd.Series((mc_result.final_values / 1.0) - 1)
        metric = metric_func(final_returns)

        results.append(
            {
                "parameter": param,
                "metric": metric,
                "mean_return": mc_result.mean_final,
                "std_return": mc_result.std_final,
            }
        )

    return pd.DataFrame(results)
