from __future__ import annotations

import numpy as np


class InfeasibleBridgeError(ValueError):
    """The requested starting distribution has no conditioned path law."""


def transition_matrix(P):
    P = np.asarray(P, dtype=float)
    if P.ndim != 2 or P.shape[0] != P.shape[1] or not len(P):
        raise ValueError("P must be a nonempty square transition matrix")
    if not np.isfinite(P).all() or (P < 0).any():
        raise ValueError("P must contain finite nonnegative probabilities")
    sums = P.sum(axis=1)

    if not np.all((sums == 0) | np.isclose(sums, 1.0, rtol=1e-8, atol=1e-12)):
        raise ValueError("Each nonzero row of P must sum to one")
    return P


def horizon_length(N):
    if isinstance(N, (bool, np.bool_)) or not isinstance(N, (int, np.integer)) or N < 1:
        raise ValueError("The sampling horizon must be a positive integer")
    return int(N)


def state_indices(states, size, name):
    values = np.asarray(list(states))
    if values.ndim != 1 or not values.size or values.dtype.kind not in "iu":
        raise ValueError(f"{name} must be a nonempty sequence of integer state indices")
    if ((values < 0) | (values >= size)).any():
        raise ValueError(f"{name} contains an out-of-range state index")
    return values.astype(int)


def weights_array(weights, N, size):
    if weights is None:
        weights = np.ones(size)
    weights = np.asarray(weights, dtype=float)
    if weights.shape == (size,):
        weights = np.broadcast_to(weights, (N, size))
    if weights.shape != (N, size):
        raise ValueError(f"Weights must have shape ({size},) or ({N}, {size})")
    if not np.isfinite(weights).all() or (weights < 0).any():
        raise ValueError("Weights must be finite and nonnegative")
    return weights


def _log_nonnegative(values):
    out = np.full(values.shape, -np.inf)
    np.log(values, out=out, where=values > 0)
    return out


def _from_log(log_values):
    with np.errstate(under="ignore", over="ignore"):
        values = np.exp(log_values)
    if not np.isfinite(values).all() or np.any(np.isfinite(log_values) & (values == 0)):
        raise FloatingPointError(
            "Backward values exceed floating-point range; no fallback or support floor was applied"
        )
    return values


def backward(P, weights, N, terminal, *, rescale=True):
    """Return h, with optional per-time scaling that leaves kernels unchanged."""
    P = transition_matrix(P)
    N = horizon_length(N)
    weights = weights_array(weights, N, len(P))
    terminal = np.asarray(terminal, dtype=float)
    if (
        terminal.shape != (len(P),)
        or not np.isfinite(terminal).all()
        or (terminal < 0).any()
    ):
        raise ValueError("Terminal weights must be a finite nonnegative vector")
    if not np.any(terminal > 0):
        raise ValueError("At least one terminal state must have positive weight")
    rows, cols = np.nonzero(P)
    log_edges = np.log(P[rows, cols])
    log_weights = _log_nonnegative(weights)
    log_h = np.full((N + 1, len(P)), -np.inf)
    log_h[N] = _log_nonnegative(terminal)
    for n in range(N - 1, -1, -1):
        edge_terms = log_edges + log_weights[n, cols] + log_h[n + 1, cols]
        np.logaddexp.at(log_h[n], rows, edge_terms)
        if rescale:
            finite = np.isfinite(log_h[n])
            if finite.any():
                log_h[n, finite] -= np.logaddexp.reduce(log_h[n, finite])
    return _from_log(log_h)


def kernel(P, weights, h_next):
    """Normalize feasible rows; return exact zero rows for infeasible states."""
    P = transition_matrix(P)
    weights = weights_array(weights, 1, len(P))[0]
    h_next = np.asarray(h_next, dtype=float)
    if h_next.shape != (len(P),) or not np.isfinite(h_next).all() or (h_next < 0).any():
        raise ValueError("The backward function must be a finite nonnegative vector")
    rows, cols = np.nonzero(P)
    scores = (
        np.log(P[rows, cols])
        + _log_nonnegative(weights)[cols]
        + _log_nonnegative(h_next)[cols]
    )
    normalizers = np.full(len(P), -np.inf)
    np.logaddexp.at(normalizers, rows, scores)
    valid = np.isfinite(scores) & np.isfinite(normalizers[rows])
    M = np.zeros_like(P)
    with np.errstate(under="ignore"):
        M[rows[valid], cols[valid]] = np.exp(scores[valid] - normalizers[rows[valid]])
    sums = M.sum(axis=1, keepdims=True)
    np.divide(M, sums, out=M, where=sums > 0)
    return M


def build(P, weights, targets, N, *, starts=None):
    P = transition_matrix(P)
    N = horizon_length(N)
    targets = np.unique(state_indices(targets, len(P), "Target set"))
    weights = weights_array(weights, N, len(P))
    terminal = np.zeros(len(P))
    terminal[targets] = 1.0 / len(targets)
    h = backward(P, weights, N, terminal)
    if starts is not None:
        starts = state_indices(starts, len(P), "Starting set")
        infeasible = starts[h[0, starts] == 0]
        if infeasible.size:
            raise InfeasibleBridgeError(
                f"No positive-weight path reaches the terminal set at horizon {N} "
                f"from starting states {infeasible.tolist()}"
            )
    return [kernel(P, weights[n], h[n + 1]) for n in range(N)], h


def checked_starts(kernels, starts):
    if not len(kernels):
        raise ValueError("At least one transition kernel is required")
    first = np.asarray(kernels[0])
    starts = state_indices(starts, len(first), "Starting set")
    rows = first[starts]
    if not np.isfinite(rows).all() or (rows < 0).any():
        raise ValueError("Invalid conditioned probabilities for a starting state")
    bad = starts[rows.sum(axis=1) == 0]
    if bad.size:
        raise InfeasibleBridgeError(
            f"Starting states {bad.tolist()} have h_0=0 at the requested horizon; "
            "the prescribed starting distribution cannot be sampled"
        )
    return starts


def sample(kernels, starts, n_paths=200, seed=0, *, rng=None):
    starts = checked_starts(kernels, starts)
    if (
        isinstance(n_paths, (bool, np.bool_))
        or not isinstance(n_paths, (int, np.integer))
        or n_paths < 0
    ):
        raise ValueError("n_paths must be a nonnegative integer")
    rng = np.random.default_rng(seed) if rng is None else rng
    size = len(kernels[0])
    paths = []
    for _ in range(n_paths):
        i = int(rng.choice(starts))
        path = [i]
        for n, matrix in enumerate(kernels):
            p = np.asarray(matrix[i], dtype=float)
            total = p.sum()
            if not np.isfinite(p).all() or (p < 0).any() or not total > 0:
                raise InfeasibleBridgeError(
                    f"Invalid or infeasible conditioned row at step {n}, state {i}"
                )
            i = int(rng.choice(size, p=p / total))
            path.append(i)
        paths.append(np.asarray(path, dtype=int))
    return paths