from __future__ import annotations

import os
import time as timer

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)
    # start iteration
    r = b - A @ x
    d = r.copy()
    rr = np.dot(r, r)
    err = 10 * tol
    counter = 0
    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
        err = np.linalg.norm(alpha * d, np.inf)
        if counter >= 10 * n:
            break
    return x


def fig9_9():
    # sparse vs non-sparse CGM for solving Ax=b

    max_isize = int(os.getenv("FIG9_9_MAX_ISIZE", "5"))
    repeats = int(os.getenv("FIG9_9_REPEATS", "6"))
    dim = np.zeros(max_isize, dtype=int)
    time_values = np.zeros(max_isize)

    for isize in range(1, max_isize + 1):
        # A from Laplace's equation
        N = 20 * isize
        M = N
        n = N * M
        lambda2 = 1
        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 = sparse.diags(
            [main, side, side, vertical, vertical],
            [0, -1, 1, -M, M],
            shape=(n, n),
            format="csr",
        )
        AA = A_sparse.toarray()

        xe = 2 * (1 - 0.5 * np.random.rand(n))
        b = AA @ xe
        dim[isize - 1] = n

        # %% nonsparse
        tic = timer.perf_counter()
        for ic in range(repeats):
            v = cgm(AA, b, xe)
        t_NS = timer.perf_counter() - tic

        # %% sparse version
        A = A_sparse
        tic = timer.perf_counter()
        for ic in range(repeats):
            v = cgm(A, b, xe)
        time_values[isize - 1] = (timer.perf_counter() - tic) / t_NS
        _ = v

    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.semilogy(dim, time_values, "--or", linewidth=1.6, markersize=9)
    ax.grid(True)
    ax.axis([0, 10000, 0.001, 1])
    ax.set_ylabel("Sparse/Non-Sparse")
    ax.set_xlabel("n-axis")
    style_axes(ax, 14)

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


if __name__ == "__main__":
    fig9_9()
