from __future__ import annotations

import os
import time

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

    # 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
    iter_values = []
    error2 = []
    error3 = []
    error = []
    while err > tol:
        counter += 1
        q = A @ d
        alpha = rr / np.dot(d, q)
        x = x + alpha * d
        r = r - alpha * q
        rr0 = rr
        rr = np.dot(r, r)
        beta = rr / rr0
        d = r + beta * d
        iter_values.append(counter)
        error2.append(np.linalg.norm(alpha * d, np.inf))
        error3.append(np.linalg.norm(r, np.inf))
        error.append(np.linalg.norm(x - sol, np.inf))
        err = error[-1]
        if counter >= 10 * n:
            break
    cgm_time = time.perf_counter() - tic

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

    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(iter_values, error, "--k", linewidth=1.6)
    ax.loglog(iter_values, error2, "-r", linewidth=1.6)
    ax.loglog(iter_values, error3, "-b", 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")
    ax.legend(
        [" Error", " Iteration Error", " Residual"],
        loc="lower left",
        fontsize=14,
        prop={"weight": "bold", "size": 14},
    )
    style_axes(ax, 14)

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


def fig9_8():
    # laplace's equation and the CGM
    # compare different error measures

    # A: from Laplace's equation
    N = int(os.getenv("FIG9_8_N", "100"))
    M = N
    n = N * M
    lambda2 = 1
    # The entries are identical to D+SD1+SD1'+SDM+SDM'; sparse storage
    # prevents the original 10000-by-10000 matrix from using excessive memory.
    main = 2 * (1 + lambda2) * np.ones(n)
    side = -lambda2 * np.ones(n - 1)
    side[np.arange(N - 1, n - 1, N)] = 0
    vertical = -np.ones(n - M)
    A = sparse.diags(
        [main, side, side, vertical, vertical],
        [0, -1, 1, -M, M],
        shape=(n, n),
        format="csr",
    )
    # A=sparse(A);

    # specify solution (xe)
    xe = 2 * (1 - 0.5 * np.random.rand(n))
    b = A @ xe

    # fprintf('\n n = %i   condition number = %3.1e \n\n',n,cond(A,inf))

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


if __name__ == "__main__":
    fig9_8()
