import numpy as np
import dislib as ds
from pycompss.api.task import task
from pycompss.api.constraint import constraint
from pycompss.api.reduction import reduction
from dislib.data.array import Array
from pycompss.api.parameter import COLLECTION_IN, COLLECTION_OUT, COLLECTION_INOUT, Depth, Type
from pycompss.api.api import compss_wait_on, compss_barrier
from time import time
from dislib.decomposition import tsqr
from random_svd import make_low_rank_matrix
import math
from rechunk_shared import rechunk_shared


@constraint(computing_units="48")
@task(
        Ur_blocks={Type: COLLECTION_OUT, Depth: 2},
        S_blocks={Type: COLLECTION_OUT, Depth: 2},
        Vr_blocks={Type: COLLECTION_OUT, Depth: 2},
    )
def _my_svd(r_block, Ur_blocks, S_blocks, Vr_blocks):
    u,s,v = np.linalg.svd(r_block, full_matrices=False)
    s = np.diag(s)
    v = v.T
    
    for i, row in enumerate(np.array_split(u, len(Ur_blocks))):
        for j, val in enumerate(np.array_split(row, len(Ur_blocks[i]), axis=1)):
            Ur_blocks[i][j] = val

    for i, row in enumerate(np.array_split(s, len(S_blocks))):
        for j, val in enumerate(np.array_split(row, len(S_blocks[i]), axis=1)):
            S_blocks[i][j] = val

    for i, row in enumerate(np.array_split(v, len(Vr_blocks))):
        for j, val in enumerate(np.array_split(row, len(Vr_blocks[i]), axis=1)):
            Vr_blocks[i][j] = val


def check_convergeces(res, tol, nsv):
    res = np.block(res).flatten()
    converged = res < tol
    return np.all(converged[:nsv]), max(res[:nsv])


@task(returns=1)
def copy_block(block):
    return block


@constraint(computing_units="48")
@task(B_blocks={Type: COLLECTION_INOUT, Depth: 2}, returns=1)
def R_ortho(B_blocks, R, k , b, B_reg_shape):
    B = np.block(B_blocks)
    R = R.astype(np.float64)
    
    for j in range(0, k, b):
        w = B[:,j:j+b]
        if j > 0:
            h = B[:,0:j].T @ w
            w = w - B[:,0:j] @ h
        w, r1 = np.linalg.qr(w)

        if j > 0:
            h2 = B[:,0:j].T @ w
            w = w - B[:,0:j] @ h2
            R[0:j,j:j+b][:] = h + h2
        w, r2 = np.linalg.qr(w)

        B[:,j:j+b] = w
        R[j:j+b,j:j+b][:] = r1 @ r2
    
    n, m = B_reg_shape
    for i in range(len(B_blocks)):
        for j in range(len(B_blocks[i])):
            B_blocks[i][j] = B[i*n:(i+1)*n, j*m:(j+1)*m]
    
    return R


def _get_all_cols(A, i):
    blocks = [[] for j in range(A._n_blocks[0])]
    shape = [A.shape[0], 0]
    for col_idx in range(i+1):
        col_shape = A._get_col_shape(col_idx)
        shape[1] += col_shape[1]
        for j in range(A._n_blocks[0]):
            blocks[j].append(A._blocks[j][i])
    return Array(blocks=blocks,
                    top_left_shape=A._reg_shape,
                    reg_shape=A._reg_shape, shape=col_shape,
                    sparse=A._sparse, delete=False)

@task(returns=1)
def _R_expand_task(R_block, b):
    return np.pad(R_block, (0, b))

@task(returns=1)
def _norm(b):
    return np.linalg.norm(b, axis=0)

@reduction(chunk_size="8")
@task(blocks={Type: COLLECTION_IN, Depth: 1}, returns=1)
def _norm_red(blocks):
    return np.linalg.norm(blocks, axis=0)

def my_norm(A):
    norm_blocks = Array._get_out_blocks(A._n_blocks)
    for i in range(A._n_blocks[0]):
        for j in range(A._n_blocks[1]):
            norm_blocks[i][j] = _norm(A._blocks[i][j])

    norms = []
    for j in range(A._n_blocks[1]):
        col = [norm_blocks[i][j] for i in range(A._n_blocks[0])]
        if len(col) > 1:
            norms.append(_norm_red(col))
        else:
            norms.append(col[0])
    return norms


def R_expand(R, b):
    R._blocks[0][0] = _R_expand_task(R._blocks[0][0], b)
    R._shape = (R.shape[0] + b, R.shape[1] + b)
    R._reg_shape = (R.shape[0] + b, R.shape[1] + b)


def nsv_tolerance(m, n, nsv, S):
	min_sval = S[nsv-1]
	total_sval = min(m, n)
	sum_nsv = sum([s**2 for s in S])
	sum_remaining = (total_sval-nsv) * min_sval**2
	tol = np.sqrt(sum_remaining / (sum_remaining + sum_nsv))
	return tol
    

def _svd_random(A, b, k, max_restarts, nsv, tol, verbose=False):
    # A: matrix
    # b: block size
    # k: number of vectors (must be multiple of b)
    # max_restarts: maximum number of restarts
    # nsv: number of desired singular values
    # tol: consider a singular converged when its residual is less than tol
    m, n = A.shape
    
    B = ds.random_array((n, k), (b, b), random_state=0)
    R = ds.zeros((k, k), (k, k))
    Q = ds.matmul(A, B)

    for iters in range(max_restarts):
        for j in range(0, k, b):
            col_idx = j // b
            q = Q._get_col_block(col_idx)
            if j > 0:
                h = ds.matmul(_get_all_cols(Q, col_idx), q, transpose_a=True)
                q = ds.data.matsubtract(q, ds.matmul(_get_all_cols(Q, col_idx), h))
                h = ds.matmul(_get_all_cols(Q, col_idx), q, transpose_a=True)
                q = ds.data.matsubtract(q, ds.matmul(_get_all_cols(Q, col_idx), h))

            q, _ = tsqr(q, mode='reduced')
            q._delete = False
            
            for i in range(Q._n_blocks[0]):
                Q._blocks[i][col_idx] = q._blocks[i][0] # copy_block(q._blocks[i][0])
            

        B = ds.matmul(A, Q, transpose_a=True)

        R._blocks[0][0] = R_ortho(B._blocks, R._blocks[0][0], k, b, B._reg_shape)
        
        Ur_blocks = Array._get_out_blocks((k//b, k//b))
        Vr_blocks = Array._get_out_blocks((k//b, k//b))
        S_blocks = Array._get_out_blocks((k//b, k//b))
        _my_svd(R._blocks[0][0], Ur_blocks, S_blocks, Vr_blocks)

        Ur = Array(Ur_blocks, (b, b), (b, b), (R.shape[0], R.shape[0]), False)
        S = Array(S_blocks, (b, b), (b, b), (R.shape[1], R.shape[1]), False)
        Vr = Array(Vr_blocks, (b, b), (b, b), (R.shape[1], R.shape[1]), False)

        U = ds.matmul(Q, Vr)
        V = ds.matmul(B, Ur)

        Q = ds.matmul(A, V)

        mmm = ds.data.matsubtract(Q, ds.matmul(U, S))

        res = np.block(compss_wait_on(my_norm(mmm))) #mmm.norm().collect()
        converged, conv_tol = check_convergeces(res, tol, nsv)
        if converged:
            if verbose:
                print("Converged in {} iterations to {} precision.".format(iters+1, conv_tol))
            break
        else:
            if verbose:
                print("Not converged in {} iterations. Current precision: {}.".format(iters+1, conv_tol))

    return U, S, V, conv_tol


def random_svd(A, iters, epsilon, tol, nsv=None, verbose=False):
    m, n = A.shape
    b = A._reg_shape[1]

    if nsv is None:
        nsv = np.argmin((np.logspace(0, -20, n) - epsilon)**2)

    if b == nsv:
        k = b
    elif b > nsv:
        k = b
    else:
        k = math.ceil(nsv / b) * b


    if verbose:
        print(f'Random-SVD: Starting computing {nsv} singular values using {k} vectors.')

    epsilon_conv = False
    while not epsilon_conv:
        U, S, V, conv_tol = _svd_random(A, b, k, iters, nsv, tol, verbose=verbose)

        if conv_tol > tol:
            print(f'WARNING: Random-SVD did not converge in {iters} iterations to {tol} precision.')
            print(f'         Maximum accuracy was {conv_tol}')
        
        S_flat = np.diag(S.collect())
        epsilon_hat = nsv_tolerance(m, n, nsv, S_flat)
        if epsilon_hat <= epsilon:
            epsilon_conv = True
        else:
            k += b
            nsv += b
            if verbose:
                print(f'NOTE: Random-SVD singular value criteria {epsilon_hat} > {epsilon}.')
                print(f'      Increasing number of vectors to {k} and {nsv} singular values.')

    if verbose:
        print(f'Random-SVD: Converged to {epsilon_hat} singular value tolerance.')
        
    return U, S, V.T, nsv

@task(returns=1, is_distributed=True)
def load_npy(file_path):
    return np.load(file_path)


def load_dataset(wildcard):
    from glob import glob
    import re, os

    files = glob(wildcard)
    files = sorted(files, key=lambda x: int(re.findall(r'\d+', os.path.basename(x))[0]))
    # files = sorted(files, key=lambda x: int(''.join(re.findall(r'\d+', os.path.basename(x)))))

    blocks = [[]]
    shape = [0, 0]
    bs = None

    for f in files:
        with open(files[0], "rb") as fid:
            version = np.lib.format.read_magic(fid)
            shape_block, _, _ = np.lib.format._read_array_header(fid, version)
            if not bs:
                bs = shape_block
            if shape[0] == 0:
                shape[0] = shape_block[0]
            shape[1] += shape_block[1]
            # blocks.append([load_npy(f)])
            blocks[0].append(load_npy(f))

    return Array(blocks, bs, bs, shape, False)

if __name__ == "__main__":
    # A = ds.random_array((3_600_000, 1200), (75_000, 500))
    
    A = load_dataset('/gpfs/projects/bsc19/bsc19756/MakeTrainingMatricesSerial/snapshots_at_step_*')
    A = A.rechunk((75_000, 500))
    # A = load_dataset('/gpfs/projects/bsc19/bsc19756/Matrix_Sebastian/*')
    # A = A.rechunk((40_000, 1450))

    print(A.shape, A._reg_shape)

    compss_barrier()
    t1 = time()

    U, S, V, nsv = random_svd(A, iters=1, nsv=456, epsilon=100, tol=1e-3, verbose=True)

    compss_barrier()

    duration = time()-t1
    print(f'Time R-SVD:', duration)
    print(np.diag(S.collect())[:nsv])
