import json
from pathlib import Path
import numpy as np


def phi(delta, eps):
    if delta == 0 and eps == 0:
        return 1.0
    return 2.0 * eps / (delta + np.sqrt(delta * delta + 4.0 * eps * eps))


def eig_dist(A, At, norm='spectral'):
    la = np.linalg.eigvalsh(A)[::-1]
    lt = np.linalg.eigvalsh(At)[::-1]
    d = np.abs(la - lt)
    if norm in ('spectral', 2, np.inf):
        return float(np.max(d))
    if norm in ('fro', 'F',  'euclidean'):
        return float(np.linalg.norm(d))
    raise ValueError(norm)


def two_by_two(gap, eps):
    A = np.array([[gap, eps], [eps, 0.0]])
    At = np.diag([gap, 0.0])
    shift = eig_dist(A, At, 'spectral')
    exact = (np.sqrt(gap * gap + 4 * eps * eps) - gap) / 2
    cert = phi(gap, eps) * eps
    return shift, exact, cert


def rank_one_check(rng, n=5, m=4, trials=100):
    worst = 0.0
    violations = 0
    for _ in range(trials):
        h1 = np.diag(np.sort(rng.uniform(2, 8, n)))
        h2 = np.diag(np.sort(rng.uniform(-4, 0, m)))
        u = rng.normal(size=n); u /= np.linalg.norm(u)
        v = rng.normal(size=m); v /= np.linalg.norm(v)
        eps = rng.uniform(0.01, 2.0)
        E = eps * np.outer(v, u)
        A = np.block([[h1, E.T], [E, h2]])
        At = np.block([[h1, np.zeros((n, m))], [np.zeros((m, n)), h2]])
        eta = np.min(np.abs(h1.diagonal()[:, None] - h2.diagonal()[None, :]))
        V = np.block([[np.zeros((n, n)), E.T], [E, np.zeros((m, m))]])
        bound = phi(eta, eps) * np.linalg.norm(V, 'fro')
        actual = eig_dist(A, At, 'fro')
        ratio = actual / bound if bound > 0 else 0
        worst = max(worst, ratio)
        violations += int(actual > bound * (1 + 1e-10))
    return {'trials': trials, 'max_actual_over_bound': worst, 'violations': violations}


def adaptive_merge_demo():
    gap = 1.0
    tau = 0.08
    rows = []
    for eps in [0.01, 0.03, 0.1, 0.3, 1.0]:
        _, _, c = two_by_two(gap, eps)
        cert_f = phi(gap, eps) * np.sqrt(2.0) * eps
        rows.append({'eps': eps, 'certificate_fro': cert_f,
                     'decision': 'retain separate blocks' if cert_f <= tau else 'merge blocks'})
    return {'gap': gap, 'threshold': tau, 'rows': rows}


def main():
    rng = np.random.default_rng(2802)
    weak = []
    for gap in [0.5, 1.0, 2.0, 4.0]:
        eps = gap * 1e-5
        actual, _, cert = two_by_two(gap, eps)
        weak.append({'gap': gap, 'eps': eps, 'observed_shift_over_eps2': actual/(eps*eps),
                     'predicted_limit_1_over_gap': 1/gap, 'certificate_ratio': actual/cert})
    strong = []
    for eps in [1, 3, 10, 30, 100]:
        actual, _, cert = two_by_two(1.0, eps)
        strong.append({'eps': eps, 'observed_shift_over_eps': actual/eps,
                       'predicted_limit': 1.0, 'certificate_ratio': actual/cert})
    zero_gap = []
    for eps in [0.01, 0.1, 1.0, 10.0]:
        actual, _, cert = two_by_two(0.0, eps)
        zero_gap.append({'eps': eps, 'observed_shift': actual, 'predicted_eps': eps,
                         'certificate_ratio': actual/cert})
    result = {'weak_coupling_prediction': weak, 'strong_coupling_prediction': strong,
              'zero_gap_prediction': zero_gap,
              'rank_one_frobenius_check': rank_one_check(rng),
              'adaptive_merge_demo': adaptive_merge_demo()}
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == '__main__':
    main()