"""The relaxation parameter: the fourth domain, and the fourth time the same
pattern appears.

omega in SOR is the oldest hand-tuned step length in numerical computing.  The
similarity form - omega = 1 + cos between successive corrections - lands
exactly in the admissible interval [1, 2] with nothing chosen, and it beats
Gauss-Seidel five-fold.  It loses to a swept omega two-fold.

A second map, derived from the theory rather than from the geometry, was tried
and is worse than the first.  Both are pinned: the working one so it is not
lost, the refuted one so it is not retried.
"""
from __future__ import annotations

import math
import os
import sys

import numpy as np
import pytest

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from relax.sor import (ContractionOmega, Fixed, SimilarityOmega, poisson_rhs,
                       residual, solve, sweep_gs)
from stage23_relaxation import anisotropic, convection, theory_check

N = 14
F = poisson_rhs(N)
W_OPT = 2.0 / (1.0 + math.sin(math.pi / (N + 1)))


def test_the_sweep_reproduces_the_sor_theory():
    """The instrument before the experiment: for the model problem the SOR
    spectral radius is exactly omega* - 1."""
    _w, theory, measured = theory_check(16)
    assert abs(theory - measured) < 0.02, (theory, measured)


def test_the_similarity_omega_lands_in_the_admissible_interval():
    """Range of the similarity and convergence condition of the method
    coincide, which is why nothing had to be clipped by hand."""
    r = SimilarityOmega()
    r.reset()
    d1 = np.array([1.0, 2.0, 3.0])
    assert r.next_omega(None, None) == 1.0
    assert r.next_omega(d1, d1) == pytest.approx(2.0)       # identical
    assert r.next_omega(d1, -d1) == pytest.approx(1.0)      # opposed
    assert r.next_omega(np.array([1.0, 0.0, 0.0]),
                        np.array([0.0, 1.0, 0.0])) == pytest.approx(1.0)


@pytest.mark.parametrize("label,coeff", [
    ("poisson", None), ("aniso", anisotropic(N, 0.1))])
def test_similarity_beats_gauss_seidel_by_a_wide_margin(label, coeff):
    gs, _ = solve(F, Fixed(1.0), coeff=coeff)
    sim, _ = solve(F, SimilarityOmega(), coeff=coeff)
    assert sim * 3 < gs, (label, sim, gs)


@pytest.mark.parametrize("label,coeff", [
    ("poisson", None), ("aniso", anisotropic(N, 0.1))])
def test_and_loses_to_a_swept_omega(label, coeff):
    """The other half, stated so the claim is never inflated."""
    sim, _ = solve(F, SimilarityOmega(), coeff=coeff)
    best = min(solve(F, Fixed(w), coeff=coeff)[0]
               for w in (1.2, 1.4, 1.6, 1.8, W_OPT))
    assert best < sim, (label, best, sim)


def test_the_similarity_rule_over_relaxes():
    """The diagnosis, not just the outcome: the cosine between successive SOR
    corrections sits near 1, so 1 + cos saturates near 2 and the rule picks a
    larger omega than the optimum."""
    _sim, mean_w = solve(F, SimilarityOmega())
    assert mean_w > W_OPT, (mean_w, W_OPT)


def test_the_contraction_map_is_refuted_and_worse():
    """The second attempt, pinned so it is not tried again.  Undamped
    feedback pins omega at the limit from the first sweep."""
    sim, _ = solve(F, SimilarityOmega())
    con, mean_w = solve(F, ContractionOmega())
    assert con > sim, (con, sim)
    assert mean_w > 1.9, mean_w


def test_convection_is_where_it_actually_hurts():
    """On a strongly non-symmetric problem Gauss-Seidel is already fast and
    the rule makes it slower - the one case where the similarity form is not
    merely second but harmful."""
    coeff = convection(N, 10.0)
    gs, _ = solve(F, Fixed(1.0), coeff=coeff)
    sim, _ = solve(F, SimilarityOmega(), coeff=coeff)
    assert sim > gs, (sim, gs)