import json
import time
from pathlib import Path
import numpy as np
from scipy.sparse import diags
from scipy.linalg import solve

SEED = 2403
rng = np.random.default_rng(SEED)
OUT = Path("results.json")


def stiffness(n):
    """Dirichlet P1 stiffness on n interior nodes, with h-scaled nodal values."""
    h = 1.0 / (n + 1)
    K = (1.0 / h) * (2*np.eye(n) - np.eye(n, k=1) - np.eye(n, k=-1))
    x = np.arange(1, n+1) * h
    return x, K


def prolongation(nc, nf):
    """P1 injection from nc interior nodes to nf interior nodes, requiring nested grids."""
    assert (nf + 1) % (nc + 1) == 0
    ratio = (nf + 1) // (nc + 1)
    P = np.zeros((nf, nc))
    for j in range(1, nc+1):
        coarse_x = j * ratio
        for i in range(max(1, coarse_x-ratio), min(nf, coarse_x+ratio)+1):
            P[i-1, j-1] = max(0.0, 1.0 - abs(i-coarse_x)/ratio)
    return P


def make_data(n=128):
    x, K = stiffness(n)
    # Manufactured exact solution u*=sin(pi x); f=-u''=pi^2 sin(pi x).
    ustar = np.sin(np.pi*x)
    f = K @ ustar  # exact for the discrete problem, avoids load quadrature error
    return x, K, ustar, f


def riesz_score(u, K, f, P):
    A = P.T @ K @ P
    b = P.T @ (f - K @ u)
    z = solve(A, b, assume_a="pos")
    return float(np.sqrt(max(0.0, z @ A @ z)))


def energy_error(u, ustar, K):
    e = u-ustar
    return float(np.sqrt(max(0.0, e @ K @ e)))


def run():
    n = 127
    x, K, ustar, f = make_data(n)
    # Candidate archive: progressively improving solutions plus oscillatory and smooth errors.
    modes = [np.sin(np.pi*x), np.sin(2*np.pi*x), np.sin(7*np.pi*x), np.sin(15*np.pi*x)]
    archive = []
    for k in range(100):
        t = k/99.0
        # Main error decays, while a few checkpoints have low-frequency bias that is
        # hard for small auxiliary spaces to detect.
        amp = 0.95*(1-t) + 0.025
        e = amp*(0.55*modes[1] + 0.25*modes[2])
        e += (0.10 + 0.35*(1-t))*modes[3]
        e += 0.018*rng.normal(size=n)
        if k in (16, 48, 76):
            e += 0.16*modes[1]
        archive.append(ustar + e)
    archive = np.asarray(archive)
    oracle = np.array([energy_error(u, ustar, K) for u in archive])

    levels = [3, 7, 15, 31, 63]
    # Full auxiliary space is the exact finite-dimensional Riesz representative.
    full_scores = np.array([riesz_score(u, K, f, np.eye(n)) for u in archive])
    full_identity_max_relerr = float(np.max(np.abs(full_scores-oracle) / np.maximum(oracle, 1e-15)))
    scores = {}
    ratios = {}
    for nc in levels:
        P = prolongation(nc, n)
        scores[nc] = np.array([riesz_score(u, K, f, P) for u in archive])
        ratios[nc] = scores[nc] / oracle

    # Prediction 1: Galerkin score is a lower bound for every candidate.
    lower_bound = {str(m): float(np.max(scores[m] <= oracle + 2e-10)) for m in levels}
    # Prediction 2: nested auxiliary spaces are pointwise monotone.
    monotone = {str(a): {
        "min_increment": float(np.min(scores[b] - scores[a])),
        "max_monotonicity_violation": float(np.max(np.maximum(0.0, scores[a] - scores[b])))
    } for a,b in zip(levels[:-1], levels[1:])}
    # Prediction 3: refinement converges to the exact energy error; report relative gap.
    convergence = {str(m): float(np.mean(np.abs(1.0-ratios[m]))) for m in levels}
    max_gap = {str(m): float(np.max(1.0-ratios[m])) for m in levels}

    # Ranking quality and selection comparison.
    oracle_rank = np.argsort(oracle)
    rank_corr = {}
    selected = {}
    for m in levels:
        # Spearman correlation computed with rank arrays, no scipy.stats dependency needed.
        r1 = np.empty(100, dtype=int); r1[np.argsort(oracle)] = np.arange(100)
        r2 = np.empty(100, dtype=int); r2[np.argsort(scores[m])] = np.arange(100)
        rank_corr[str(m)] = float(np.corrcoef(r1, r2)[0,1])
        j = int(np.argmin(scores[m]))
        selected[str(m)] = {"checkpoint": j, "score": float(scores[m][j]), "oracle_error": float(oracle[j])}

    # A deliberately realistic raw training loss proxy: sampled pointwise residual,
    # with noise, and a pointwise residual monitor on a separate coarse sample.
    raw_loss = np.array([np.mean((f-K@u)**2) for u in archive])
    raw_loss += rng.normal(0, 0.03*np.std(raw_loss), size=100)
    raw_j = int(np.argmin(raw_loss))
    # Measure the actual cheap post-processing wall time on this tiny problem.
    t0 = time.perf_counter()
    for nc in levels:
        P = prolongation(nc, n)
        _ = [riesz_score(u, K, f, P) for u in archive]
    riesz_seconds = time.perf_counter() - t0
    # Riesz postprocessing cost is measured as number of tiny solves.
    result = {
        "seed": SEED, "n_fine": n, "levels": levels,
        "full_space_identity_max_relative_error": full_identity_max_relerr,
        "riesz_postprocess_seconds": riesz_seconds,
        "predictions": {
            "lower_bound_fraction": lower_bound,
            "nested_min_increment": monotone,
            "mean_relative_gap_to_oracle": convergence,
            "max_relative_gap_to_oracle": max_gap,
            "interpretation": "All scores must be <= oracle; nested increments must be >= 0; gaps should decrease with level."
        },
        "selection": {
            "raw_training_loss": {"checkpoint": raw_j, "oracle_error": float(oracle[raw_j])},
            "riesz": selected,
            "oracle_best": {"checkpoint": int(np.argmin(oracle)), "oracle_error": float(np.min(oracle))},
            "spearman_rank_correlation": rank_corr
        },
        "oracle_error_range": [float(np.min(oracle)), float(np.max(oracle))],
        "postprocess_solves": len(levels)*len(archive)
    }
    OUT.write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    run()