from __future__ import annotations

import os
import time

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

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


def pcgm(A, b, M, sol):
    # %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    # subfunction for the PCGM
    # uses M (not factor of M)

    # set PCGM parameters
    tol = 1e-12

    n = len(b)
    x = np.zeros(n)

    # start iteration
    r = b - A @ x
    z = np.linalg.solve(M, r)
    d = z.copy()
    rz = np.dot(r, z)
    err = 10 * tol
    counter = 0
    iter_values = []
    error = []
    errorI = []
    while err > tol and counter < 2 * n:
        counter += 1
        q = A @ d
        alpha = rz / np.dot(d, q)
        x = x + alpha * d
        r = r - alpha * q
        z = np.linalg.solve(M, r)
        rz0 = rz
        rz = np.dot(r, z)
        beta = rz / rz0
        d = z + beta * d
        iter_values.append(counter)
        error.append(np.linalg.norm(x - sol, np.inf))
        errorI.append(np.linalg.norm(alpha * d, np.inf))
        err = error[-1]
    return np.asarray(iter_values), np.asarray(error), np.asarray(errorI)


def pcgmLU(A, b, L, U, sol):
    # %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    # subfunction for the PCGM
    # uses M (not factor of M)

    # set PCGM parameters
    tol = 1e-12

    n = len(b)
    x = np.zeros(n)

    # start iteration
    r = b - A @ x
    # z=M\r;
    y = linalg.solve_triangular(L, r, lower=True)
    z = linalg.solve_triangular(U, y, lower=False)
    d = z.copy()
    rz = np.dot(r, z)
    err = 10 * tol
    counter = 0
    iter_values = []
    error = []
    errorI = []
    while err > tol and counter < 2 * n:
        counter += 1
        q = A @ d
        alpha = rz / np.dot(d, q)
        x = x + alpha * d
        r = r - alpha * q
        # z=M\r;
        y = linalg.solve_triangular(L, r, lower=True)
        z = linalg.solve_triangular(U, y, lower=False)
        rz0 = rz
        rz = np.dot(r, z)
        beta = rz / rz0
        d = z + beta * d
        iter_values.append(counter)
        error.append(np.linalg.norm(x - sol, np.inf))
        errorI.append(np.linalg.norm(alpha * d, np.inf))
        err = error[-1]
    return np.asarray(iter_values), np.asarray(error), np.asarray(errorI)


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

    # set CGM parameters
    tol = 1e-12

    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
    iter_values = []
    error = []
    errorI = []
    while err > tol and counter < 2 * n:
        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)
        error.append(np.linalg.norm(x - sol, np.inf))
        errorI.append(np.linalg.norm(alpha * d, np.inf))
        err = error[-1]
    return np.asarray(iter_values), np.asarray(error), np.asarray(errorI)


def fig9_12():
    # comparing PCGM and CGM: figures 9.12 and 9.13

    icase = int(os.getenv("FIG9_12_ICASE", "2"))
    print("icase =", icase)

    if icase == 2:
        # 100x100 random
        n = int(os.getenv("FIG9_12_N", "100"))
        print("n =", n)
        Q, R = np.linalg.qr(np.random.randn(n, n))
        D = np.diag(np.logspace(0, 4, n))
        A = Q.T @ D @ Q
        Imax = 2 * n
    else:
        # tri-diagonal
        n = int(os.getenv("FIG9_12_N", "1000"))
        D = 2 * np.eye(n)
        SD1 = np.diag(-np.ones(n - 1), -1)
        A = D + SD1 + SD1.T
        # A=sparse(A);
        Imax = n + 20

    cond_A = np.linalg.cond(A, np.inf)
    print("cond_A =", cond_A)
    sol = np.random.rand(n)
    sol = 2 * (1 - 0.5 * np.random.rand(n))
    b = A @ sol

    # use SSOR PCGM
    LM = np.tril(A) + np.diag(np.diag(A))
    UM = np.eye(n) + np.diag(1 / np.diag(A)) @ np.tril(A).T
    M = LM @ UM
    tic = time.perf_counter()
    # [iters,errors,errorIs]=pcgm(A,b,M,sol);
    iters, errors, errorIs = pcgmLU(A, b, LM, UM, sol)
    ssor_time = time.perf_counter() - tic
    print("ssor_time =", ssor_time)
    B = np.linalg.cholesky(np.linalg.inv(M)).T
    cond_BAB = np.linalg.cond(B @ A @ B.T, 2)
    print("cond_BAB =", cond_BAB)

    # use CGM to solve matrix equation
    tic = time.perf_counter()
    iter_values, error, errorI = cgm(A, b, sol)
    cgm_time = time.perf_counter() - tic
    print("cgm_time =", cgm_time)

    fig1 = matlab_figure(1, 658, 230)
    ax1 = subaxis(1, 1, 1, 1, "MT", 0.01, "MB", 0.18, "MR", -0.02, "ML", 0.06, "P", 0.04)
    set_color_order(ax1)

    ax1.semilogy(iter_values, error, "-b", linewidth=1.6)
    ax1.grid(True)
    ax1.grid(False, which="minor")
    ax1.axis([1, Imax, 1e-12, 4])
    ax1.set_yticks([1e-12, 1e-8, 1e-4, 1])
    ax1.set_ylabel("Error")
    ax1.set_xlabel("Iteration Steps")
    # legend(' CGM','Location','SouthWest','FontSize',12,'FontWeight','bold');
    style_axes(ax1, 14)

    fig2 = matlab_figure(2, 658, 230)
    ax2 = subaxis(1, 1, 1, 1, "MT", 0.01, "MB", 0.18, "MR", -0.02, "ML", 0.06, "P", 0.04)
    set_color_order(ax2)

    ax2.semilogy(iter_values, errorI, "-b", linewidth=1.6)
    ax2.grid(True)
    ax2.grid(False, which="minor")
    ax2.axis([1, Imax, 1e-12, 4])
    ax2.set_yticks([1e-12, 1e-8, 1e-4, 1])
    ax2.set_ylabel("Iterative Error")
    ax2.set_xlabel("Iteration Steps")
    # legend(' CGM','Location','SouthWest','FontSize',12,'FontWeight','bold')
    style_axes(ax2, 14)

    # pause

    ax1.semilogy(iters, errors, "-r", linewidth=1.6)
    ax1.legend(
        [" CGM", " PCGM"],
        loc="lower left",
        fontsize=14,
        prop={"weight": "bold", "size": 14},
    )

    # if icase==1
    #     exportgraphics(gcf,'/Users/mark/Desktop/pcgmA.eps')
    # elseif icase==2
    #     exportgraphics(gcf,'/Users/mark/Desktop/pcgm1.eps')
    # end

    ax2.semilogy(iters, errorIs, "-r", linewidth=1.6)
    ax2.legend(
        [" CGM", " PCGM"],
        loc="lower left",
        fontsize=14,
        prop={"weight": "bold", "size": 14},
    )

    # if icase==1
    #     exportgraphics(gcf,'/Users/mark/Desktop/pcgmB.eps')
    # elseif icase==2
    #     exportgraphics(gcf,'/Users/mark/Desktop/pcgm2.eps')
    # end
    finish_figure(fig1, "pcgmA.png" if icase == 1 else "pcgm1.png")
    finish_figure(fig2, "pcgmB.png" if icase == 1 else "pcgm2.png")
    plt.show()


if __name__ == "__main__":
    fig9_12()
