"""Stage 23 - the relaxation parameter, which is where this principle belongs.

The step rule was stuck on equations that bisection also solves.  The wider
statement is: wherever an iterative method has a step length, and there is a
current result and a required one, the step should follow how SIMILAR the two
are.  The oldest hand-tuned step length in numerical computing is the
relaxation factor omega of SOR.

    convergence needs   0 < omega < 2
    the optimum         known in closed form only for model problems
    everywhere else     guessed, swept, or left at folklore values

The similarity form.  SOR over-relaxes because the Gauss-Seidel correction
under-shoots consistently in the same direction.  How consistent that
direction is, is the cosine between successive corrections, and

    omega = 1 + max(0, cos(d_k, d_{k-1}))   lands exactly in [1, 2]

so the range of the similarity and the admissible interval of the theory
coincide.  No constant is chosen.

The instrument is checked first: on the model Poisson problem the SOR spectral
radius is exactly omega* - 1, so the measured contraction of the oracle must
reproduce it.  If it does not, the sweep is wrong and nothing below counts.

Then three problem classes, in increasing order of how unavailable the optimum
is:

    model Poisson       omega* is a theorem            - we should LOSE
    anisotropic         omega* has no closed form      - open
    variable coefficient / convection - no theory at all - open
"""
from __future__ import annotations

import math

import numpy as np

from relax.sor import (Fixed, OptimalPoisson, SimilarityOmega, poisson_rhs,
                       residual, solve, sweep_gs)

TOL = 1e-8


def theory_check(n=16):
    """The SOR spectral radius for the model problem is omega* - 1."""
    w = 2.0 / (1.0 + math.sin(math.pi / (n + 1)))
    f = poisson_rhs(n)
    h2 = 1.0 / (n + 1) ** 2
    u = np.zeros_like(f)
    rs = []
    for _ in range(60):
        sweep_gs(u, f, h2, w)
        rs.append(float(np.linalg.norm(residual(u, f, h2))))
    ratios = [rs[i + 1] / rs[i] for i in range(30, len(rs) - 1) if rs[i] > 0]
    measured = sum(ratios) / len(ratios) if ratios else float("nan")
    return w, w - 1.0, measured


def anisotropic(n, eps):
    """-u_xx - eps*u_yy: the classic problem where a single omega is wrong."""
    cw = np.full((n, n), 1.0)
    ce = np.full((n, n), 1.0)
    cs = np.full((n, n), eps)
    cn = np.full((n, n), eps)
    cc = np.full((n, n), 2.0 + 2.0 * eps)
    return (cw, ce, cs, cn, cc)


def variable(n, rng):
    """A jumping diffusion coefficient - no closed-form optimum exists."""
    a = np.exp(rng.uniform(-2.0, 2.0, (n, n)))
    cw = np.roll(a, 1, 0) * 0.5 + a * 0.5
    ce = np.roll(a, -1, 0) * 0.5 + a * 0.5
    cs = np.roll(a, 1, 1) * 0.5 + a * 0.5
    cn = np.roll(a, -1, 1) * 0.5 + a * 0.5
    cc = cw + ce + cs + cn
    return (cw, ce, cs, cn, cc)


def convection(n, pe):
    """Convection-diffusion, upwind: non-symmetric, and the SOR theory for
    symmetric problems does not apply at all."""
    cw = np.full((n, n), 1.0 + pe)
    ce = np.full((n, n), 1.0)
    cs = np.full((n, n), 1.0)
    cn = np.full((n, n), 1.0)
    cc = np.full((n, n), 4.0 + pe)
    return (cw, ce, cs, cn, cc)


def main():
    n = 24
    w_opt, rho_theory, rho_measured = theory_check(n)
    print("instrument check on the model Poisson problem")
    print(f"   omega* = {w_opt:.6f}   spectral radius: theory "
          f"{rho_theory:.6f}, measured {rho_measured:.6f}")
    if abs(rho_theory - rho_measured) > 0.02:
        print("   the sweep does NOT reproduce the theory - stopping")
        return
    print("   the sweep reproduces the theory\n")

    rng = np.random.default_rng(7)
    f = poisson_rhs(n, rng)
    grid = [
        ("model Poisson", None, "omega* is a theorem"),
        ("anisotropic eps=0.1", anisotropic(n, 0.1), "no closed form"),
        ("anisotropic eps=0.01", anisotropic(n, 0.01), "no closed form"),
        ("variable coefficient", variable(n, np.random.default_rng(3)),
         "no theory"),
        ("convection Pe=2", convection(n, 2.0), "non-symmetric"),
        ("convection Pe=10", convection(n, 10.0), "non-symmetric"),
    ]
    sweep_grid = (1.0, 1.2, 1.4, 1.6, 1.8, 1.9, w_opt)

    print("=" * 96)
    print("sweeps to ||r|| <= 1e-8 ||r0||    (lower is better; 20000 = did "
          "not converge)")
    print("=" * 96)
    print(f"{'problem':24s}{'why':16s}{'Gauss-Seidel':>14s}"
          f"{'best fixed w':>14s}{'(that w)':>10s}{'similarity':>12s}"
          f"{'mean w':>8s}")
    print("-" * 96)
    wins = total = 0
    for label, coeff, why in grid:
        gs, _ = solve(f, Fixed(1.0), tol=TOL, coeff=coeff)
        best, bw = 10 ** 9, None
        for w in sweep_grid:
            k, _ = solve(f, Fixed(w), tol=TOL, coeff=coeff)
            if k < best:
                best, bw = k, w
        sim, mw = solve(f, SimilarityOmega(), tol=TOL, coeff=coeff)
        total += 1
        wins += sim < best
        print(f"{label:24s}{why:16s}{gs:>14d}{best:>14d}{bw:>10.3f}"
              f"{sim:>12d}{mw:>8.3f}", flush=True)
    print("-" * 96)
    print(f"\n   similarity beats the best swept omega on {wins} of {total} "
          f"problems")
    print("   the swept omega is an oracle: it costs seven full solves to "
          "find,\n   and the similarity rule costs none.")


if __name__ == "__main__":
    main()