from __future__ import annotations

import os

import matplotlib.pyplot as plt
import numpy as np
from scipy import sparse

from plot_utils import finish_figure, matlab_figure, set_color_order, style_axes
from subaxis import subaxis


def cgm(A, b, sol):
    # %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    # CGM

    # stopping tolerance
    tol = 1000 * np.finfo(float).eps

    # starting point
    n = len(b)
    x = np.zeros(n)

    # CGM
    r = b - A @ x
    d = r.copy()
    rr = np.dot(r, r)
    err = 10 * tol
    counter = 0
    error = []
    while err > tol and counter < 10 * n:
        counter += 1
        q = A @ d
        alpha = rr / np.dot(d, q)
        bad = alpha * d
        x = x + bad
        r = r - alpha * q
        rr0 = rr
        rr = np.dot(r, r)
        beta = rr / rr0
        d = r + beta * d
        err = np.linalg.norm(bad, np.inf)
        value = np.linalg.norm(x - sol, np.inf)
        if value == 0:
            value = np.nextafter(0.0, 1.0)
        error.append(value)
    return np.asarray(error)


def fig9_10():
    # check CGM for solving Ax=b and
    # try different matrices

    n = int(os.getenv("FIG9_10_N", "1000"))
    print(f"\n n = {n:d} \n")

    # diagonal matrix
    A1 = sparse.diags(np.arange(1, n + 1, dtype=float), format="csr")
    cond_A1 = float(n)
    print("cond_A1 =", cond_A1)

    # matrix with exactly 3 eigenvalues a, a-1, a+1
    a = 3
    A2 = sparse.eye(n, format="lil") * a
    A2[0, n - 1] = 1
    A2[n - 1, 0] = 1
    A2 = A2.tocsr()
    cond_A2 = (a + 1) / (a - 1)
    print("cond_A2 =", cond_A2)

    # tri-diagonal
    D = 2 * np.ones(n)
    SD1 = -np.ones(n - 1)
    A3 = sparse.diags([D, SD1, SD1], [0, -1, 1], format="csr")
    eigen_min = 2 - 2 * np.cos(np.pi / (n + 1))
    eigen_max = 2 - 2 * np.cos(n * np.pi / (n + 1))
    cond_A3 = eigen_max / eigen_min
    print("cond_A3 =", cond_A3)

    # specify solution (xe)
    xe = 2 * (1 - 0.5 * np.random.rand(n))
    # xe=ones(n,1);
    b1 = A1 @ xe
    b2 = A2 @ xe
    b3 = A3 @ xe

    # use CGM and return error values
    A1_err = cgm(A1, b1, xe)
    A2_err = cgm(A2, b2, xe)
    A3_err = cgm(A3, b3, xe)

    steps1 = np.arange(1, len(A1_err) + 1)
    steps2 = np.arange(1, len(A2_err) + 1)
    steps3 = np.arange(1, len(A3_err) + 1)
    fig = matlab_figure(None, 658, 230)
    ax = subaxis(1, 1, 1, 1, "MT", 0.01, "MB", 0.18, "MR", -0.02, "ML", 0.06, "P", 0.04)
    set_color_order(ax)

    ax.loglog(steps1, A1_err, "-r", linewidth=2)
    ax.grid(True)
    ax.grid(False, which="minor")
    ax.axis([1, 1200, 1e-12, 10])
    ax.set_yticks([1e-12, 1e-8, 1e-4, 1])
    ax.set_ylabel("Error")
    ax.set_xlabel("Iteration Steps (k)")
    # legend(' A_1','Location','South','FontSize',14,'FontWeight','bold')
    style_axes(ax, 14)

    ax.loglog(steps2, A2_err, "--k", linewidth=1.6, markersize=9)
    # legend(' A_1',' A_2','Location','South','FontSize',14,'FontWeight','bold')

    ax.loglog(steps3, A3_err, "-b", linewidth=2.4)
    ax.legend(
        [" A_1", " A_2", " A_3"],
        loc="lower center",
        fontsize=14,
        prop={"weight": "bold", "size": 14},
    )

    # exportgraphics(gcf,'/Users/mark/Desktop/cgm_error.eps')
    finish_figure(fig, "cgm_error.png")
    plt.show()


if __name__ == "__main__":
    fig9_10()
