import json
import numpy as np


def build_q(h, q, theta):
    h = max(float(h), 1e-6)
    theta = max(float(theta), 1e-6)
    c = np.sqrt(h * theta)
    Q = np.array([
        [0, 0, -q/(2*theta), c/theta, -c/theta],
        [0, 0, 0, 1, 1],
        [0, 1, 0, 0, 0],
        [0, 0, 1, 0, 0],
        [1, 0, 0, 0, 0]], dtype=float)
    # This is the algebraic inverse of the displayed Q.  Solving the
    # first/fourth rows gives a PLUS q/(4c) term in rows 4 and 5.
    Qi = np.array([
        [0, 0, 0, 0, 1],
        [0, 0, 1, 0, 0],
        [0, 0, 0, 1, 0],
        [theta/(2*c), .5, 0, q/(4*c), 0],
        [-theta/(2*c), .5, 0, -q/(4*c), 0]], dtype=float)
    return Q, Qi


def displayed_inverse(h, q, theta):
    h = max(float(h), 1e-6); theta = max(float(theta), 1e-6)
    c = np.sqrt(h * theta)
    return np.array([
        [0, 0, 0, 0, 1], [0, 0, 1, 0, 0], [0, 0, 0, 1, 0],
        [theta/(2*c), .5, 0, -q/(4*c), 0],
        [-theta/(2*c), .5, 0, -q/(4*c), 0]], dtype=float)


def characteristic_gate(a, alpha=2.0, eps=1e-8):
    left, right = a[:-1], a[1:]
    s = np.abs(right-left) / (np.abs(right)+np.abs(left)+eps)
    g = 1.0 / (1.0 + alpha*s)
    return left * g, g


def scalar_tv_gate(a, alpha=2.0, eps=1e-8):
    left, right = a[:-1], a[1:]
    num = np.linalg.norm(right-left, axis=1)
    den = np.linalg.norm(right, axis=1)+np.linalg.norm(left, axis=1)+eps
    g = 1.0/(1.0+alpha*num/den)
    return left*g[:, None], np.broadcast_to(g[:, None], left.shape)


def metrics(pred, gates, target, base):
    acoustic_rms = float(np.sqrt(np.mean(pred[:,3:]**2)))
    base_acoustic = float(np.sqrt(np.mean(base[:,3:]**2)))
    return {
        'acoustic_rms': acoustic_rms,
        'acoustic_suppression_fraction': 1-acoustic_rms/base_acoustic,
        'equilibrium_mse_to_clean': float(np.mean((pred[:,:3]-target[:,:3])**2)),
        'total_mse_to_clean': float(np.mean((pred-target)**2)),
        'mean_gate_equilibrium': float(np.mean(gates[:,:3])),
        'mean_gate_acoustic': float(np.mean(gates[:,3:])),
        'min_gate': float(np.min(gates)), 'max_gate': float(np.max(gates)),
    }


def main():
    rng = np.random.default_rng(631)
    identity_errors = []
    literal_errors = []
    for _ in range(1000):
        h = 10**rng.uniform(-6, 1); theta = 10**rng.uniform(-6, 1)
        q = rng.normal() * np.sqrt(h*theta)
        Q, Qi = build_q(h, q, theta)
        identity_errors.append(max(np.max(np.abs(Q@Qi-np.eye(5))), np.max(np.abs(Qi@Q-np.eye(5)))))
        Qil = displayed_inverse(h, q, theta)
        literal_errors.append(max(np.max(np.abs(Q@Qil-np.eye(5))), np.max(np.abs(Qil@Q-np.eye(5)))))

    n = 256; x = np.linspace(-1, 1, n)
    shock = .5*(1+np.tanh(x/.045))
    a = np.zeros((n,5))
    a[:,0] = 1.0 + .4*shock
    a[:,1] = .15*np.sin(2*np.pi*x)
    a[:,2] = .08*np.cos(np.pi*x)
    envelope = np.exp(-(x/.22)**2)
    ringing = envelope*((-1.0)**np.arange(n))
    a[:,3] = .18*ringing; a[:,4] = -.15*ringing
    a += rng.normal(0, .002, a.shape)
    clean = a.copy(); clean[:,3:] = 0.0
    base, base_g = a[:-1], np.ones((n-1,5)); target = clean[:-1]
    gated, g_char = characteristic_gate(a)
    scalar, g_scalar = scalar_tv_gate(a)

    Q, Qi = build_q(2.0, .7, 1.3); r = rng.normal(size=5)
    result = {
        'math': {
            'corrected_inverse_max_identity_error': float(max(identity_errors)),
            'corrected_inverse_median_identity_error': float(np.median(identity_errors)),
            'literal_displayed_inverse_max_identity_error': float(max(literal_errors)),
            'literal_displayed_inverse_median_identity_error': float(np.median(literal_errors)),
            'representative_roundtrip_error': float(np.max(np.abs(Q @ (Qi @ r) - r))),
            'gate_min': float(np.min(g_char)), 'gate_max': float(np.max(g_char)),
            'all_gates_bounded': bool(np.all((g_char > 0) & (g_char <= 1))),
        },
        'baseline_no_gate': metrics(base, base_g, target, base),
        'characteristic_gate': metrics(gated, g_char, target, base),
        'scalar_tv_gate': metrics(scalar, g_scalar, target, base),
        'setup': {'n': n, 'alpha': 2.0, 'seed': 631,
                  'description': 'smooth equilibrium jump plus localized alternating acoustic ringing'}
    }
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    main()