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

from ezfio import ezfio
from datetime import datetime
import time
import numpy as np
import subprocess
from scipy.linalg import eig, eigh


# ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ !
def get_energy():
    buffer = subprocess.check_output( ['qmcchem', 'result', '-e', 'e_loc', EZFIO_file]
                                      , encoding='UTF-8' )
    if buffer.strip() != "":
        buffer = buffer.splitlines()[-1]
        _, energy, error = [float(x) for x in buffer.split()]
        return energy, error
    else:
      return None, None
# ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ !


# ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ !
def run_qmc():
    return subprocess.check_output(['qmcchem', 'run', EZFIO_file])

def stop_qmc():
    subprocess.check_output(['qmcchem', 'stop', EZFIO_file])

def set_vmc_params():
    subprocess.check_output(['qmcchem', 'edit', '-c', '-j', 'None'
                                      , '-m', 'SRMC'
                                      , '-s', 'Brownian'
                                      , '-l', str(block_time)
                                      , '--time-step=0.0005'
                                      , '-t', str(total_time)
                                      , EZFIO_file])
# ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ !


# ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ !
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 check_svd_error( A, U, S, Vt ):
    norm_A  = np.linalg.norm(A, ord='fro')
    A_SVD   = np.linalg.multi_dot([ U, np.diag(S), Vt ])
    err_SVD = 100. * np.linalg.norm( A - A_SVD, ord="fro") / norm_A
    print(' error between A_SVD and Aref = {} %\n'.format(err_SVD) )
# ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ !


# ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ !
def TSVD_save_EZFIO():

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

    ezfio.set_spindeterminants_n_svd_coefs( low_rank )
    ezfio.set_spindeterminants_psi_svd_alpha( U_toEZFIO )
    ezfio.set_spindeterminants_psi_svd_beta ( V_toEZFIO )
    ezfio.set_spindeterminants_psi_svd_coefs( S_TSVD )

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




if __name__ == '__main__':

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

    # EZFIO file
    EZFIO_file = "/home/aammar/qp2/src/svdwf/h2o_work/FN_test/cc_pCVDZ/h2o_dz_nsvd"
    ezfio.set_file(EZFIO_file)
    print(" filename = {}".format(EZFIO_file))

    # parameters
    energ_nuc = 9.194966082434476

    # get spindeterminant data from EZFIO
    n_alpha = ezfio.get_spindeterminants_n_det_alpha()
    n_beta  = ezfio.get_spindeterminants_n_det_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() )
    n_det   = A_rows.shape[0]
    print(' ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~')
    print(' matrix: {} x {}'.format(n_alpha,n_beta))
    print(' n_det = {}'.format(n_det))
    print(' ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~')
    # construct intial dense matrix
    Aref = get_Aref()
    # perform initial Full SVD
    print(' ----- Performing Full SVD ----- ')
    U_FSVD, S_FSVD, Vt_FSVD  = np.linalg.svd(Aref, full_matrices=False)
    V_FSVD = Vt_FSVD.T
    #print(' initial psi coeff = {}\n'.format(S_FSVD))
    check_svd_error( Aref, U_FSVD, S_FSVD, Vt_FSVD )

    # run QMC before SVD
    block_time = 500  # in sec
    total_time = 3000 # in sec
    set_vmc_params()
    run_qmc()
    t_read = time.time()
    print(' getting QMCCHEM results from {}'.format(EZFIO_file) )
    results =  subprocess.check_output(['qmcchem', 'result', EZFIO_file],  encoding='UTF-8')
    print(' getting results after {} minutes \n'.format( (time.time()-t_read)/60. ))
    Eloc, ErrEloc = get_energy()
    print(' Eloc = {} +/- {}'.format(Eloc, ErrEloc))
    print(' {} <= Eloc <= {} \n'.format(Eloc-ErrEloc,Eloc+ErrEloc))

    #_________________________________________________________________________________________
    #
    # loop over truncated rank
    #_________________________________________________________________________________________

    low_rank_min = 5
    low_rank_max = 31

    for low_rank in range( low_rank_min, low_rank_max+1, 2 ):

        t_it = time.time()
        print(" SVD iteration rank = {}".format(low_rank))

        # SVD naive truncation
        U_TSVD  = U_FSVD[:, 0:low_rank+1]
        V_TSVD  = V_FSVD[:, 0:low_rank+1]
        Vt_TSVD = V_TSVD.T
        S_TSVD  = S_FSVD[0:low_rank+1]

        # check truncated SVD error
        check_svd_error( Aref, U_TSVD, S_TSVD, Vt_TSVD )

        # save new coefficients to EZFIO
        TSVD_save_EZFIO()

        # set QMC settings and run
        block_time = 500  # in sec
        total_time = 3000 # in sec
        set_vmc_params()
        run_qmc()

        # read QMC=CHEM results
        t_read = time.time()
        print(' getting QMCCHEM results from {}'.format(EZFIO_file) )
        results =  subprocess.check_output(['qmcchem', 'result', EZFIO_file],  encoding='UTF-8')
        print(' getting results after {} minutes \n'.format( (time.time()-t_read)/60. ))

        # < E_loc >
        Eloc, ErrEloc = get_energy()
        print(' Eloc = {} +/- {}'.format(Eloc, ErrEloc))
        print(' {} <= Eloc <= {} \n'.format(Eloc-ErrEloc,Eloc+ErrEloc))

        print(" iteration time = {:.3f} minutes \n".format((time.time()-t_it)/60.) )
    #_________________________________________________________________________________________

    print(" end after {:.3f} minutes".format((time.time()-t0)/60.) )

# end