import json, math, os, random
import numpy as np

SEED = 123
np.random.seed(SEED)
random.seed(SEED)


def resolvent_iter(H, L, gamma, K):
    x = np.zeros_like(H, dtype=float)
    for _ in range(K):
        x = H + gamma * L.dot(x)
    return x


def toy_verification():
    # Scalar eigenmode: xi = H/(1-gamma*lambda), and finite-K xi_K.
    gamma = 0.9
    H = np.array([1.0])
    rows = []
    # Prediction 1: amplification is exactly 1/(1-gamma lambda).
    for lam in [0.0, 0.2, 0.5, 0.8, 0.95, 1.0]:
        L = np.array([[lam]])
        x = resolvent_iter(H, L, gamma, 1000)
        pred = 1.0 / (1.0 - gamma * lam)
        rows.append({"lambda": lam, "observed_gain": float(x[0]), "predicted_gain": pred,
                     "relative_error": abs(float(x[0])-pred)/pred})
    # Prediction 2: truncation residual after K is (gamma lambda)^K.
    lam = 0.8
    q = gamma * lam
    trunc = []
    for K in [1, 2, 4, 8, 12, 20]:
        observed = abs(float(resolvent_iter(H, np.array([[lam]]), gamma, K)[0] - 1/(1-q)))
        predicted = q**K / (1-q)
        trunc.append({"K": K, "observed_residual": observed, "predicted_residual": predicted,
                      "relative_error": abs(observed-predicted)/max(predicted, 1e-15)})
    # Prediction 3: instability at gamma*lambda > 1, with finite-K growth.
    boundary = []
    for lam in [0.8, 1.0, 1.05, 1.2]:
        q = gamma * lam
        vals = [abs(float(resolvent_iter(H, np.array([[lam]]), gamma, K)[0])) for K in [5, 10, 20]]
        boundary.append({"lambda": lam, "gamma_lambda": q, "abs_x_K5_K10_K20": vals,
                         "predicted": "bounded" if q < 1 else "diverges"})
    # Covariance prediction using scalar Gaussian/bootstrap samples.
    rng = np.random.default_rng(SEED)
    lam = 0.7; q = gamma*lam; sigma_h = 0.35
    hs = rng.normal(0, sigma_h, size=200000)
    xis = hs/(1-q)
    observed_var = float(np.var(xis)); predicted_var = sigma_h**2/(1-q)**2
    covariance = {"lambda": lam, "observed_variance": observed_var,
                  "predicted_variance": predicted_var,
                  "relative_error": abs(observed_var-predicted_var)/predicted_var}
    return {"gain_sweep": rows, "truncation_sweep": trunc, "boundary_sweep": boundary,
            "variance_check": covariance}


def offline_critic_experiment():
    # Small deterministic chain with noisy one-step rewards. The target variance is
    # estimated from bootstrap reward replicas and propagated through P.
    rng = np.random.default_rng(SEED + 1)
    n = 24; gamma = 0.9; episodes = 180; horizon = 12
    # Fixed policy transitions: mostly move right, with a terminal reward at last state.
    P = np.zeros((n,n))
    for s in range(n):
        if s == n-1: P[s,s] = 1
        else:
            P[s,s+1] = 0.82; P[s,s] += 0.18
    r_true = np.zeros(n); r_true[-1] = 1.0
    V_true = np.linalg.solve(np.eye(n)-gamma*P, r_true)
    # Dataset has state-dependent reward noise and sparse visits to late states.
    data = []
    for _ in range(episodes):
        s = 0
        for t in range(horizon):
            ns = int(rng.choice(n, p=P[s]))
            noise = rng.normal(0, 0.05 + 0.35*(s > 7))
            r = r_true[s] + noise
            data.append((s, ns, r)); s = ns
    # Replay counts and bootstrap target variance (known next-state here).
    counts = np.bincount([x[0] for x in data], minlength=n)
    B = 32
    samples = [[] for _ in range(n)]
    for s, ns, r in data:
        samples[s].append((ns,r))
    Hvar = np.zeros(n)
    Hmean = np.zeros(n)
    for s in range(n):
        if samples[s]:
            vals = np.array([r for _,r in samples[s]])
            # bootstrap means approximate one-step empirical-process uncertainty
            boot = np.array([rng.choice(vals, size=len(vals), replace=True).mean() for _ in range(B)])
            Hvar[s] = np.var(boot, ddof=1)
            Hmean[s] = vals.mean() - r_true[s]
        else:
            Hvar[s] = 0.5
    # Resolvent of standard deviation (diagonal approximation requested by idea).
    # Use transition propagation for the point uncertainty magnitude.
    Hstd = np.sqrt(Hvar + 1e-9)
    u = np.zeros(n)
    for _ in range(40): u = Hstd + gamma * P.dot(u)
    # Fitted value iteration with a shared low-dimensional critic. This avoids the
    # degenerate tabular case where each state has its own intercept and weights
    # cancel from the per-state mean.
    z = np.arange(n, dtype=float) / (n - 1)
    Phi = np.column_stack([np.ones(n), z, z*z, (z > 0.55).astype(float)])
    def fit(weighted):
        theta = np.zeros(Phi.shape[1])
        for _ in range(80):
            numer = np.zeros(Phi.shape[1]); denom = np.zeros((Phi.shape[1], Phi.shape[1]))
            for s,ns,r in data:
                y = r + gamma * float(Phi[ns].dot(theta))
                w = 1.0/(u[s] + 0.05) if weighted else 1.0
                numer += w * Phi[s] * y
                denom += w * np.outer(Phi[s], Phi[s])
            theta_new = np.linalg.solve(denom + 1e-5*np.eye(Phi.shape[1]), numer)
            theta = 0.7*theta + 0.3*theta_new
        return Phi.dot(theta)
    vb = fit(False); vi = fit(True)
    # Test error is against known population value, and Bellman residual uses true P.
    def metrics(v):
        bell = r_true + gamma*P.dot(v) - v
        return {"value_rmse": float(np.sqrt(np.mean((v-V_true)**2))),
                "bellman_rmse": float(np.sqrt(np.mean(bell**2))),
                "late_state_rmse": float(np.sqrt(np.mean((v[8:]-V_true[8:])**2)))}
    # Calibration: correlation of propagated uncertainty with absolute value error.
    cal = float(np.corrcoef(u, np.abs(vb-V_true))[0,1])
    return {"baseline": metrics(vb), "resolvent_weighted": metrics(vi),
            "uncertainty_error_correlation": cal,
            "mean_uncertainty_early": float(np.mean(u[:8])),
            "mean_uncertainty_late": float(np.mean(u[8:])),
            "dataset_count_early": int(np.sum(counts[:8])), "dataset_count_late": int(np.sum(counts[8:]))}


def main():
    out = {"seed": SEED, "toy": toy_verification(), "offline": offline_critic_experiment()}
    with open("results.json", "w") as f: json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))

if __name__ == '__main__': main()