"""Exact discrete Wright-Fisher Markov chain solver for fixation probabilities and conditional absorption times."""
import json
import math
import os


def solve_wf(N, s):
    size = N - 1
    P = [[0.0] * (N + 1) for _ in range(N + 1)]
    for i in range(1, N):
        p_prime = (i * (1.0 + s)) / (N + i * s)
        for j in range(N + 1):
            P[i][j] = math.comb(N, j) * (p_prime**j) * ((1.0 - p_prime) ** (N - j))
    P[0][0] = 1.0
    P[N][N] = 1.0

    # Solve for u (fixation probabilities from states 1..N-1)
    A = [
        [(1.0 - P[i][i] if i == j else -P[i][j]) for j in range(1, N)]
        for i in range(1, N)
    ]
    r = [P[i][N] for i in range(1, N)]

    for col in range(size):
        pivot = max(range(col, size), key=lambda row: abs(A[row][col]))
        A[col], A[pivot] = A[pivot], A[col]
        r[col], r[pivot] = r[pivot], r[col]
        pv = A[col][col]
        for row in range(col + 1, size):
            factor = A[row][col] / pv
            for k in range(col, size):
                A[row][k] -= factor * A[col][k]
            r[row] -= factor * r[col]

    u = [0.0] * size
    for row in range(size - 1, -1, -1):
        total = r[row]
        for col in range(row + 1, size):
            total -= A[row][col] * u[col]
        u[row] = total / A[row][row]

    # Conditional time to fixation (Doob h-transform conditioned on absorption into state N)
    A_cond = [
        [(1.0 - P[i][i] if i == j else -P[i][j] * (u[j - 1] / u[i - 1])) for j in range(1, N)]
        for i in range(1, N)
    ]
    b = [1.0] * size

    for col in range(size):
        pivot = max(range(col, size), key=lambda row: abs(A_cond[row][col]))
        A_cond[col], A_cond[pivot] = A_cond[pivot], A_cond[col]
        b[col], b[pivot] = b[pivot], b[col]
        pv = A_cond[col][col]
        for row in range(col + 1, size):
            factor = A_cond[row][col] / pv
            for k in range(col, size):
                A_cond[row][k] -= factor * A_cond[col][k]
            b[row] -= factor * b[col]

    tau = [0.0] * size
    for row in range(size - 1, -1, -1):
        total = b[row]
        for col in range(row + 1, size):
            total -= A_cond[row][col] * tau[col]
        tau[row] = total / A_cond[row][row]

    # Continuous diffusion approximation (Kimura 1962)
    if s == 0:
        kimura = 1.0 / N
    else:
        kimura = (1.0 - math.exp(-2.0 * s)) / (1.0 - math.exp(-2.0 * N * s))

    return u[0], kimura, tau[0]


def main():
    pop_sizes = [20, 50, 100]
    s_values = [-0.05, -0.02, -0.01, 0.0, 0.01, 0.02, 0.05, 0.10]

    grid = {}
    for N in pop_sizes:
        grid[N] = {}
        for s in s_values:
            u_exact, u_kim, tau = solve_wf(N, s)
            rel_err_pct = (u_kim - u_exact) / u_exact * 100.0 if u_exact > 0 else 0.0
            grid[N][s] = {
                "u_exact": u_exact,
                "u_kimura": u_kim,
                "rel_err_pct": rel_err_pct,
                "tau": tau,
            }

    os.makedirs("results", exist_ok=True)

    # R1: Fixation probabilities and Kimura diffusion comparison
    r1 = {
        "n_small": 20,
        "n_mid": 50,
        "n_large": 100,
        "s_weak": 0.01,
        "s_mod": 0.05,
        "s_strong": 0.1,
        "pop_sizes": pop_sizes,
        "err_s01_n100": round(grid[100][0.01]["rel_err_pct"], 2),
        "err_s10_n20": round(grid[20][0.10]["rel_err_pct"], 2),
        "u_exact_n100_s05": round(grid[100][0.05]["u_exact"], 5),
        "u_kimura_n100_s05": round(grid[100][0.05]["u_kimura"], 5),
        "err_n100_s05": round(grid[100][0.05]["rel_err_pct"], 2),
        "u_exact_n50_s01": round(grid[50][0.01]["u_exact"], 5),
        "u_kimura_n50_s01": round(grid[50][0.01]["u_kimura"], 5),
        "err_n50_s01": round(grid[50][0.01]["rel_err_pct"], 2),
    }
    with open("results/R1.json", "w") as f:
        json.dump(r1, f, indent=2)
        f.write("\n")

    # R2: Trajectory time symmetry (conditional time to fixation)
    tau_diffs = []
    for N in pop_sizes:
        for s_pos in [0.01, 0.02, 0.05]:
            tau_neg = grid[N][-s_pos]["tau"]
            tau_pos = grid[N][s_pos]["tau"]
            pct_diff = abs(tau_pos - tau_neg) / ((tau_pos + tau_neg) / 2.0) * 100.0
            tau_diffs.append(pct_diff)

    r2 = {
        "max_tau_diff_pct": round(max(tau_diffs), 2),
        "symmetry_bound_pct": 3.05,
        "tau_n50_neg05": round(grid[50][-0.05]["tau"], 2),
        "tau_n50_pos05": round(grid[50][0.05]["tau"], 2),
        "diff_n50_s05_pct": round(abs(grid[50][0.05]["tau"] - grid[50][-0.05]["tau"]) / ((grid[50][0.05]["tau"] + grid[50][-0.05]["tau"]) / 2.0) * 100.0, 2),
        "tau_n100_neg02": round(grid[100][-0.02]["tau"], 2),
        "tau_n100_pos02": round(grid[100][0.02]["tau"], 2),
        "diff_n100_s02_pct": round(abs(grid[100][0.02]["tau"] - grid[100][-0.02]["tau"]) / ((grid[100][0.02]["tau"] + grid[100][-0.02]["tau"]) / 2.0) * 100.0, 2),
    }
    with open("results/R2.json", "w") as f:
        json.dump(r2, f, indent=2)
        f.write("\n")

    # R3: Neutral limit and asymptotic convergence
    res_neutrals = [abs(grid[N][0.0]["u_exact"] - 1.0 / N) for N in pop_sizes]
    diff_2n_n100 = abs(grid[100][0.0]["tau"] - 200.0) / 200.0 * 100.0
    r3 = {
        "neutral_res_max": round(max(res_neutrals), 6),
        "tau_neutral_n20": round(grid[20][0.0]["tau"], 2),
        "tau_neutral_n50": round(grid[50][0.0]["tau"], 2),
        "tau_neutral_n100": round(grid[100][0.0]["tau"], 2),
        "kimura_ohta_bound_n100": 200.0,
        "max_diff_kimura_ohta_pct": round(diff_2n_n100, 2),
    }
    with open("results/R3.json", "w") as f:
        json.dump(r3, f, indent=2)
        f.write("\n")

    print("Computation complete. Wrote R1.json, R2.json, R3.json.")


if __name__ == "__main__":
    main()