PA-extra

import numpy as np

np.set_printoptions(precision=3, suppress=True)


def eliminate_state(P, k):
    m = len(P)
    P_kk = P[k, k]

    info = {
        'P_kk': P_kk,
        'col_k': P[:, k].copy()
    }

    states = [i for i in range(m) if i != k]
    P_new = np.zeros((m - 1, m - 1))

    for i_new, i in enumerate(states):
        for j_new, j in enumerate(states):
            P_new[i_new, j_new] = P[i, j] + (P[i, k] * P[k, j]) / (1 - P_kk)

    return P_new, info


def reduction_phase(P):
    n = len(P)
    current_P = P.copy()
    elim_info = []

    print("Reduction Phase")
    print(f"\nOriginal Matrix:\n{current_P}\n")

    for step in range(n - 1):
        state_num = n - step
        last_idx = len(current_P) - 1

        current_P, info = eliminate_state(current_P, last_idx)
        info['state_num'] = state_num
        elim_info.append(info)

        print(f"After eliminating state {state_num}:")
        print(current_P)
        print()

    return elim_info


def enlargement_phase(elim_info):
    pi = np.array([1.0])
    states = [1]

    print("Enlargement Phase")
    print(f"Start: states={states}, π={pi}\n")

    for info in reversed(elim_info):
        state_num = info['state_num']
        P_kk = info['P_kk']
        col_k = info['col_k']

        # Compute flow into eliminated state
        flow = sum(pi[i] * col_k[s - 1] for i, s in enumerate(states))
        pi_k = flow / (1 - P_kk)

        # Insert new state in sorted order
        states.append(state_num)
        states.sort()
        pos = states.index(state_num)
        pi = np.insert(pi, pos, pi_k)
        pi = pi / pi.sum()

        print(f"Add state {state_num}: states={states}")
        print(f"π = {pi}\n")

    return pi


def verify(pi, P):
    print("VERIFICATION")
    print("=" * 50)

    sum_check = abs(pi.sum() - 1.0) < 1e-6
    print(f"Sum = {pi.sum():.10f} {'✓' if sum_check else '✗'}")

    pi_P = pi @ P
    error = np.linalg.norm(pi_P - pi)
    steady_check = error < 1e-6
    print(f"||πP - π|| = {error:.2e} {'✓' if steady_check else '✗'}")

    print(f"\nFinal π = {pi}")
    print(f"Expected: [1/6, 1/4, 7/36, 1/6, 8/36]")
    print(f"         = [{1/6:.3f}, {1/4:.3f}, {7/36:.3f}, {1/6:.3f}, {8/36:.3f}]")

    return sum_check and steady_check


if __name__ == "__main__":
    # 5-state Markov chain from Sheskin (1985) Section 4
    P = np.array([
        [0.0,  0.5,  0.0,  0.0,  0.5],
        [0.25, 0.0,  0.25, 0.5,  0.0],
        [0.0,  0.5,  0.0,  0.5,  0.0],
        [0.0,  0.25, 0.25, 0.0,  0.5],
        [0.5,  0.0,  0.0,  0.5,  0.0]
    ])

    elim_info = reduction_phase(P)
    pi = enlargement_phase(elim_info)
    valid = verify(pi, P)

    print(f"\n{'SUCCESS' if valid else 'FAILED'}")