from __future__ import annotations

import numpy as np

from scoobdoob import _conditioning
from scoobdoob._conditioning import InfeasibleBridgeError as InfeasibleBridgeError


def doob_kernels(P, S_end, horizon, w=None):
    """Time-sliced Doob kernels {Ph[n]} for n = 0..horizon-1.

    Backward recursion  h[n] = P (w * h[n+1]),  h[horizon] = uniform on S_end.
    Doob kernel         Ph[n]_ij = P_ij w_j h[n+1]_j / (P (w h[n+1]))_i.
    """

    return _conditioning.build(P, w, S_end, horizon)[0]


def sample_bridge_paths(Ph, S_start, n_paths=200, seed=0):
    """Sample paths, rejecting any infeasible start rather than dropping it."""
    return _conditioning.sample(Ph, S_start, n_paths, seed)


def sample_unconditional_paths(P, S_start, horizon, n_paths=200, seed=0):
    """Baseline: free MSM rollouts of length ``horizon`` from the start set."""
    P = np.asarray(P, dtype=np.float64)
    rng = np.random.default_rng(seed)
    K = P.shape[0]
    starts = np.asarray(list(S_start), dtype=int)
    paths = []
    for _ in range(n_paths):
        i = int(rng.choice(starts))
        path = [i]
        for _ in range(horizon):
            p = P[i]
            s = p.sum()
            p = p / s if s > 0 else np.full(K, 1.0 / K)
            i = int(rng.choice(K, p=p))
            path.append(i)
        paths.append(np.asarray(path, dtype=int))
    return paths


def reached(paths, S_end):
    """Boolean array: did each path's final state land in S_end?"""
    end = set(int(s) for s in S_end)
    return np.array([int(p[-1]) in end for p in paths], dtype=bool)


def bootstrap_ci(success_bool, n_boot=2000, alpha=0.05, seed=0):
    """Percentile bootstrap CI for the success rate of a boolean array."""
    x = np.asarray(success_bool, dtype=float)
    n = len(x)
    rate = float(x.mean()) if n else float("nan")
    if n == 0:
        return rate, (float("nan"), float("nan"))
    rng = np.random.default_rng(seed)
    boot = x[rng.integers(0, n, size=(n_boot, n))].mean(axis=1)
    lo, hi = np.percentile(boot, [100 * alpha / 2, 100 * (1 - alpha / 2)])
    return rate, (float(lo), float(hi))


def success_with_ci(paths, S_end, n_boot=2000, alpha=0.05, seed=0):
    """Convenience: success rate + bootstrap CI for a set of sampled paths."""
    s = reached(paths, S_end)
    rate, ci = bootstrap_ci(s, n_boot=n_boot, alpha=alpha, seed=seed)
    return {"rate": rate, "ci": ci, "n": len(paths), "n_success": int(s.sum())}