from __future__ import annotations

import os
import time

import matplotlib.pyplot as plt
import numpy as np

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


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

    # set CGM parameters
    tol = 1000 * np.finfo(float).eps

    n = len(b)
    x = np.zeros(n)
    tic = time.perf_counter()
    # start iteration
    r = b - A @ x
    d = r.copy()
    rr = np.dot(r, r)
    err = 10 * tol
    counter = 0
    steps = []
    error2 = []
    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
        steps.append(counter)
        error2.append(np.linalg.norm(bad, np.inf))
        error.append(np.linalg.norm(x - sol, np.inf))
        err = error2[-1]
    cgm_time = time.perf_counter() - tic

    print(
        f"\n  Number of CGM Iterations = {counter:d}     "
        f"Error =  {error[-1]:e}    Time = {cgm_time:e} \n"
    )

    fig = matlab_figure(None, 658, 230)
    ax = subaxis(1, 1, 1, 1, "MT", 0.01, "MB", 0.185, "MR", -0.02, "ML", 0.06, "P", 0.04)
    set_color_order(ax)

    ax.loglog(steps, error, "b", linewidth=1.6)
    ax.loglog(steps, error2, "r", linewidth=1.6)

    ax.grid(True)
    ax.grid(False, which="minor")
    ax.axis([1, 1000, 1e-12, 10])
    ax.set_yticks([1e-12, 1e-8, 1e-4, 1])
    ax.set_ylabel("Error")
    ax.set_xlabel("Iteration Steps (k)")
    ax.legend(
        [" Error", " Iterative Error"],
        loc="lower left",
        fontsize=14,
        prop={"weight": "bold", "size": 14},
    )
    style_axes(ax, 14)

    # exportgraphics(gcf,'/Users/mark/Desktop/nonconv.eps')
    finish_figure(fig, "nonconv.png")
    return x


def fig9_11():
    # check CGM for solving Ax=b and
    # for poor converging matrix

    # create symmetric matrix with specififed eigenvalues
    n = int(os.getenv("FIG9_11_N", "100"))
    Q, R = np.linalg.qr(np.random.randn(n, n))
    D = np.diag(np.logspace(0, 6, n))  # log spacing
    # D=diag(randi(10^6,1,n));  # random
    # D=diag(linspace(1,10^6,n));  # uniform spacing
    A = Q.T @ D @ Q

    cond_A = np.linalg.cond(A, 2)
    print("cond_A =", cond_A)

    # specify solution (xe)
    xe = np.ones(n)
    b = A @ xe

    print(f"\n n = {n:d} \n")

    # use CGM to solve matrix equation
    x = cgm(A, b, xe)
    _ = x
    plt.show()


if __name__ == "__main__":
    fig9_11()
