# import numpy as np 
import CTL.funcs.xplib as xplib
import CTL.funcs.funcs as funcs

def SVDDecomposition(a, chi, returnSV = False, errorOrder = 2):
    """
    SVD decomposition of a given matrix, keep at most chi singular values. Singular values are divided into square roots and go to each side.

    Parameters
    ----------
    a : 2-D ndarray of float
        The matrix to decompose, shape (n, m).
    chi : int
        The maximum number of non-zero singular values to be kept.
    returnSV : bool, default False
        If True, then return two isometries with singular values, otherwise(by default) return two matrices, product of which gives an approximation of original matrix.
    errorOrder : int, default 2
        The order of error in singular values. The error will be calculated as 1 - (s[:chi] ** errorOrder).sum() / (s ** errorOrder).sum().

    Returns
    uRes : 2-D ndarray of float
        Left part of the decomposition, shape (n, k), where k is generated by n, m, k, and the number of non-zero singular values.
    s : 1-D ndarray of float(only when returnSV is True)
        The singular values, shape (k,)
    vRes : 2-D ndarray of float
        Right part of the decomposition, shape (k, m). a approximately equal to uRes @ vRes.
    error : float
        The loss of truncating singular values: sum(square of truncated values) / sum(square of all values).
    """

    u, s, vh = xplib.xp.linalg.svd(a)
    
    sLen = len(s)
    if (chi >= sLen):
        error = 0.0
        chi = sLen
    else:
        error = xplib.xp.sum(s[chi:] ** errorOrder) / xplib.xp.sum(s ** errorOrder)

    # print('error = {}'.format(error))

    u = u[:, :chi]
    s = s[:chi]
    vh = vh[:chi]

    if (returnSV):
        return u, s, vh, error

    sqrtSMat = xplib.xp.diag(xplib.xp.sqrt(s))

    uRes = u @ sqrtSMat
    vRes = funcs.transposeConjugate(sqrtSMat @ vh)
    return uRes, vRes, error





