"""Module providing a rank-SVD implementations."""

import os
import math
import numpy as np
from .utils import *

# set random seed
np.random.seed(42)

##########################
## Gradient computation ##
##########################

def compute_gradient_x(M_l, U, Vt, l, x, m, n):
    '''
    Computes the gradient of the loss function with respect to x
    
    @param M: np array of shape (m, n), input matrix
    @param U: np array of shape (m, rank), left singular column vectors
    @param Vt: np array of shape (rank, n), right singular column vectors
    @param l: current rank being computed
    @param x: np array of shape (m,1)
    @param m: # of rows in M
    @param n: # of cols in M
    @return: np array of shape (m,1)
    '''
    compute_MT_x(n, M_l, x, Vt[l:l+1,:])
    compute_M_xT(m, M_l, Vt[l:l+1,:], U[:,l:l+1])
    return (np.linalg.norm(x)**2)*x - U[:,l:l+1]

##########################
##  SVD via Gradient Descent (algorithm 4)
##########################
def gd_svd(M,
            rank = 1,
            eps = 1e-15,
            eta = 1/2,
            record_conv = None,
            U_true = None,
            S_true = None,
            Vt_true = None,
            seed = 42,
            ):
    '''
    Computes rank-SVD via gradient descent
    
    @param M: np array of shape (m, n), matrix that is being approximated with singular values
    @param rank: int, rank of the approximation
    @param max_iters: int, max number of steps to run algorithm for)
    @param eps: float, tolerance for stopping condition
    @param record_conv: bool, whether to record convergence metrics (slows down computation)
    @param U_true: np array of shape (m, rank), true left singular vectors
    @param S_true: np array of shape (rank,), true singular values
    @param Vt_true: np array of shape (rank, n), true right singular vectors

    @return: U (m, rank), S (rank,), Vt (rank, n), record (list of tab separated strings, first line is convergence iters, second is header)
    '''
    if record_conv:
        assert U_true is not None
        assert S_true is not None
        assert Vt_true is not None

    rng = np.random.default_rng(seed)

    # initialize U, S, Vt, x
    m,n = M.shape
    U = np.zeros(shape=(m, rank))
    S = np.ones(rank)
    Vt = rng.standard_normal((rank,n))

    if record_conv:
        with open(record_conv,'w') as f:
            f.write("rank\titer\tMerr\tSerr\tx_norm\tUerr\tU_ortho\teps_1\tVerr\tV_ortho\teps_2\n")
    record = [np.zeros(rank, dtype = int)]

    # initialize M_l
    M_l = np.memmap("_svd_M_l_temp.dat", dtype=np.float64, mode='w+', shape=(m,n))
    np.copyto(M_l, M)

    for l in range(rank):
        # initialize Vt, x, U
        Vt[l:l+1,:] /= np.linalg.norm(Vt[l:l+1,:])
        x = M_l @ Vt[l:l+1,:].T
        U[:,l:l+1] = x / np.linalg.norm(x)

        # initialize convergence conditions
        stop1 = 1
        stop2 = 1

        if record_conv:
            with open(record_conv,'a') as f:
                f.write(f"{l+1}\t{record[0][l]}\t{M_err(M, U, S, Vt)}\t{S_err(S, S_true, l+1)}\t{np.linalg.norm(x)}\t{U_err(U, U_true, l+1)}\t{U_ortho(U,l+1)}\t{stop1}\t{Vt_err(Vt, Vt_true, l+1)}\t{Vt_ortho(Vt,l+1)}\t{stop2}\n")

        while ((stop1 > eps) or (stop2 > eps)) and (np.linalg.norm(x) > 10*eps*rank):
            # compute gradient of x
            stop2 = S[l].copy()
            stop1 = U[:,l:l+1].copy()
            
            grad_x = (np.linalg.norm(x)**2)*x - M_l @ (M_l.T @ x)

            # step x
            x -= eta/(np.linalg.norm(x)**2) * grad_x

            # update singular value/vectors
            S[l] = np.linalg.norm(x)
            U[:,l:l+1] = x/np.linalg.norm(x)
            # compute update for vT
            Vt[l:l+1,:] = U[:,l:l+1].T @ M
            Vt[l:l+1,:] /= S[l]

            # compute stopping conditions
            stop1 = np.linalg.norm(U[:,l:l+1] - stop1)
            stop2 = abs(S[l] - stop2)

            record[0][l] +=1
            if record_conv:
                with open(record_conv,'a') as f:
                    f.write(f"{l+1}\t{record[0][l]}\t{M_err(M, U, S, Vt)}\t{S_err(S, S_true, l+1)}\t{np.linalg.norm(x)}\t{U_err(U, U_true, l+1)}\t{U_ortho(U,l+1)}\t{stop1}\t{Vt_err(Vt, Vt_true, l+1)}\t{Vt_ortho(Vt,l+1)}\t{stop2}\n")

        M_l -= S[l] * U[:,l:l+1] @ Vt[l:l+1,:]
        M_l.flush()
    
    # remove temporary file
    os.remove("_svd_M_l_temp.dat")
    return U,S,Vt


##########################
##  SVD via Gradient Descent (algorithm 3)
##########################
def gd_svd3(M,
            rank = 1,
            eps = 1e-15,
            eta = 1/2,
            record_conv = None,
            U_true = None,
            S_true = None,
            Vt_true = None,
            seed = 42,
            ):
    '''
    Computes rank-SVD via gradient descent
    
    @param M: np array of shape (m, n), matrix that is being approximated with singular values
    @param rank: int, rank of the approximation
    @param max_iters: int, max number of steps to run algorithm for)
    @param eps: float, tolerance for stopping condition
    @param record_conv: bool, whether to record convergence metrics (slows down computation)
    @param U_true: np array of shape (m, rank), true left singular vectors
    @param S_true: np array of shape (rank,), true singular values
    @param Vt_true: np array of shape (rank, n), true right singular vectors

    @return: U (m, rank), S (rank,), Vt (rank, n), record (list of tab separated strings, first line is convergence iters, second is header)
    '''
    if record_conv:
        assert U_true is not None
        assert S_true is not None
        assert Vt_true is not None

    rng = np.random.default_rng(seed)

    # initialize U, S, Vt, x
    m,n = M.shape
    U = np.zeros(shape=(m, rank))
    S = np.ones(rank)
    Vt = rng.standard_normal((rank,n))

    if record_conv:
        with open(record_conv,'w') as f:
            f.write("rank\titer\tMerr\tSerr\tstop\tUerr\tU_ortho\tx_norm\tVerr\tV_ortho\tgrad_x_norm\n")
    record = [np.zeros(rank, dtype = int)]

    # initialize M_l
    M_l = np.memmap("_svd_M_l_temp.dat", dtype=np.float64, mode='w+', shape=(m,n))
    np.copyto(M_l, M)

    for l in range(rank):
        # normalize Vt
        Vt[l:l+1,:] /= np.linalg.norm(Vt[l:l+1,:])

        # initialize x, grad_x
        x = M_l @ Vt[l:l+1,:].T
        grad_x = np.ones((m,1))

        # initialize convergence conditions
        stop = 1

        if record_conv:
            with open(record_conv,'a') as f:
                f.write(f"{l+1}\t{record[0][l]}\t{M_err(M, U, S, Vt)}\t{S_err(S, S_true, l+1)}\t{stop}\t{U_err(U, U_true, l+1)}\t{U_ortho(U, l+1)}\t{np.linalg.norm(x)}\t{Vt_err(Vt, Vt_true, l+1)}\t{Vt_ortho(Vt, l+1)}\t{np.linalg.norm(grad_x)}\n")

        while (stop > eps) and (np.linalg.norm(x) > 10*eps*rank):
            # compute gradient of x
            grad_x = (np.linalg.norm(x)**2)*x - M_l @ (M_l.T @ x)
            stop = np.linalg.norm(grad_x)/(np.linalg.norm(x)**3)

            # step x
            x -= eta/(np.linalg.norm(x)**2) * grad_x

            # update singular value/vectors
            S[l] = np.linalg.norm(x)
            U[:,l:l+1] = x/np.linalg.norm(x)
            # compute update for vT
            Vt[l:l+1,:] = U[:,l:l+1].T @ M
            Vt[l:l+1,:] /= S[l]

            record[0][l] +=1
            if record_conv:
                with open(record_conv,'a') as f:
                    f.write(f"{l+1}\t{record[0][l]}\t{M_err(M, U, S, Vt)}\t{S_err(S, S_true, l+1)}\t{stop}\t{U_err(U, U_true, l+1)}\t{U_ortho(U, l+1)}\t{np.linalg.norm(x)}\t{Vt_err(Vt, Vt_true, l+1)}\t{Vt_ortho(Vt, l+1)}\t{np.linalg.norm(grad_x)}\n")

        M_l -= S[l] * U[:,l:l+1] @ Vt[l:l+1,:]
        M_l.flush()
    
    # remove temporary file
    os.remove("_svd_M_l_temp.dat")
    return U,S,Vt
    
##########################
##    Power Method
##########################
def power_method(M,
                rank = 1,
                eps = 1e-15,
                record_conv = None,
                U_true = None,
                S_true = None,
                Vt_true = None,
                seed = 42,
                ):
    '''
    Computes rank-SVD via power method
    
    @param M: np array of shape (m, n), matrix that is being approximated with singular values
    @param rank: int, rank of the approximation
    @param max_iters: int, max number of steps to run algorithm fol+1)
    @param eps: float, tolerance for stopping condition
    @param record_conv: bool, whether to record convergence metrics (slows down computation)
    @param U_true: np array of shape (m, rank), true left singular vectors
    @param S_true: np array of shape (rank,), true singular values
    @param Vt_true: np array of shape (rank, n), true right singular vectors

    @return: U (m, rank), S (rank,), Vt (rank, n), record (list of tab separated strings, first line is convergence iters, second is header)
    '''
    if record_conv:
        assert U_true is not None
        assert S_true is not None
        assert Vt_true is not None

    rng = np.random.default_rng(seed)
        
    # initialize U, S, Vt, x
    m,n = M.shape
    S = np.ones(rank)
    U = np.zeros((m, rank))
    Vt = np.zeros((rank, n))

    if record_conv:
        with open(record_conv,'w') as f:
            f.write("rank\titer\tMerr\tSerr\tx_norm\tUerr\tU_ortho\teps_1\tVerr\tV_ortho\teps_2\n")
    record = [np.zeros(rank, dtype = int)]
    
    # initialize M_l
    M_l = np.memmap("_svd_M_l_temp.dat", dtype=np.float64, mode='w+', shape=(m,n))
    np.copyto(M_l, M)

    for l in range(rank):
        # initialize x
        x = rng.standard_normal((m, 1))
        x /= np.linalg.norm(x)
        x = M_l @ (M_l.T @ x)

        U[:,l:l+1] = x / np.linalg.norm(x)
        
        # initialize convergence conditions
        stop1 = 1
        stop2 = 1

        if record_conv:
            with open(record_conv,'a') as f:
                f.write(f"{l+1}\t{record[0][l]}\t{M_err(M, U, S, Vt)}\t{S_err(S, S_true, l+1)}\t{np.linalg.norm(x)}\t{U_err(U, U_true, l+1)}\t{U_ortho(U,l+1)}\t{stop1}\t{Vt_err(Vt, Vt_true, l+1)}\t{Vt_ortho(Vt,l+1)}\t{stop2}\n")

        while ((stop1 > eps) or (stop2 > eps)) and (math.sqrt(np.linalg.norm(x)) > 10*eps*rank):
            stop2 = S[l].copy()
            stop1 = U[:,l:l+1].copy()
            
            # compute update for x
            x = (M_l @ (M_l.T @ x)) / np.linalg.norm(x)

            # update singular value/vectors
            S[l] = math.sqrt(np.linalg.norm(x))
            U[:,l:l+1] = x / np.linalg.norm(x)
            # compute update for vT
            Vt[l:l+1,:] = U[:,l:l+1].T @ M
            Vt[l:l+1,:] /= S[l]

            # compute stopping conditions
            stop1 = np.linalg.norm(U[:,l:l+1] - stop1)
            stop2 = abs(S[l] - stop2)

            record[0][l] +=1

            if record_conv:
                with open(record_conv,'a') as f:
                    f.write(f"{l+1}\t{record[0][l]}\t{M_err(M, U, S, Vt)}\t{S_err(S, S_true, l+1)}\t{np.linalg.norm(x)}\t{U_err(U, U_true, l+1)}\t{U_ortho(U,l+1)}\t{stop1}\t{Vt_err(Vt, Vt_true, l+1)}\t{Vt_ortho(Vt,l+1)}\t{stop2}\n")

        M_l -= S[l] * U[:,l:l+1] @ Vt[l:l+1,:]
        M_l.flush()
    
    os.remove("_svd_M_l_temp.dat")
    return U,S,Vt
