import json, math, time
import numpy as np
from scipy.linalg import solve

SEED = 2117
rng = np.random.default_rng(SEED)

# Row-generator convention: p'(t)=p(t)Q.  State 0 is source, state 3 target.
Q0 = np.array([
    [-1.20, 0.90, 0.30, 0.00],
    [ 0.10,-1.00, 0.70, 0.20],
    [ 0.05, 0.05,-0.55, 0.45],
    [ 0.30, 0.00, 0.10,-0.40],
], dtype=float)
# Perturbed directed edge n=1 -> m=2, baseline rate 0.70.
N, M, SRC, TGT = 1, 2, 0, 3

def generator(h=0.0):
    Q = Q0.copy()
    # perturb rate additively; diagonal is adjusted to preserve row sums
    Q[N, M] = Q0[N, M] + h
    Q[N, N] = -np.sum(Q[N, np.arange(4) != N])
    return Q

def mfpt(Q, source=SRC, target=TGT):
    transient = [x for x in range(len(Q)) if x != target]
    A = Q[np.ix_(transient, transient)]
    vals = solve(A, -np.ones(len(transient)))
    return float(vals[transient.index(source)])

def stationary(Q):
    # solve Q^T pi=0 with one row replaced by normalization
    A = Q.T.copy(); b = np.zeros(len(Q))
    A[-1, :] = 1.0; b[-1] = 1.0
    return solve(A, b)

def auxiliary(Q, reset_rate):
    A = Q.copy()
    # Add a fast edge target -> source, as in the paper's auxiliary system.
    A[TGT, SRC] += reset_rate
    A[TGT, TGT] -= reset_rate
    return A

def aux_response(h, reset_rate, eps=1e-6):
    # Exact finite difference of -T(0)*d log pi_target / dh.
    pplus = stationary(auxiliary(generator(eps), reset_rate))[TGT]
    pminus = stationary(auxiliary(generator(-eps), reset_rate))[TGT]
    return -mfpt(generator(), SRC, TGT) * (math.log(pplus)-math.log(pminus))/(2*eps)

def rollout_mfpt(Q, nroll=3000):
    # Gillespie first passage samples, used as the expensive baseline.
    out = np.empty(nroll)
    for r in range(nroll):
        s, tm = SRC, 0.0
        while s != TGT and tm < 1e7:
            rates = Q[s].copy(); rates[s] = 0.0
            total = rates.sum()
            tm += rng.exponential(1.0/total)
            u = rng.random() * total
            s = int(np.searchsorted(np.cumsum(rates), u))
        out[r] = tm
    return out.mean(), out.std(ddof=1)/math.sqrt(nroll)

def main():
    baseT = mfpt(generator())
    # Formula (3a): derivative wrt W_nm.
    pi = stationary(generator())
    # Use direct all-pairs MFPT to evaluate the paper's exact response formula.
    def all_tau(Q): return np.array([[mfpt(Q,j,i) if i != j else 0.0 for j in range(4)] for i in range(4)])
    tau = all_tau(generator())
    # target k=TGT, source l=SRC, edge n->m
    R_formula = -pi[N] * (tau[TGT,N]-tau[TGT,M]) * (tau[N,TGT]+tau[TGT,SRC]-tau[N,SRC])
    eps=1e-5
    R_fd=(mfpt(generator(eps))-mfpt(generator(-eps)))/(2*eps)

    # Prediction 1: fast-reset error decreases approximately as 1/reset_rate.
    Ks=np.array([2.,5.,10.,20.,50.,100.,200.,500.,1000.])
    aux=np.array([aux_response(0.0,k) for k in Ks])
    relerr=np.abs(aux-R_fd)/max(abs(R_fd),1e-12)
    slope=np.polyfit(np.log(Ks[-5:]), np.log(np.maximum(relerr[-5:],1e-16)), 1)[0]

    # Prediction 2: exact finite perturbation curve obeys theorem's rational denominator.
    hs=np.array([-.45,-.30,-.15,.15,.30,.45])
    Sigma=tau[TGT,N]+tau[N,TGT]-tau[TGT,M]
    U=tau[N,TGT]+tau[TGT,SRC]-tau[N,SRC]
    G=tau[TGT,N]-tau[TGT,M]
    pred=[]; obs=[]; err=[]
    for h in hs:
        d=mfpt(generator(h))-baseT
        formula=-pi[N]*h*G*U/(1+pi[N]*h*Sigma)
        obs.append(d); pred.append(formula); err.append(abs(d-formula))
    # Prediction 3: symmetry/sign: make target distances equal by constructing a symmetric
    # two-branch chain and check derivative is zero; also sign follows G.
    Qsym=np.array([[-1, .5, .5, 0],[0,-1,0,1],[0,0,-1,1],[.2,0,0,-.2]],float)
    # perturb 1->2 changes branch balance; target distances from 1 and 2 equal.
    def symmf(h):
        Q=Qsym.copy(); Q[1,2]=h; Q[1,1]=-(Q[1,2]+Q[1,3]); return mfpt(Q,0,3)
    symR=(symmf(1e-6)-symmf(-1e-6))/(2e-6)
    # Sign check on two edges: edge 1->2 helpful (negative), edge 2->1 harmful (positive)
    def edgeR(a,b):
        def f(h):
            Q=generator(); Q[a,b]+=h; Q[a,a]-=h; return mfpt(Q)
        return (f(1e-5)-f(-1e-5))/(2e-5)
    signRs={'1_to_2':edgeR(1,2),'2_to_1':edgeR(2,1)}

    # Mini experiment: finite-difference rollout baseline versus stationary auxiliary estimate.
    t0=time.time(); nroll=2500; hh=.12
    p0,se0=rollout_mfpt(generator(hh),nroll); m0,se_m0=rollout_mfpt(generator(-hh),nroll)
    mc_fd=(p0-m0)/(2*hh); mc_se=math.sqrt(se0**2+se_m0**2)/(2*hh)
    exact_aux=aux_response(0,1000)
    result={
      'seed':SEED,'base_mfpt':baseT,'exact_fd_derivative':R_fd,'formula_3a':R_formula,
      'formula_abs_error':abs(R_formula-R_fd),
      'prediction_1_reset_rates':Ks.tolist(),'prediction_1_aux_response':aux.tolist(),
      'prediction_1_relative_error':relerr.tolist(),'prediction_1_loglog_slope':float(slope),
      'prediction_2_G':G,'prediction_2_U':U,'prediction_2_Sigma':Sigma,
      'prediction_2_observed_delta':obs,'prediction_2_predicted_delta':pred,
      'prediction_2_max_abs_error':max(err),
      'prediction_3_symmetric_derivative':float(symR),'prediction_3_sign_derivatives':signRs,
      'baseline_rollout_fd':mc_fd,'baseline_rollout_standard_error':mc_se,
      'idea_auxiliary_fd':float(exact_aux),'idea_abs_error_vs_exact':abs(exact_aux-R_fd),
      'runtime_sec':time.time()-t0
    }
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()