import json
import numpy as np
from scipy.linalg import solve_discrete_lyapunov, svd

SEED = 7


def compress(Y, tau=1e-12):
    if Y.shape[1] == 0:
        return np.zeros((Y.shape[0], 0))
    U, s, _ = svd(Y, full_matrices=False)
    keep = s > (tau * s[0] if s[0] else 0.0)
    return U[:, keep] * s[keep]


def gramian_factors(F, B, C, K=100, tau=1e-12):
    Zp = np.zeros((F.shape[0], 0))
    Zq = np.zeros_like(Zp)
    for _ in range(K):
        Zp = compress(np.concatenate((F @ Zp, B), axis=1), tau)
        Zq = compress(np.concatenate((F.T @ Zq, C.T), axis=1), tau)
    return Zp, Zq


def balance(F, B, C, r, K=100):
    Zp, Zq = gramian_factors(F, B, C, K)
    U, s_all, Vt = svd(Zq.T @ Zp, full_matrices=False)
    rr = min(r, len(s_all))
    keep = s_all[:rr] > max(s_all[0] * 1e-13, 1e-15)
    rr = int(np.sum(keep))
    U, V, s = U[:, :rr], Vt.T[:, :rr], s_all[:rr]
    invsqrt = np.diag(1.0 / np.sqrt(s))
    Vr = Zp @ V @ invsqrt
    Wr = Zq @ U @ invsqrt
    return Wr.T @ F @ Vr, Wr.T @ B, C @ Vr, s_all, Vr, Wr


def response(F, B, C, ws):
    I = np.eye(F.shape[0])
    vals = []
    for w in ws:
        vals.append(C @ np.linalg.solve(np.exp(1j*w) * I - F, B))
    return np.asarray(vals)


def relerr(A, B):
    return float(np.linalg.norm(A - B) / max(np.linalg.norm(A), 1e-15))


def main():
    rng = np.random.default_rng(SEED)
    results = {}

    # Diagonal SSM: P_ii=1/(1-a_i^2), and finite-horizon error is a_i^(2K).
    a = np.array([.20, .55, .80, .90, .96, .985, .995])
    F = np.diag(a)
    B = np.ones((len(a), 1))
    C = np.ones((1, len(a)))
    P = solve_discrete_lyapunov(F, B @ B.T)
    Q = solve_discrete_lyapunov(F.T, C.T @ C)
    stein_p = relerr(P, F @ P @ F.T + B @ B.T)
    stein_q = relerr(Q, F.T @ Q @ F + C.T @ C)
    results['stein_relative_residuals'] = {'P': stein_p, 'Q': stein_q}

    # Prediction 1: finite-horizon Gramian tail is approximately rho^(2K).
    finite = np.zeros_like(P)
    rows = []
    for K in [1, 2, 4, 8, 16, 32, 64, 128]:
        finite = sum(np.linalg.matrix_power(F, i) @ B @ B.T @ np.linalg.matrix_power(F.T, i) for i in range(K))
        tail = relerr(P, finite)
        # Exact relative spectral-norm tail for this diagonal system.
        tail_exact = np.outer(a, a)**K / (1 - np.outer(a, a))
        exact_tail = np.linalg.norm(tail_exact) / np.linalg.norm(P)
        rows.append({'K': K, 'observed_tail': tail, 'predicted_exact_tail': float(exact_tail),
                     'asymptotic_slowest_mode': float(a[-1]**(2*K))})
    results['prediction_1_finite_horizon_geometric_tail'] = rows

    # Prediction 2: controllability/observability and Hankel singular values diverge as 1/(1-rho^2).
    scaling = []
    for rho in [.50, .70, .80, .90, .95, .98, .99, .995]:
        f = np.array([[rho]])
        p = solve_discrete_lyapunov(f, np.ones((1, 1)))[0, 0]
        q = solve_discrete_lyapunov(f, np.ones((1, 1)))[0, 0]
        scaling.append({'rho': rho, 'P': float(p), 'Q': float(q), 'predicted_1_over_1_minus_rho2': 1.0/(1-rho*rho)})
    results['prediction_2_near_instability_scaling'] = scaling

    # Prediction 3: balanced truncation error is bounded by 2*sum(discarded HSVs).
    n = 12
    eigs = np.array([.15, .25, .35, .45, .55, .62, .68, .74, .80, .86, .91, .96])
    R = rng.normal(size=(n, n)); U0, _, _ = np.linalg.svd(R)
    f = U0 @ np.diag(eigs) @ U0.T
    b = rng.normal(size=(n, 2)); c = rng.normal(size=(2, n))
    full = response(f, b, c, np.linspace(0, np.pi, 301))
    trunc_rows = []
    for r in [2, 4, 6, 8, 10]:
        fr, br, cr, hsv, Vr, Wr = balance(f, b, c, r, K=160)
        red = response(fr, br, cr, np.linspace(0, np.pi, 301))
        err = float(np.max([np.linalg.norm(x, 2) for x in (full-red)]))
        bound = float(2*np.sum(hsv[r:])) if r < len(hsv) else 0.0
        # HSV list from the balanced realization is sorted; compare the actual retained rank.
        trunc_rows.append({'r': r, 'max_frequency_error': err, 'bound_2_sum_discarded_HSV': bound,
                           'biorthogonality_error': relerr(Wr.T @ Vr, np.eye(Vr.shape[1]))})
    results['prediction_3_balanced_truncation_bound'] = trunc_rows

    # Compression mini-experiment: Gramian balancing versus magnitude state pruning.
    r = 4
    fr, br, cr, hsv, Vr, Wr = balance(f, b, c, r, K=160)
    gram_err = float(np.max([np.linalg.norm(x, 2) for x in (full-response(fr, br, cr, np.linspace(0, np.pi, 301)))]))
    # Prune states with smallest row/column energy, then preserve the same coordinates.
    score = np.sum(b*b, axis=1) + np.sum(c*c, axis=0)
    keep = np.argsort(score)[-r:]
    fp = f[np.ix_(keep, keep)]; bp = b[keep]; cp = c[:, keep]
    prune_err = float(np.max([np.linalg.norm(x, 2) for x in (full-response(fp, bp, cp, np.linspace(0, np.pi, 301)))]))
    results['compression_comparison'] = {'rank': r, 'gramian_max_frequency_error': gram_err,
                                         'magnitude_pruning_max_frequency_error': prune_err,
                                         'full_state_update_cost_units': n*n + n*2 + 2*n,
                                         'reduced_projection_state_dimension': r,
                                         'ideal_state_state_transition_cost_ratio': (r*r)/(n*n)}

    # Verify the factor iteration itself against exact Gramians.
    zp, zq = gramian_factors(f, b, c, K=160)
    results['factor_iteration_relative_errors'] = {'P': relerr(zp@zp.T, solve_discrete_lyapunov(f,b@b.T)),
                                                    'Q': relerr(zq@zq.T, solve_discrete_lyapunov(f.T,c.T@c))}
    with open('results.json', 'w') as out:
        json.dump(results, out, indent=2)
    print(json.dumps(results, indent=2))


if __name__ == '__main__':
    main()