# Independent recomputation: Wright-Fisher haploid chain with selection, fundamental-matrix approach.
import numpy as np
from scipy.stats import binom

def chain(N, s):
    i = np.arange(N + 1)
    p = i * (1 + s) / (N + i * s)
    P = np.array([binom.pmf(np.arange(N + 1), N, pi) for pi in p])
    Q = P[1:N, 1:N]
    R = P[1:N, [0, N]]
    F = np.linalg.inv(np.eye(N - 1) - Q)       # fundamental matrix
    B = F @ R                                    # absorption probabilities
    u = B[:, 1]                                  # fixation from states 1..N-1
    # conditional time to fixation: E[T 1{fix}] / P(fix); E[T 1{fix}] = F @ u
    t_cond = (F @ u) / u
    return u[0], t_cond[0]

def kimura(N, s):
    return 1 / N if s == 0 else (1 - np.exp(-2 * s)) / (1 - np.exp(-2 * N * s))

err = lambda N, s: (kimura(N, s) - chain(N, s)[0]) / chain(N, s)[0] * 100
print("err_s01_n100", round(err(100, 0.01), 4), "err_s10_n20", round(err(20, 0.10), 4))
diffs = []
for N in (20, 50, 100):
    for s in (0.01, 0.02, 0.05):
        tp, tn = chain(N, s)[1], chain(N, -s)[1]
        diffs.append(abs(tp - tn) / ((tp + tn) / 2) * 100)
print("max_tau_diff_pct", round(max(diffs), 4))
print("neutral_res_max", max(abs(chain(N, 0.0)[0] - 1 / N) for N in (20, 50, 100)))
t100 = chain(100, 0.0)[1]
print("tau_neutral_n100", round(t100, 4), "diff_pct", round(abs(t100 - 200) / 200 * 100, 4))
# all-positive check of C1's 'systematically overestimates' across the author's grid
over = [(N, s, round(err(N, s), 3)) for N in (20, 50, 100) for s in (-0.05, -0.02, -0.01, 0.01, 0.02, 0.05, 0.10)]
print("signs of Kimura error over grid:", over)