#!/usr/bin/env python3
# !!!
import os, sys
# !!!
#QP_PATH=os.environ["QMCCHEM_PATH"]
#sys.path.insert(0,QMCCHEM_PATH+"/EZFIO/Python/")
# !!!
from ezfio import ezfio
from datetime import datetime
import numpy as np
from scipy.sparse.linalg import svds
from R3SVD_LiYu  	       import R3SVD_LiYu
from RSVD                import powit_RSVD  
from R3SVD_AMMAR         import R3SVD_AMMAR
import time 
# !!!
fmt = '%5d' + 2 * ' %e'
# !!!
if __name__ == "__main__":
    # !!!
    if len(sys.argv) != 2:
        print("Usage: %s <EZFIO_DIRECTORY>"%sys.argv[0])
        sys.exit(1)
    filename = sys.argv[1]
    ezfio.set_file(filename)
    # !!!
    N_det  = ezfio.get_spindeterminants_n_det()
    A_rows = np.array(ezfio.get_spindeterminants_psi_coef_matrix_rows())
    A_cols = np.array(ezfio.get_spindeterminants_psi_coef_matrix_columns())
    A_vals = np.array(ezfio.get_spindeterminants_psi_coef_matrix_values())
    nrows, ncols  = ezfio.get_spindeterminants_n_det_alpha(), ezfio.get_spindeterminants_n_det_beta()
    Y = np.zeros( (nrows, ncols) )
    for k in range(N_det):
        i = A_rows[k] - 1
        j = A_cols[k] - 1
        Y[i,j] = A_vals[0][k]
    print("# # # # # # # # # # # # # # # # # # # # # #")
    print('matrix dimensions = {} x {}'.format(nrows, ncols))
    print("# # # # # # # # # # # # # # # # # # # # # # \n")
    normY = np.linalg.norm(Y, ord='fro')
    print( normY )
    # !!!
    print('Full SVD:')
    t_beg = time.time()
    U, S_FSVD, VT = np.linalg.svd(Y, full_matrices=0)
    t_end         = time.time()
    rank          = S_FSVD.shape[0]
    energy        = np.sum(np.square(S_FSVD)) / normY**2
    err_SVD       = 100. * np.linalg.norm(Y - np.dot(U,np.dot(np.diag(S_FSVD),VT)), ord='fro') / normY
    print('rank = {}, energy = {}, error = {}%, CPU time = {} \n'.format(rank, energy, err_SVD, t_end-t_beg))
    # !!!
    np.savetxt('results_python/h2o_pseudo/SingValues_FullSVD.txt', np.transpose([ np.array(range(rank))+1, S_FSVD  ]), fmt='%5d' + ' %e', delimiter=' ')
    # !!!
    t       = 50
    delta_t = 10
    npow    = 15
    err_thr = 1e-3
    maxit   = 10
    # !!!
    print('RRR SVD Li & Yu:')
    t_beg          = time.time()
    U, S_R3SVD, VT = R3SVD_LiYu(Y, t, delta_t, npow, err_thr, maxit)
    t_end          = time.time()
    rank           = S_R3SVD.shape[0]
    energy         = np.sum( np.square(S_R3SVD) ) / normY**2
    err_SVD        = 100. * np.linalg.norm(Y - np.dot(U,np.dot(np.diag(S_R3SVD),VT)), ord='fro') / normY
    print('nb Pow It = {}, rank = {}, energy = {}, error = {} %, CPU time = {}\n'.format(npow, rank, energy, err_SVD, t_end-t_beg))
    # !!!
    err_R3SVD = np.zeros( (rank) )
    for i in range(rank):
        err_R3SVD[i] = 100.0 * abs( S_FSVD[i] - S_R3SVD[i] ) / S_FSVD[i]
    np.savetxt('results_python/h2o_pseudo/SingValues_R3SVD.txt', np.transpose([ np.array(range(rank))+1, S_R3SVD, err_R3SVD ]), fmt=fmt, delimiter=' ')
    # !!!
    nb_oversamp = 10
    tol_SVD     = 1e-10
    print('RRR SVD my version:')
    t_beg          = time.time()
    U, S_MRSVD, VT = R3SVD_AMMAR(Y, t, delta_t, npow, nb_oversamp, err_thr, maxit, tol_SVD)
    t_end          = time.time()
    rank           = S_MRSVD.shape[0]
    energy         = np.sum( np.square(S_MRSVD) ) / normY**2
    err_SVD        = 100. * np.linalg.norm(Y - np.dot(U,np.dot(np.diag(S_MRSVD),VT)), ord='fro') / normY
    print('nb Pow It = {}, rank = {}, energy = {}, error = {} %, CPU time = {}\n'.format(npow, rank, energy, err_SVD, t_end-t_beg))
    # !!!
    err_MRSVD = np.zeros( (rank) )
    for i in range(rank):
        err_MRSVD[i] = 100.0 * abs( S_FSVD[i] - S_MRSVD[i] ) / S_FSVD[i]
    np.savetxt('results_python/h2o_pseudo/SingValues_MRSVD.txt', np.transpose([ np.array(range(rank))+1, S_MRSVD, err_MRSVD ]), fmt=fmt, delimiter=' ')
    # !!!
    trank = rank
    print("Truncated RSVD (pre-fixed rank = {} & oversampling parameter = {}):".format(trank,nb_oversamp))
    t_beg         = time.time()
    U, S_RSVD, VT = powit_RSVD(Y, trank, npow, nb_oversamp)
    t_end         = time.time()
    rank          = S_RSVD.shape[0]
    energy        = np.sum( np.square(S_RSVD) ) / normY**2
    err_SVD       = 100. * np.linalg.norm( Y - np.dot(U,np.dot(np.diag(S_RSVD),VT)), ord="fro") / normY
    print('nb Pow It = {}, rank = {}, energy = {}, error = {} %, CPU time = {}\n'.format(npow, rank, energy, err_SVD, t_end-t_beg))
    # !!!
    err_RSVD = np.zeros( (rank) )
    for i in range(rank):
        err_RSVD[i] = 100.0 * abs( S_FSVD[i] - S_RSVD[i] ) / S_FSVD[i]
    np.savetxt('results_python/h2o_pseudo/SingValues_RSVD.txt', np.transpose([ np.array(range(rank))+1, S_RSVD, err_RSVD ]), fmt=fmt, delimiter=' ')
    # !!!
    print("Truncated SVD (scipy):")
    t_beg         = time.time()
    U, S_TSVD, VT = svds(Y, min(trank, min(Y.shape[0],Y.shape[1])-1 ), which='LM')
    t_end         = time.time()
    rank          = S_TSVD.shape[0]
    energy        = np.sum( np.square(S_TSVD) ) / normY**2
    err_SVD       = 100. * np.linalg.norm( Y - np.dot(U, np.dot(np.diag(S_TSVD),VT) ), ord="fro") / normY
    print('rank = {}, energy = {}, error = {} %, CPU time = {}\n'.format(rank, energy, err_SVD, t_end-t_beg))
    # !!!
    err_TSVD = np.zeros( (rank) )
    for i in range(rank-1):
        for j in range(i+1,rank):
            if( S_TSVD[j] > S_TSVD[i]):
                S_TSVD[i], S_TSVD[j] = S_TSVD[j], S_TSVD[i]
    for i in range(rank):
        err_TSVD[i] = 100.0 * abs( S_FSVD[i] - S_TSVD[i] ) / S_FSVD[i]
    np.savetxt('results_python/h2o_pseudo/SingValues_TSVD.txt', np.transpose([ np.array(range(rank))+1, S_TSVD, err_TSVD ]), fmt=fmt, delimiter=' ')
    # !!!
# !!!