"""Shared utilities: trajectory bounds, nested splits, Wilson CI, seeding."""
import math
import numpy as np


def get_traj_bounds(terminals):
    """Return list of (start, end_exclusive) for each trajectory."""
    ends = np.where(np.asarray(terminals) > 0.5)[0]
    bounds, start = [], 0
    for e in ends:
        bounds.append((int(start), int(e) + 1))
        start = e + 1
    return bounds


def nested_traj_order(n_traj, seed=0):
    """Fixed permutation of trajectory ids so that prefixes are NESTED across demo sizes."""
    return np.random.default_rng(seed).permutation(n_traj)


def wilson_ci(k, n, z=1.96):
    """95% Wilson score interval for a binomial proportion."""
    if n == 0:
        return 0.0, 0.0, 0.0
    p = k / n
    denom = 1.0 + z * z / n
    center = (p + z * z / (2 * n)) / denom
    half = (z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n))) / denom
    return p, max(0.0, center - half), min(1.0, center + half)


def seed_everything(seed):
    import torch
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)