# Generated by scripts/export_project_setup.py from _includes/linalg/week6_setup.md, _includes/projects/week6_setup.md
# Requires NumPy and Matplotlib. No network access is needed.
# Hjelpefunksjoner for ukens forsøk: hver bane lagrer alle løsningsforslag.
# Les særlig koblingen mellom residual, korreksjon og stoppkontroll.
# Oppsettet kjøres automatisk; det er ikke en ekstra implementasjonsoppgave.

import numpy as np
import matplotlib.pyplot as plt

# GS korrigerer én koordinat om gangen og lagrer også mellomstegene.
def gs_path(A, b, x0, sweeps=12):
    A, b = np.asarray(A, float), np.asarray(b, float)
    x = np.array(x0, dtype=float, copy=True)
    if np.any(np.diag(A) == 0):
        raise ValueError("GS krever ikke-null diagonalelementer")
    path = [x.copy()]
    for k in range(sweeps):
        for i in range(len(b)):
            # Bruk gjeldende x, inkludert koordinatene som allerede er oppdatert i sveipet.
            x[i] += (b[i] - A[i] @ x) / A[i, i]
            path.append(x.copy())
    return np.array(path)

# Bratteste nedstigning velger ny residualretning ved hvert steg, uten CGs hukommelse.
def descent_path(A, b, x0, steps=25, rtol=1e-10):
    x = np.array(x0, dtype=float, copy=True)
    path = [x.copy()]
    target = rtol * np.linalg.norm(b)
    for k in range(steps):
        # Residualen er ubalansen i de opprinnelige likningene, og kan beregnes uten fasit.
        r = b - A @ x
        if np.linalg.norm(r) <= target:
            break
        Ar = A @ r
        if r @ Ar <= 0:
            raise ValueError("Positiv krumning mangler; bruk SPD i dette forsøket")
        # Kvotienten gir minimum langs linjen x + alpha*r for en SPD-matrise.
        x += (r @ r)/(r @ Ar)*r
        path.append(x.copy())
    return np.array(path)

# Tegn nivåkurver for den kvadratiske funksjonen og baner fra valgte metoder.
def bowl_plot(ax, A, b, paths, bounds=(-1, 3, -1, 3)):
    lo, hi, bottom, top = bounds
    u, v = np.meshgrid(np.linspace(lo, hi, 180), np.linspace(bottom, top, 180))
    Z = .5*(A[0,0]*u*u + 2*A[0,1]*u*v + A[1,1]*v*v) - b[0]*u - b[1]*v
    # Direkte løsning brukes bare som referansepunkt i figuren, ikke av iterasjonen.
    star = np.linalg.solve(A, b)
    minimum = .5*star @ A @ star - b @ star
    span = max(float(Z.max()-minimum), 1.)
    ax.contour(u, v, Z, levels=minimum+np.geomspace(.01, span, 14), colors='#adb5bd')
    for name, values in paths.items():
        ax.plot(values[:,0], values[:,1], 'o-', markersize=3, label=name)
    ax.plot(*star, '*', color='black', markersize=12, label='løsning')
    ax.set(xlabel='x₁', ylabel='x₂', xlim=(lo, hi), ylim=(bottom, top))
    ax.set_aspect('equal')
    ax.legend()


def cg(A, b, x0=None, rtol=1e-8, atol=0., max_steps=1000):
    A, b = np.asarray(A, float), np.asarray(b, float)
    if A.shape != (b.size, b.size) or not np.allclose(A, A.T):
        raise ValueError('Bruk en symmetrisk kvadratisk matrise')
    x = np.zeros_like(b) if x0 is None else np.array(x0, float, copy=True)
    if x.shape != b.shape or rtol <= 0 or atol < 0 or max_steps < 1:
        raise ValueError('Kontroller startvektor, toleranser og maksimalgrense')
    if not all(np.all(np.isfinite(t)) for t in [A, b, x]):
        raise ValueError('Bruk endelige tall')
    # Residualen er ubalansen i de opprinnelige likningene, og kan beregnes uten fasit.
    r = b - A @ x
    # Startvektoren teller som første lagrede punkt, men ikke som et iterasjonssteg.
    path, residuals = [x.copy()], [np.linalg.norm(r)]
    # Absolutt margin pluss margin relativt til b; samme krav brukes ved sammenligning.
    target = atol + rtol*np.linalg.norm(b)
    matvecs = 1
    # Kontroller også startforslaget: en startvektor som oppfyller residualkravet skal gi stopp før noen divisjon.
    if residuals[-1] <= target:
        return {'path':np.array(path), 'residuals':np.array(residuals),
                'converged':True, 'matvecs':matvecs, 'preconditioner_calls':0}
    p = r.copy()
    rr = r @ r
    for k in range(max_steps):
        Ap = A @ p
        matvecs += 1
        # p.T @ A @ p er positiv for en ikke-null retning når A er SPD.
        curvature = p @ Ap
        if curvature <= 0 or not np.isfinite(curvature):
            raise ValueError('CG trenger positiv krumning; kontroller SPD')
        # Velg minimum langs søkeretningen; rr er r.T @ r i vanlig CG.
        alpha = rr / curvature
        x = x + alpha*p
        # Kort oppdatering av residualen; avrunding kan gi avvik fra direkte beregnet b-Ax.
        r = r - alpha*Ap
        # Direkte kontroll i det opprinnelige systemet.
        actual = np.linalg.norm(b - A @ x)
        matvecs += 1
        path.append(x.copy()); residuals.append(actual)
        if actual <= target:
            break
        rr_new = r @ r
        if rr_new == 0:
            break  # Rapporter manglende konvergens hvis direkte kontroll ikke var liten.
        # Kombiner ny residual med forrige retning for å bevare A-konjugerthet i eksakt regning.
        p = r + (rr_new/rr)*p
        rr = rr_new
    return {'path':np.array(path), 'residuals':np.array(residuals),
            'converged':residuals[-1] <= target, 'matvecs':matvecs,
            'preconditioner_calls':0}

# Prosjektverktøy: CG, NumPy og Matplotlib kommer fra week6-setup.
# Matematiske data og kontroller beskrives i selve prosjektet.

def coupled_problem(n, span, coupling):
    if n < 2 or span < 1 or not 0 <= coupling < 1:
        raise ValueError('Bruk n >= 2, span >= 1 og 0 <= coupling < 1')
    d = np.geomspace(1., span, n)
    B = 2*np.eye(n) - coupling*(np.eye(n,k=1)+np.eye(n,k=-1))
    return d[:,None]*B*d[None,:]

def make_week6_problems(n=60):
    """To SPD-systemer med samme kjente løsning og ulik skalering/kobling."""
    rng = np.random.default_rng(6)
    star = rng.normal(size=n)
    problems = {
        'Ujevn skalering': coupled_problem(n, 100., .25),
        'Poisson': 2*np.eye(n) - np.eye(n,k=1) - np.eye(n,k=-1),
    }
    return star, problems


def show_coordinate_change(A, b, At, bt):
    """Samme kvadratiske problem i x- og y-koordinater, for oppvarmingen."""
    fig, axes = plt.subplots(2,1,figsize=(7,9))
    bowl_plot(axes[0], A, b, {}, bounds=(-2,4,-1,3))
    bowl_plot(axes[1], At, bt, {}, bounds=(-3,5,6,14))
    axes[0].set(title='Opprinnelige x-koordinater', xlabel='x₁', ylabel='x₂')
    axes[1].set(title='Nye y-koordinater', xlabel='y₁', ylabel='y₂')
    fig.tight_layout(); plt.show()


def show_spectra(problems, transformed):
    """Egenverdier og kondisjonstall før/etter symmetrisk skalering."""
    fig, axes = plt.subplots(len(problems), 1, figsize=(7,4*len(problems)), squeeze=False)
    for ax, (name, A) in zip(axes[:,0], problems.items()):
        At = transformed[name][0]
        lam, lamt = np.linalg.eigvalsh(A), np.linalg.eigvalsh(At)
        ax.semilogy(lam, '.', label='A')
        ax.semilogy(lamt, '.', label='Symmetrisk skalert A')
        ax.set(title=name, xlabel='Sortert indeks', ylabel='Egenverdi')
        ax.legend()
        print(name, 'κ₂ før:', lam[-1]/lam[0], 'etter:', lamt[-1]/lamt[0])
    fig.tight_layout(); plt.show()


def compare_runs(A, b, star, runs):
    fig, axes = plt.subplots(2,1,figsize=(7,8))
    print(f'{"Metode":<8} {"Steg":>5} {"A-prod.":>8} {"M-løsn.":>8} '
          f'{"Rel. residual":>14} {"Rel. feil":>12} {"Konvergert":>11}')
    for name, out in runs.items():
        path = out['path']
        # I disse implementasjonene: én startberegning og to produkter per steg.
        work = 1 + 2*np.arange(len(path))
        residuals = out['residuals']/np.linalg.norm(b)
        # axis=1 tar lengden av hver rad: én feilnorm per lagret løsningsforslag.
        errors = np.linalg.norm(path-star,axis=1)/np.linalg.norm(star)
        axes[0].semilogy(work, np.maximum(residuals,1e-16), label=name)
        axes[1].semilogy(work, np.maximum(errors,1e-16), label=name)
        print(f'{name:<8} {len(path)-1:5d} {out["matvecs"]:8d} '
              f'{out["preconditioner_calls"]:8d} {residuals[-1]:14.3e} '
              f'{errors[-1]:12.3e} {str(out["converged"]):>11}')
    for ax, label in zip(axes,['Original relativ residual','Relativ løsningsfeil']):
        ax.set(xlabel='Antall A-vektor-produkter inkl. stoppkontroll',ylabel=label)
        ax.legend()
    fig.tight_layout(); plt.show()


def check_pcg(pcg_method):
    """Fire kontroller av studentens fullførte PCG før hovedforsøket."""
    A = np.diag([1.,100.]); star = np.ones(2); b = A @ star
    one = pcg_method(A, b, np.diag(A), rtol=1e-12)
    assert one['converged'] and len(one['path']) == 2, 'Diagonaltesten skal bruke ett steg.'
    assert np.allclose(one['path'][-1], star), 'Kontroller diagonalløsningen.'
    A = np.array([[1.,2.],[2.,100.]]); b = A @ star
    ref = cg(A, b, rtol=1e-12)
    identity = pcg_method(A, b, np.ones(2), rtol=1e-12)
    assert identity['converged'] and np.allclose(identity['path'],ref['path']), 'M=I skal gi CG.'
    exact = pcg_method(A, b, np.diag(A), x0=star, rtol=1e-12)
    assert exact['converged'] and len(exact['path']) == 1, 'Eksakt start skal gi null steg.'
    w = 1/np.sqrt(np.diag(A))
    At = w[:,None]*A*w[None,:]; bt = w*b
    direct = pcg_method(A, b, np.diag(A), rtol=1e-12)
    changed = cg(At, bt, rtol=1e-12)
    assert direct['converged'] and changed['converged'], 'Begge koordinatbeskrivelser skal konvergere.'
    assert np.allclose(direct['path'], changed['path']*w), 'Banene må stemme etter omregning.'
    print('Fire PCG-kontroller bestått: ett steg, M=I, eksakt start og koordinatskifte.')