import json
from pathlib import Path
import numpy as np
from experiment import phi


def make_problem(eps_scale, seed=2802):
    rng = np.random.default_rng(seed)
    blocks = [np.diag([1.0, 1.4]), np.diag([2.0, 2.4]),
              np.diag([4.0, 4.5]), np.diag([7.0, 7.8])]
    n = 8
    H = np.zeros((n, n))
    for i, B in enumerate(blocks): H[2*i:2*i+2, 2*i:2*i+2] = B
    for i in range(3):
        eps = eps_scale * (1.0 if i != 1 else 0.7)
        u = np.array([1., .3]); u /= np.linalg.norm(u)
        v = np.array([.4, 1.]); v /= np.linalg.norm(v)
        E = eps * np.outer(v, u)
        H[2*i+2:2*i+4, 2*i:2*i+2] = E
        H[2*i:2*i+2, 2*i+2:2*i+4] = E.T
    return H, rng.normal(size=n), blocks


def groups_fixed(blocks):
    return [list(range(2*i, 2*i+2)) for i in range(len(blocks))]


def groups_adaptive(H, blocks, tau):
    groups = groups_fixed(blocks)
    while True:
        for j in range(len(groups)-1):
            a, b = groups[j], groups[j+1]
            E = H[np.ix_(b, a)]
            eps = np.linalg.norm(E, 2)
            ea = np.linalg.eigvalsh(H[np.ix_(a, a)])
            eb = np.linalg.eigvalsh(H[np.ix_(b, b)])
            eta = np.min(np.abs(ea[:, None] - eb[None, :]))
            cert = phi(eta, eps) * np.sqrt(2) * eps
            if cert > tau:
                groups[j] = a + b
                groups.pop(j+1)
                break
        else:
            return groups


def preconditioner(H, groups):
    P = np.zeros_like(H)
    for g in groups:
        P[np.ix_(g, g)] = np.linalg.inv(H[np.ix_(g, g)])
    return P


def run(H, b, groups, steps=8):
    P = preconditioner(H, groups)
    x = np.zeros_like(b)
    optimum = np.linalg.solve(H, b)
    opt_loss = .5*optimum@H@optimum - b@optimum
    for _ in range(steps): x -= 0.9 * P @ (H @ x - b)
    loss = .5*x@H@x - b@x
    return float(max(loss - opt_loss, 0.0))


def main():
    out = []
    for coupling in [0.03, 0.3, 0.8]:
        H, b, blocks = make_problem(coupling)
        fixed = groups_fixed(blocks)
        adaptive = groups_adaptive(H, blocks, tau=0.08)
        full = [list(range(8))]
        out.append({'coupling': coupling, 'steps': 8,
                    'fixed_groups': fixed, 'adaptive_groups': adaptive,
                    'full_loss': run(H,b,full), 'fixed_loss': run(H,b,fixed),
                    'adaptive_loss': run(H,b,adaptive), 'full_matrix_entries': 64,
                    'fixed_preconditioner_entries': sum(len(g)**2 for g in fixed),
                    'adaptive_preconditioner_entries': sum(len(g)**2 for g in adaptive)})
    Path('optimizer_results.json').write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == '__main__': main()