import numpy as np

SEED = 1149

def stationary(T):
    w, v = np.linalg.eig(T.T)
    x = np.real(v[:, np.argmin(abs(w - 1))])
    if x.sum() < 0: x = -x
    x = np.maximum(x, 0); return x / x.sum()

def make_chain():
    # Eight states grouped into four adjacent coarse states.
    pc = np.array([18., 12., 6., 1.]) / 37.
    pi = np.repeat(pc / 2, 2)
    T = np.zeros((8, 8))
    for i in range(7):
        # Equalize neighboring probability flows, making pi stationary.
        T[i, i+1] = min(.22, .12 * min(1., pi[i+1]/pi[i]))
        T[i+1, i] = min(.22, .12 * min(1., pi[i]/pi[i+1]))
    T[np.diag_indices(8)] = 1 - T.sum(1)
    Q = T.copy(); Q[6:8] = [.5, .5, 0, 0, 0, 0, 0, 0]
    return T, Q, pi

def coarse(T, pi, alpha=0.):
    groups = [np.array([2*i, 2*i+1]) for i in range(4)]
    den = np.array([pi[g].sum() for g in groups])
    out = np.zeros((4, 4))
    for i, gi in enumerate(groups):
        for j, gj in enumerate(groups):
            out[i,j] = (pi[gi,None] * T[np.ix_(gi,gj)]).sum() / max(den[i], 1e-15)
    return out, den

def hit_time(T, start_dist, sink=(6,7)):
    keep = [i for i in range(8) if i not in sink]
    h = np.linalg.solve(np.eye(len(keep))-T[np.ix_(keep,keep)], np.ones(len(keep)))
    return sum(start_dist[i] * h[keep.index(i)] for i in keep if start_dist[i] > 0)

def sample_counts(T, n, rg, alpha=.5):
    s = stationary(T); starts = rg.choice(8, n, p=s)
    ends = np.array([rg.choice(8, p=T[i]) for i in starts])
    C = np.full((8,8), alpha); np.add.at(C, (starts, ends), 1)
    return C / C.sum(1, keepdims=True)

def main():
    P, Q, pi = make_chain(); piq = stationary(Q)
    _, pic = coarse(P, pi); _, piqc = coarse(Q, piq)
    reset = np.array([.5,.5,0,0,0,0,0,0])
    direct = hit_time(P, reset)
    hill = 1/piqc[3] - 1
    print('CORE_CHECK')
    print('P_stationarity_error %.3e' % max(abs(pi@P-pi)))
    print('Q_stationarity_error %.3e' % max(abs(piq@Q-piq)))
    print('equilibrium_target', np.round(pic,6))
    print('ness_target', np.round(piqc,6))
    print('reset_process_MFPT direct %.6f hill %.6f abs_error %.3e' % (direct,hill,abs(direct-hill)))

    print('PREDICTION_1 lag-invariant matched stationary error: lag, equilibrium, NESS')
    for lag in [1,2,4,8,16,32]:
        ep = max(abs(stationary(coarse(np.linalg.matrix_power(P,lag),pi)[0])-pic))
        eq = max(abs(stationary(coarse(np.linalg.matrix_power(Q,lag),piq)[0])-piqc))
        print(lag, '%.3e %.3e' % (ep,eq))

    print('PREDICTION_2 reset-strength: mix, dual_MFPT_error, single_MFPT_error, sink_rate')
    for mix in [0,.25,.5,.75,1.]:
        reset = np.array([1-mix,mix,0,0,0,0,0,0])
        q = P.copy(); q[6:8] = reset
        sq = stationary(q); _, sc = coarse(q,sq)
        truth = hit_time(P,reset); dual = 1/sc[3]-1
        single = 1/pic[3]-1
        print('%.2f %.6f %.6f %.6f' % (mix,abs(dual-truth),abs(single-truth),sc[3]))

    print('PREDICTION_3 sample-size RMS NESS occupancy error')
    ns=[400,1600,6400,25600]; vals=[]
    for n in ns:
        de=[]; se=[]
        for r in range(30):
            rg=np.random.default_rng(SEED+100000+n+r)
            pe=sample_counts(P,n,rg); qe=sample_counts(Q,n,rg)
            spe=stationary(pe); sqe=stationary(qe)
            _, pce=coarse(pe,spe); _, qce=coarse(qe,sqe)
            de.append(np.linalg.norm(qce-piqc)); se.append(np.linalg.norm(pce-piqc))
        d=float(np.sqrt(np.mean(np.square(de)))); s=float(np.sqrt(np.mean(np.square(se))))
        vals.append(d); print(n,'dual %.6f single %.6f'%(d,s))
    print('dual_loglog_slope %.3f predicted -0.5' % np.polyfit(np.log(ns),np.log(vals),1)[0])

if __name__ == '__main__': main()