#!/usr/bin/env python3

import os, sys

#QP_PATH=os.environ["QMCCHEM_PATH"]
#sys.path.insert(0,QMCCHEM_PATH+"/EZFIO/Python/")

import scipy
from   scipy import linalg

from   ezfio import ezfio
from   datetime import datetime
import numpy as np
import time 




# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
def get_Aref():

    Aref = np.zeros( (n_alpha, n_beta) )
    for k in range(n_det):
        i         = A_rows[k] - 1
        j         = A_cols[k] - 1
        Aref[i,j] = A_vals[0][k]

    return( Aref )
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
def powit_RSVD(X, n_TSVD, nb_powit, nb_oversamp):

    print(" --- begin powit_RSVD --- ")
    print(" n_TSVD = {}".format(n_TSVD))
    print(" pow it = {} & nb oversampling = {}".
          format(nb_powit,nb_oversamp))

    G  = np.random.randn(X.shape[1], n_TSVD+nb_oversamp)
    Q  = QR_fact( np.dot(X,G) )

    for i in range(nb_powit):
        ti = time.time()
        print(" start pow it = {}".format(i))

        Q = QR_fact( np.dot(X.T,Q) )
        Q = QR_fact( np.dot(X,Q) )

        tf = time.time()
        dt = (tf-ti)/60.
        print(" end pow it = {} after {} min".format(i,dt))

    Y = np.dot(Q.T,X)

    U, S, VT = np.linalg.svd(Y, full_matrices=1)
    U = np.dot(Q,U)

    print( " --- end powit_RSVD --- \n")
    return U, S, VT
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
def QR_fact(X):

    Q, _ = linalg.qr(X, mode="full")
    #Q,R      = np.linalg.qr(X, mode="complete")
    #D        = np.diag( np.sign( np.diag(R) ) )
    Qunique  = Q #np.dot(Q,D)
    #Runique = np.dot(D,R)

    return(Qunique)
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
def TSVD_save_EZFIO():

    U_toEZFIO = np.zeros( ( 1, U.shape[1], U.shape[0] ) )
    V_toEZFIO = np.zeros( ( 1, V.shape[1], V.shape[0] ) )
    U_toEZFIO[0,:,:] = U_TSVD.T
    V_toEZFIO[0,:,:] = V_TSVD.T

    ezfio.set_spindeterminants_n_svd_coefs( n_TSVD )
    ezfio.set_spindeterminants_psi_svd_alpha( U_toEZFIO )
    ezfio.set_spindeterminants_psi_svd_beta ( V_toEZFIO )
    ezfio.set_spindeterminants_psi_svd_coefs( S_RSVD )

    print(' SVD vectors & coeff are saved to EZFIO ')
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -




if __name__ == "__main__":


    print("")
    print(" Today's date:", datetime.now() )

    # EZFIO file
    #EZFIO_file = "/home/aammar/qp2/src/svdwf/h2o_work/FN_test/cc_pCVDZ/h2o_dz"
    EZFIO_file = "/home/aammar/qp2/src/svdwf/h2o_work/FN_test/cipsi_calcul/h2o_dz_fci"
    ezfio.set_file(EZFIO_file)
    print(" EZFIO = {}\n".format(EZFIO_file))

    #read_wf = True
    #ezfio.read_wf = True
    #TOUCH read_wf


    n_det = ezfio.get_spindeterminants_n_det()
    print(' n_det = {}'.format(n_det))

    n_alpha = ezfio.get_spindeterminants_n_det_alpha()
    n_beta  = ezfio.get_spindeterminants_n_det_beta()
    print(' matrix dimensions = {} x {} = {} \n'.format(n_alpha, n_beta, n_alpha*n_beta))

    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())
    Aref   = get_Aref()
    A_norm = np.linalg.norm(Aref, ord='fro')


    npow        = 15
    nb_oversamp = 10
    n_TSVD      = 100 #min(n_alpha,n_beta)

    t_beg = time.time()
    U, S_RSVD, Vt = powit_RSVD(Aref, n_TSVD, npow, nb_oversamp)
    print(' powit_RSVD time = {}\n'.format((time.time()-t_beg)/60.))
    
    S_mat = np.zeros((n_alpha,n_beta))
    for i in range(n_TSVD):
        S_mat[i,i] = S_RSVD[i]
    err_SVD = 100. * np.linalg.norm( Aref - np.dot(U,np.dot(S_mat,Vt)), ord="fro") / A_norm
    print(' powit_RSVD error (%) = {} \n'.format(err_SVD))
#______________________________________________________________________________________________________________________