import numpy as np
import scipy as sp
from scipy import linalg as la
import scipy.sparse.linalg as spla
import matplotlib.pyplot as plt
%matplotlib inline
import matplotlib as mpl
mpl.rcParams['font.size'] = 14
mpl.rcParams['axes.labelsize'] = 20
mpl.rcParams['xtick.labelsize'] = 14
mpl.rcParams['ytick.labelsize'] = 14

"""
Explanation: <center>
    <h1> ILI285 - Computación Científica I  / INF285 - Computación Científica </h1>
    <h2> Sylvester Equation with GMRes </h2>
    <h2> <a href="#acknowledgements"> [S]cientific [C]omputing [T]eam </a> </h2>
    <h2> Version: 1.00</h2>
</center>
Table of Contents

Introduction
Short reminder about Least Squares
GMRes
Theoretical Problems
Practical Problems
Acknowledgements
End of explanation
"""


n = 10
np.random.seed(0)
A = np.random.rand(n,n)+2*np.eye(n)
#print(np.linalg.eigvals(A))
B = np.random.rand(n,n)
#print(np.linalg.eigvals(B))
C = np.random.rand(n,n)

"""
Explanation: <div id='intro' />

Sylvester Equation
https://en.wikipedia.org/wiki/Sylvester_equation
$$A\,X+X\,B=C$$
where $A\in\mathbb{R}^{n\times n}$, $B\in\mathbb{R}^{n\times n}$, $C\in\mathbb{R}^{n\times n}$ and $X\in\mathbb{R}^{n\times n}$.
$A$, $B$ and $C$ are given, the problem is to find $X$.
<div id='GMR' />

GMRes
GMRes is a member of the family of Krylov methods. It finds an approximation of $\mathbf{x}$ restricted to live on the Krylov sub-space $\mathcal{K_k}$,  where $\mathcal{K_k}={\mathbf{r}_0, A\,\mathbf{r}_0, A^2\,\mathbf{r}_0, \cdots, A^{k-1}\,\mathbf{r}_0}$ and $\mathbf{r}_0 = \mathbf{b} - A\,\mathbf{x}_0$ is the residual vector of the initial guess.
The idea behind this method is to look for improvements to the initial guess $\mathbf{x}_0$ in the Krylov space. At the $k$-th iteration, we enlarge the Krylov space by adding $A^k\,\mathbf{r}_0$, reorthogonalize the basis, and then use least squares to find the best improvement to add to $\mathbf{x}_0$.
The algorithm is as follows:
Generalized Minimum Residual Method
$\mathbf{x}0$ = initial guess<br>
$\mathbf{r}$ = $\mathbf{b} - A\,\mathbf{x}_0$ = $\mathbf{b} - $<span style="color:blue">afun</span>$(\mathbf{x}_0)$<br>
$\mathbf{q}_1$ = $\mathbf{r} / \|\mathbf{r}\|_2$<br>
for $k = 1, ..., m$<br>
$\qquad \ \ \mathbf{y} = A\,\mathbf{q}_k$ = <span style="color:blue">afun</span>$(\mathbf{q}_k)$ <br>
$\qquad$ for $j = 1,2,...,k$ <br>
$\qquad \qquad$ $h{jk} = \mathbf{q}j^*\,\mathbf{y}$<br>
$\qquad \qquad$ $\mathbf{y} = \mathbf{y} - h{jk}\, \mathbf{q}j$<br>
$\qquad$ end<br>
$\qquad \ h{k+1,k} = \|y\|2 \qquad$ (If $h{k+1,k} = 0$ , skip next line and terminate at bottom.) <br>
$\qquad \ \mathbf{q}{k+1} = \mathbf{y}/h{k+1,k}$ <br>
$\qquad$ Minimize $\left\|\widehat{H}_k\, \mathbf{c}_k - [\|\mathbf{r}\|_2 \ 0 \ 0 \ ... \ 0]^T \right\|_2$ for $\mathbf{c}_k$ <br>
$\qquad$ $\mathbf{x}_k = Q_k \, \mathbf{c}_k + \mathbf{x}_0$ <br>
end
Generating an instance of the Sylvester Equation
End of explanation
"""


def solve_JGS_iterative_Sylvester(A,B,C,m,alg=1):
    if alg==1:
        # Algorithm 1
        # AX+XB=C
        # X=A^{-1}(C-XB)
        # X^{(i+1)}=A^{-1}(C-X^{(i)} B)
        X0 = np.zeros_like(A)
        X1 = np.zeros_like(A)
        for i in range(m):
            X1=np.linalg.solve(A,C-np.dot(X0,B))
            X0=X1
            print(np.linalg.norm(np.dot(A,X1)+np.dot(X1,B)-C))
        return X1
#    elif algo==2: # TO DO 1!!!!!!!!!!!!!
        # Algorithm 2
        # AX+XB=C
        # X = (C-AX)B^{-1}
        # X^{(i+1)}=(C-A X^{(i)})B^{-1}
        # How do we implement this? Hint: You only need to use np.linalg.solve in a convenient way.

X_JGS=solve_JGS_iterative_Sylvester(A,B,C,10)

print(X_JGS)

"""
Explanation: Using a Jacobi/Gauss-Seidel iterative solver
End of explanation
"""


def compute_matrix_vector_product(x,A,B,n):
    X = np.reshape(x,(n,n))
    out = np.dot(A,X)+np.dot(X,B)
    return out.flatten()
Ax = lambda x: compute_matrix_vector_product(x,A,B,n)
afun = spla.LinearOperator((n**2, n**2), matvec=Ax)
x, exitCode = spla.gmres(afun, C.flatten(), tol=1e-10)
X_GMRes = np.reshape(x,(n,n))
print(X_GMRes)

print(np.linalg.norm(X_JGS-X_GMRes))

"""
Explanation: Using GMRes of SciPy
End of explanation
"""


Ax_JGS = Ax(X_JGS.flatten())
Ax_GMRes = Ax(X_GMRes.flatten())
c =  C.flatten()
print(np.linalg.norm(Ax_JGS-c)/np.linalg.norm(c))
print(np.linalg.norm(Ax_GMRes-c)/np.linalg.norm(c))

"""
Explanation: Computing the relative residues
End of explanation
"""


# This is a very instructive implementation of GMRes.
def GMRes_Ax(A, b, x0=np.array([0.0]), m=10, flag_display=True, threshold=1e-12):
    n = len(b)
    if len(x0)==1:
        x0=np.zeros(n)
    r0 = b - np.dot(A, x0)
    nr0=np.linalg.norm(r0)
    out_res=np.array(nr0)
    Q = np.zeros((n,n))
    H = np.zeros((n,n))
    Q[:,0] = r0 / nr0
    flag_break=False
    for k in np.arange(np.min((m,n))):
        y = np.dot(A, Q[:,k])
        if flag_display:
            print('||y||=',np.linalg.norm(y))
        for j in np.arange(k+1):
            H[j][k] = np.dot(Q[:,j], y)
            if flag_display:
                print('H[',j,'][',k,']=',H[j][k])
            y = y - np.dot(H[j][k],Q[:,j])
            if flag_display:
                print('||y||=',np.linalg.norm(y))
        # All but the last equation are treated equally. Why?
        if k+1<n:
            H[k+1][k] = np.linalg.norm(y)
            if flag_display:
                print('H[',k+1,'][',k,']=',H[k+1][k])
            if (np.abs(H[k+1][k]) > 1e-16):
                Q[:,k+1] = y/H[k+1][k]
            else:
                print('flag_break has been activated')
                flag_break=True
            # Do you remember e_1? The canonical vector.
            e1 = np.zeros((k+1)+1)        
            e1[0]=1
            H_tilde=H[0:(k+1)+1,0:k+1]
        else:
            H_tilde=H[0:k+1,0:k+1]
        # Solving the 'SMALL' least square problem. 
        # This could be improved with Givens rotations!
        ck = np.linalg.lstsq(H_tilde, nr0*e1)[0] 
        if k+1<n:
            x = x0 + np.dot(Q[:,0:(k+1)], ck)
        else:
            x = x0 + np.dot(Q, ck)
        # Why is 'norm_small' equal to 'norm_full'?
        norm_small=np.linalg.norm(np.dot(H_tilde,ck)-nr0*e1)
        out_res = np.append(out_res,norm_small)
        if flag_display:
            norm_full=np.linalg.norm(b-np.dot(A,x))
            print('..........||b-A\,x_k||=',norm_full)
            print('..........||H_k\,c_k-nr0*e1||',norm_small);
        if flag_break:
            if flag_display: 
                print('EXIT: flag_break=True')
            break
        if norm_small<threshold:
            if flag_display:
                print('EXIT: norm_small<threshold')
            break
    return x,out_res

"""
Explanation: To Do 2: What do we need to change in our implementation of GMRes to be able to use the lambda function "Ax"?
End of explanation
"""


