import json
import math
from pathlib import Path
import numpy as np


def make_chain(q: float, n_states: int = 10) -> np.ndarray:
    P = np.zeros((n_states, n_states), dtype=float)
    for z in range(n_states - 1):
        P[z, z] = 1.0 - q
        P[z, z + 1] = q
    P[-1, -1] = 1.0
    return P


def backward_messages(P: np.ndarray, horizon: int, target: int) -> np.ndarray:
    h = np.zeros((horizon + 1, P.shape[0]), dtype=float)
    h[horizon, target] = 1.0
    for t in range(horizon - 1, -1, -1):
        h[t] = P @ h[t + 1]
    return h


def doob_kernel(P: np.ndarray, h: np.ndarray, t: int, z: int) -> np.ndarray:
    d = h[t, z]
    if d <= 0:
        raise ValueError("zero-feasibility state")
    raw = P[z] * h[t + 1] / d
    return raw / raw.sum()


def sample_doob(P, h, horizon, start, rng):
    z = int(start)
    path = [z]
    for t in range(horizon):
        z = int(rng.choice(P.shape[0], p=doob_kernel(P, h, t, z)))
        path.append(z)
    return np.asarray(path, dtype=int)


def conditioned_path_probability(P, h, path):
    prob = 1.0
    for t, (z, zp) in enumerate(zip(path[:-1], path[1:])):
        prob *= doob_kernel(P, h, t, int(z))[int(zp)]
    return float(prob)


def transition_kl(p, q):
    mask = p > 0
    return float(np.sum(p[mask] * (np.log(p[mask]) - np.log(np.maximum(q[mask], 1e-300)))))


def expected_conditioned_kl(P, h, horizon, start, noise):
    """Expected KL of exact conditioned paths vs a log-message-perturbed sampler."""
    occ = np.zeros(P.shape[0]); occ[start] = 1.0
    total = 0.0
    for t in range(horizon):
        nxt_occ = np.zeros_like(occ)
        for z in range(P.shape[0]):
            if occ[z] == 0 or h[t, z] <= 0:
                continue
            exact = doob_kernel(P, h, t, z)
            vals = P[z] * h[t + 1] * np.exp(noise)
            approx = vals / vals.sum()
            total += occ[z] * transition_kl(exact, approx)
            nxt_occ += occ[z] * exact
        occ = nxt_occ
    return float(total)


def rejection_cost(event_probability, n=20000, seed=0):
    """Monte Carlo mean rejection cost using geometric draws, not a huge loop."""
    rng = np.random.default_rng(seed)
    # Number of Bernoulli trials through each success is geometric(p).
    return float(rng.geometric(event_probability, size=n).mean())


def main():
    rng = np.random.default_rng(2460)
    n, target, start = 10, 9, 0
    # Verification 1: backward recursion gives exact event probability and kernels normalize.
    prob_rows = []
    for horizon in [9, 12, 20, 30, 50]:
        q = 0.5
        P = make_chain(q, n)
        h = backward_messages(P, horizon, target)
        max_resid = max(abs(doob_kernel(P, h, t, z).sum() - 1)
                         for t in range(horizon) for z in range(target + 1) if h[t, z] > 0)
        prob_rows.append({"T": horizon, "event_probability": float(h[0, start]), "kernel_max_normalization_residual": float(max_resid), "predicted_probability": float(sum(math.comb(horizon, k) * q**k * (1-q)**(horizon-k) for k in range(target, horizon+1)))})

    # Verification 2: exact Doob trajectories always hit target at the horizon.
    violation_rows = []
    for horizon in [9, 12, 20, 30, 50]:
        P = make_chain(0.5, n); h = backward_messages(P, horizon, target)
        paths = np.array([sample_doob(P, h, horizon, start, rng) for _ in range(3000)])
        violation_rows.append({"T": horizon, "violations": int(np.sum(paths[:, -1] != target)), "mean_final_state": float(paths[:, -1].mean())})

    # Verification 3: exact path law agrees with the conditional original law.
    # Enumerate all binary advance/stay strings for T=12 and compare empirical
    # Doob frequencies with original path probabilities divided by Pr(C).
    P = make_chain(0.5, n); h = backward_messages(P, 12, target)
    Tlaw = 12; samples = [sample_doob(P, h, Tlaw, start, rng) for _ in range(20000)]
    keys = [tuple(x.tolist()) for x in samples]
    empirical = {}
    for k in keys: empirical[k] = empirical.get(k, 0) + 1 / len(keys)
    exact = {}
    for bits in np.ndindex(*(2,) * Tlaw):
        z = 0; path = [z]; prob = 1.0
        for b in bits:
            zp = z + int(b) if z < target else target
            prob *= P[z, zp]; z = zp; path.append(z)
        if z == target:
            exact[tuple(path)] = prob / h[0, start]
    support = set(empirical) | set(exact)
    tv = 0.5 * sum(abs(empirical.get(k, 0.0) - exact.get(k, 0.0)) for k in support)
    law_check = {"T": Tlaw, "sample_count": len(samples),
                 "sampled_terminal_violations": int(sum(x[-1] != target for x in samples)),
                 "exact_conditioned_paths": len(exact), "empirical_exact_path_TV": float(tv)}

    # Verification 4: error grows quadratically with small log-message error.
    P = make_chain(0.5, n); h = backward_messages(P, 20, target)
    kl_rows = []
    for eps in [0.0, 0.01, 0.02, 0.05, 0.10, 0.20]:
        # deterministic bounded perturbation, applied only where message is positive
        noise = np.zeros(n)
        noise[:-1] = eps * np.linspace(-1, 1, n-1)
        kl = expected_conditioned_kl(P, h, 20, start, noise)
        kl_rows.append({"epsilon": eps, "path_KL": kl, "KL_over_epsilon2": (kl/(eps*eps) if eps else 0.0)})

    # Rare-event signature: q^9 probability; Doob costs T transitions, rejection costs 1/p.
    cost_rows = []
    for q in [0.5, 0.3, 0.2, 0.1]:
        T = 9; P = make_chain(q, n); h = backward_messages(P, T, target)
        p = float(h[0, start])
        doob_transitions = T
        rej = rejection_cost(p, n=3000, seed=123)
        cost_rows.append({"q": q, "event_probability": p, "doob_transitions_per_sample": doob_transitions, "rejection_draws_per_sample": rej, "predicted_rejection_1_over_p": 1/p})

    out = {"probability_and_normalization": prob_rows, "terminal_constraint": violation_rows, "path_law_check": law_check, "message_error_scaling": kl_rows, "rare_event_cost": cost_rows}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()