from __future__ import annotations

from dataclasses import dataclass
from typing import Callable, List

import numpy as np
import pandas as pd
from scipy.stats import rankdata


@dataclass(slots=True)
class SobolSpec:
    params: List[str]
    bounds: List[List[float]]  # [[low, high], ...]
    n_base: int
    seed: int = 1


def sobol_sensitivity(
    spec: SobolSpec, model: Callable[[np.ndarray], float]
) -> pd.DataFrame:
    try:
        from SALib.analyze import sobol
        from SALib.sample import saltelli
    except ImportError as exc:
        raise RuntimeError(
            "SALib is required for Sobol analysis. Install with: uv add SALib"
        ) from exc

    problem = {
        "num_vars": len(spec.params),
        "names": spec.params,
        "bounds": spec.bounds,
    }
    # SALib versions < 1.4 do not support the 'seed' kwarg. Seed numpy instead for compatibility.
    rng_state = np.random.get_state()
    np.random.seed(spec.seed)
    try:
        X = saltelli.sample(problem, spec.n_base, calc_second_order=False)
    finally:
        np.random.set_state(rng_state)
    Y = np.zeros(X.shape[0])
    for i, row in enumerate(X):
        Y[i] = float(model(row))
    Si = sobol.analyze(problem, Y, calc_second_order=False, print_to_console=False)
    out = pd.DataFrame(
        {
            "Parameter": spec.params,
            "S1": Si["S1"],
            "S1_conf": Si["S1_conf"],
            "ST": Si["ST"],
            "ST_conf": Si["ST_conf"],
        }
    )
    return out


def ceac(
    effect_samples: np.ndarray,
    wtp_grid: np.ndarray,
    cost_samples: np.ndarray | None = None,
) -> pd.DataFrame:
    """Compute CEAC from samples.

    effect_samples: delta effect (e.g., cases averted); positive is better
    cost_samples: delta cost; defaults to 0
    WTP grid: values to apply to effect to form NMB = WTP * effect - cost
    """
    if cost_samples is None:
        cost_samples = np.zeros_like(effect_samples)
    out = []
    for w in wtp_grid:
        nmb = w * effect_samples - cost_samples
        prob_ce = float(np.mean(nmb > 0))
        out.append({"WTP": float(w), "PrCostEffective": prob_ce})
    return pd.DataFrame(out)


@dataclass(slots=True)
class MorrisSpec:
    params: List[str]
    bounds: List[List[float]]  # [[low, high], ...]
    levels: int = 4
    trajectories: int = 10
    seed: int = 1


def morris_sensitivity(
    spec: MorrisSpec, model: Callable[[np.ndarray], float]
) -> pd.DataFrame:
    """Run Morris screening (elementary effects) if SALib is available.

    Requires a real scalar model. Returns per-parameter mu_star and sigma.
    """
    try:
        from SALib.analyze import morris as morris_analyze
        from SALib.sample import morris as morris_sample
    except ImportError as exc:
        raise RuntimeError(
            "SALib is required for Morris analysis. Install with: uv add SALib"
        ) from exc

    problem = {
        "num_vars": len(spec.params),
        "names": spec.params,
        "bounds": spec.bounds,
    }
    rng_state = np.random.get_state()
    np.random.seed(spec.seed)
    try:
        X = morris_sample.sample(
            problem,
            N=spec.trajectories,
            num_levels=spec.levels,
            optimal_trajectories=None,
        )
    finally:
        np.random.set_state(rng_state)

    Y = np.zeros(X.shape[0])
    for i, row in enumerate(X):
        Y[i] = float(model(row))
    Si = morris_analyze.analyze(problem, X, Y, print_to_console=False)
    return pd.DataFrame(
        {"Parameter": spec.params, "mu_star": Si["mu_star"], "sigma": Si["sigma"]}
    )


def prcc(X: np.ndarray, y: np.ndarray, param_names: List[str]) -> pd.DataFrame:
    """Partial Rank Correlation Coefficients.

    X: shape (n_samples, n_params) of input samples
    y: shape (n_samples,) scalar outputs
    Returns PRCC per parameter with signs matching monotone associations.
    """
    if X.ndim != 2 or y.ndim != 1 or X.shape[0] != y.shape[0]:
        raise ValueError("Shapes must be X:(n, p), y:(n,).")
    if X.shape[1] != len(param_names):
        raise ValueError("param_names length must match X columns.")

    # Rank-transform
    Xr = np.column_stack([rankdata(X[:, j]) for j in range(X.shape[1])]).astype(float)
    yr = rankdata(y).astype(float)
    # Correlation matrix of [Xr | yr]
    Z = np.column_stack([Xr, yr])
    C = np.corrcoef(Z, rowvar=False)
    P = np.linalg.pinv(C)  # precision matrix
    prcc_vals = []
    for j in range(X.shape[1]):
        # -P_jy / sqrt(P_jj * P_yy)
        val = -P[j, -1] / np.sqrt(P[j, j] * P[-1, -1])
        prcc_vals.append(val)
    return pd.DataFrame({"Parameter": param_names, "PRCC": prcc_vals})


def cost_effectiveness_plane(
    delta_effect: np.ndarray, delta_cost: np.ndarray
) -> pd.DataFrame:
    """Return a DataFrame for the cost-effectiveness plane scatter plot.

    Inputs are sample-wise differences (strategy − comparator).
    """
    if delta_effect.shape != delta_cost.shape:
        raise ValueError("delta_effect and delta_cost must have the same shape")
    return pd.DataFrame(
        {
            "DeltaEffect": delta_effect.astype(float),
            "DeltaCost": delta_cost.astype(float),
        }
    )


def evpi(nmb_samples: np.ndarray) -> float:
    """Compute EVPI at a single WTP from NMB samples (n_samples, n_strategies)."""
    exp_max = float(np.mean(np.max(nmb_samples, axis=1)))
    max_exp = float(np.max(np.mean(nmb_samples, axis=0)))
    return exp_max - max_exp


def evppi(
    samples: pd.DataFrame, outcome_col: str, group_cols: List[str]
) -> pd.DataFrame:
    # Simple EVPI-like subgroup value (not full Bayesian) as illustration
    # EVPI = E[max_j E(NMB|j)] - max_j E(NMB) over decisions j; here we show variance reduction by grouping
    overall = samples[outcome_col].mean()
    grouped = samples.groupby(group_cols)[outcome_col].mean().reset_index(name="Mean")
    return grouped.assign(Overall=overall, Delta=lambda df: df["Mean"] - overall)