import json
import numpy as np


def quadratic_loss(H, x):
    return float(0.5 * x @ H @ x)


def orthogonal_case(rng):
    """Reflection-invariant SPD quadratic with deliberately mismatched sectors."""
    signs = np.array([1, 1, 1, -1, -1, -1.0])
    S0 = np.diag(signs)
    H0 = np.zeros((6, 6))
    H0[:3, :3] = np.diag([1.0, 2.0, 3.0])
    H0[3:, 3:] = np.diag([18.0, 27.0, 45.0])
    U1, _ = np.linalg.qr(rng.normal(size=(3, 3)))
    U2, _ = np.linalg.qr(rng.normal(size=(3, 3)))
    Q = np.zeros((6, 6)); Q[:3, :3], Q[3:, 3:] = U1, U2
    H, S = Q @ H0 @ Q.T, Q @ S0 @ Q.T
    I = np.eye(6)
    Pp, Pm = (I + S) / 2, (I - S) / 2
    cross = np.linalg.norm(Pp @ H @ Pm)
    lp = np.linalg.eigvalsh(Pp @ H @ Pp)[-1]
    lm = np.linalg.eigvalsh(Pm @ H @ Pm)[-1]

    # Exact one-step spectral test, including the stiffest mode in each sector.
    # Restrict to a basis of the sector; do not include the zero modes of P.
    se, V = np.linalg.eigh(S)
    Vp, Vm = V[:, se > 0], V[:, se < 0]
    def sector_stability(Vs, eta):
        Hs = Vs.T @ H @ Vs
        T = np.eye(Hs.shape[0]) - eta * Hs
        rho = max(abs(np.linalg.eigvalsh(T)))
        return float(rho), bool(rho < 1.0)

    stability = {}
    for name, Vs, lam in [('plus', Vp, lp), ('minus', Vm, lm)]:
        ec = 2.0 / lam
        rho_lo, ok_lo = sector_stability(Vs, .99 * ec)
        rho_hi, ok_hi = sector_stability(Vs, 1.01 * ec)
        stability[name] = {
            'predicted_boundary': float(ec),
            'rho_at_0.99_boundary': rho_lo, 'stable_at_0.99_boundary': ok_lo,
            'rho_at_1.01_boundary': rho_hi, 'stable_at_1.01_boundary': ok_hi,
        }

    eta_scalar = 1.8 / lm
    eta_p, eta_m = 1.8 / lp, 1.8 / lm
    x0 = rng.normal(size=6)

    def run(sectorwise, steps=30):
        x = x0.copy(); losses = []
        for _ in range(steps):
            losses.append(quadratic_loss(H, x))
            g = H @ x
            if sectorwise:
                x -= eta_p * (Pp @ g) + eta_m * (Pm @ g)
            else:
                x -= eta_scalar * g
        return losses, quadratic_loss(H, x)

    scalar, scalar_final = run(False)
    sector, sector_final = run(True)
    return {
        'cross_block_frobenius_norm': float(cross),
        'lambda_max_plus': float(lp), 'lambda_max_minus': float(lm),
        'curvature_ratio_minus_over_plus': float(lm / lp),
        'critical_eta_plus': float(2 / lp), 'critical_eta_minus': float(2 / lm),
        'baseline_global_eta': float(eta_scalar),
        'idea_eta_plus': float(eta_p), 'idea_eta_minus': float(eta_m),
        'baseline_loss_step_10': float(scalar[10]),
        'idea_loss_step_10': float(sector[10]),
        'baseline_final_loss': float(scalar_final),
        'idea_final_loss': float(sector_final),
        'stability_scan': stability,
    }


def nonorthogonal_case():
    """Verify that a non-orthogonal parity basis needs its metric G."""
    S0 = np.diag([1., 1., -1., -1.])
    H0 = np.diag([2., 5., 12., 20.])
    B = np.array([[1., .7, 0, 0], [0, 1., 0, 0],
                  [0, 0, 1., .5], [0, 0, 0, 1.]])
    S = B @ S0 @ np.linalg.inv(B)
    H = np.linalg.inv(B).T @ H0 @ np.linalg.inv(B)
    plus, minus = np.array([0, 1]), np.array([2, 3])
    Gp, Gm = B[:, plus].T @ B[:, plus], B[:, minus].T @ B[:, minus]
    Hp, Hm = B[:, plus].T @ H @ B[:, plus], B[:, minus].T @ H @ B[:, minus]
    gp = np.linalg.eigvals(np.linalg.solve(Gp, Hp))
    gm = np.linalg.eigvals(np.linalg.solve(Gm, Hm))
    def residual(vals, Hs, Gs):
        return max(abs(np.linalg.det(Hs - v * Gs)) for v in vals)
    return {
        'involution_error': float(np.linalg.norm(S @ S - np.eye(4))),
        'metric_nonorthogonality': float(np.linalg.norm(B.T @ B - np.eye(4))),
        'generalized_plus_eigenvalues': np.sort(gp).tolist(),
        'generalized_minus_eigenvalues': np.sort(gm).tolist(),
        'pencil_residual_plus': float(residual(gp, Hp, Gp)),
        'pencil_residual_minus': float(residual(gm, Hm, Gm)),
        'naive_euclidean_plus_eigenvalues': np.sort(np.linalg.eigvalsh(Hp)).tolist(),
        'naive_euclidean_minus_eigenvalues': np.sort(np.linalg.eigvalsh(Hm)).tolist(),
    }


if __name__ == '__main__':
    rng = np.random.default_rng(3026)
    print(json.dumps({'orthogonal': orthogonal_case(rng),
                      'nonorthogonal': nonorthogonal_case()}, indent=2))