import json
import numpy as np
from scipy.linalg import eigh


def make_gram(n=120, seed=7):
    rng = np.random.default_rng(seed)
    rows = []
    for i in range(n - 1):
        row = np.zeros(n)
        row[i], row[i + 1] = 1.0, -1.0
        rows.append(row)
    for i in range(n):
        row = np.zeros(n)
        row[i] = 0.08
        rows.append(row)
    G = np.asarray(rows)
    b = G.T @ rng.normal(size=G.shape[0])
    return G, b


def row_closures(G, blocks, tol=1e-12):
    supports = [set(np.flatnonzero(np.abs(row) > tol)) for row in G]
    closures, touching = [], []
    for block in blocks:
        js = [j for j, s in enumerate(supports) if s.intersection(block)]
        om = sorted(set().union(*(supports[j] for j in js)))
        touching.append(np.asarray(js, dtype=int))
        closures.append(np.asarray(om, dtype=int))
    return closures, touching


def build_preconditioners(G, lam=1e-4, block_size=10, keep=1):
    m, n = G.shape
    A0 = G.T @ G
    A = A0 + lam * np.eye(n)
    blocks = [np.arange(i, min(i + block_size, n)) for i in range(0, n, block_size)]
    oms, touching = row_closures(G, blocks)
    mult = np.zeros(m, dtype=int)
    for js in touching:
        mult[js] += 1
    # Exact Gram splitting, used here as a first numerical verification.
    split = np.zeros((n, n))
    local_data = []
    for om, js in zip(oms, touching):
        H = G[np.ix_(js, om)] / np.sqrt(mult[js])[:, None]
        local0 = H.T @ H
        local = local0 + lam * np.eye(len(om))
        split[np.ix_(om, om)] += local0
        local_data.append((om, local, local0))
    # P contains low generalized-energy modes on each aggregate, with the
    # row-closure Gram Schur complement providing the local energy metric.
    cols = []
    for block, (om, local, local0) in zip(blocks, local_data):
        pos = {v: k for k, v in enumerate(om)}
        ii = np.asarray([pos[v] for v in block])
        jj = np.asarray([k for k, v in enumerate(om) if v not in set(block)])
        Lbb = local0[np.ix_(ii, ii)]
        if len(jj):
            Lbg = local0[np.ix_(ii, jj)]
            Lgg = local0[np.ix_(jj, jj)]
            schur = local0[np.ix_(ii, ii)] - Lbg @ np.linalg.pinv(Lgg) @ Lbg.T
        else:
            schur = Lbb
        # D is the damped local diagonal, as in the generalized local problem.
        Dfull = np.diag(np.diag(local))
        D = Dfull[np.ix_(ii, ii)]
        vals, vecs = eigh(schur + lam * np.eye(len(ii)), D)
        take = np.argsort(vals)[:min(keep, len(ii))]
        for q in take:
            v = np.zeros(n)
            v[block] = vecs[:, q]
            v /= np.linalg.norm(v)
            cols.append(v)
    P = np.column_stack(cols) if cols else np.zeros((n, 0))
    # Columns have disjoint aggregate support, hence are already independent;
    # QR makes the construction robust to any future duplicate modes.
    P, _ = np.linalg.qr(P, mode='reduced')
    Ac = P.T @ A @ P

    def block_jacobi(v):
        z = np.zeros(n)
        counts = np.zeros(n)
        for om, local, _ in local_data:
            z[om] += np.linalg.solve(local, v[om])
            counts[om] += 1.0
        return z / counts

    def two_level(v):
        # Additive Schwarz plus Galerkin coarse correction. The coarse solve
        # uses A_c=(GP)^T(GP)+lambda P^T P exactly, without Hessian assembly.
        z = block_jacobi(v)
        if P.shape[1]:
            z += P @ np.linalg.solve(Ac, P.T @ v)
        return z

    return A, split, P, Ac, block_jacobi, two_level, mult


def pcg(A, b, M=None, tol=1e-9, maxit=500):
    x = np.zeros_like(b)
    r = b - A @ x
    z = M(r) if M else r.copy()
    p = z.copy()
    rz = float(r @ z)
    history = [np.linalg.norm(r)]
    for k in range(1, maxit + 1):
        Ap = A @ p
        alpha = rz / float(p @ Ap)
        x += alpha * p
        r -= alpha * Ap
        history.append(np.linalg.norm(r))
        if history[-1] <= tol * history[0]:
            return x, k, history
        z = M(r) if M else r.copy()
        rz_new = float(r @ z)
        p = z + (rz_new / rz) * p
        rz = rz_new
    return x, maxit, history


def main():
    G, b = make_gram()
    A, split, P, Ac, bj, tl, mult = build_preconditioners(G)
    n = A.shape[0]
    gram_error = np.linalg.norm(split - G.T @ G) / np.linalg.norm(G.T @ G)
    coarse_error = np.linalg.norm(Ac - ((G @ P).T @ (G @ P) + 1e-4 * (P.T @ P))) / max(1, np.linalg.norm(Ac))
    results = {}
    for name, M in [('unpreconditioned', None), ('block_jacobi', bj), ('two_level', tl)]:
        _, it, hist = pcg(A, b, M)
        results[name] = {'iterations': int(it), 'final_relative_residual': float(hist[-1] / hist[0]), 'history': hist}
    # A small damping sensitivity check at fixed construction settings.
    damping = {}
    for lam in [1e-6, 1e-4, 1e-2]:
        A2, _, _, _, bj2, tl2, _ = build_preconditioners(G, lam=lam)
        damping[str(lam)] = {}
        for name, M in [('block_jacobi', bj2), ('two_level', tl2)]:
            _, it, hist = pcg(A2, b, M)
            damping[str(lam)][name] = int(it)
    out = {'n': n, 'm': G.shape[0], 'coarse_dimension': int(P.shape[1]),
           'max_row_touch_multiplicity': int(mult.max()),
           'relative_gram_split_error': float(gram_error),
           'relative_coarse_gram_error': float(coarse_error),
           'results': results, 'damping_iterations': damping}
    print(json.dumps(out, indent=2))


if __name__ == '__main__':
    main()