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


from ezfio import ezfio
from math import sqrt
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 get_Ci_h_matrix_svd():

    Ci_h_matrix_svd = np.zeros( (n_svd,n_svd) )
    SErr_H          = np.zeros( (n_svd,n_svd) )

    f = open(VMCmat_file, 'r')
    Lines = f.readlines()
    for line in Lines:
        line = line.split()
        ii                      = int  ( line[0] )
        jj                      = int  ( line[1] )
        Ci_h_matrix_svd[ii][jj] = float( line[2] )
        SErr_H         [ii][jj] = float( line[3] )

    return(Ci_h_matrix_svd,SErr_H)
# ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! 



# ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! 
def get_Ci_overlap_matrix_svd():

    Ci_overlap_matrix_svd = np.zeros( (n_svd,n_svd) )
    SErr_O                = np.zeros( (n_svd,n_svd) )

    f = open(VMCmat_file, 'r')
    Lines = f.readlines()
    for line in Lines:
        line = line.split()
        ii                            = int  ( line[0] )
        jj                            = int  ( line[1] )
        Ci_overlap_matrix_svd[ii][jj] = float( line[4] )
        SErr_O               [ii][jj] = float( line[5] )

    return(Ci_overlap_matrix_svd,SErr_O)
# ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! 



# ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! 
def check_symmetric(a, tol=1e-3):
    return np.all(np.abs(a-a.T) < tol)
# ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! 



# ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! 
def get_Esvd():

    # read CI_SVD matrices
    Ci_h_matrix_svd      , SErr_H = get_Ci_h_matrix_svd()
    Ci_overlap_matrix_svd, SErr_O = get_Ci_overlap_matrix_svd()

    #print( 'Ci_h_matrix_svd is symmetric ? {}' .format(check_symmetric(Ci_h_matrix_svd)) )
    #print( 'Ci_overlap_matrix_svd is symmetric ? {}' .format(check_symmetric(Ci_overlap_matrix_svd)) )

    # symmetrise and diagonalise
    aa = Ci_h_matrix_svd
    aa = 0.5*( aa + aa.T )
    bb = Ci_overlap_matrix_svd
    eigvals_svd, vr = eig(aa, bb, left=False, right=True, overwrite_a=True, overwrite_b=True,
            check_finite=True, homogeneous_eigvals=False)

    #print( eigvals_svd + E_toadd )
    recouvre_svd = np.abs(psi_svd_coeff @ vr)
    ind_gssvd = np.argmax(recouvre_svd)

    E_svd = eigvals_svd[ind_gssvd] + E_toadd

    return( E_svd, vr[:,ind_gssvd] )
# ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! 




if __name__ == '__main__':

    t0 = time.time()

    # ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ #
    #EZFIO_file = "Cr2_TSVD30"
    #EZFIO_file = "Cr2_TSVD100"
    #EZFIO_file = "Cr2_TSVD200"
    #EZFIO_file = "Cr2_TSVD924"
    #EZFIO_file = "Cr2_TSVD924_1"

    #EZFIO_file = "Cr2_TSVD30_it1"
    #EZFIO_file = "Cr2_TSVD100_it1"
    #EZFIO_file = "Cr2_TSVD200_it1"
    EZFIO_file = "Cr2_TSVD924_1_it1"

    #EZFIO_file = "Cr2_TSVD200_minmaxit1"

    #EZFIO_file = "Cr2_TSVD30_it2"

    E_toadd = 64.8242130309350
    # ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ #

    ezfio.set_file(EZFIO_file)
    n_svd         = ezfio.get_spindeterminants_n_svd_coefs()
    psi_svd_coeff = np.array(ezfio.get_spindeterminants_psi_svd_coefs())

    print(" ")
    print(" Today's date:", datetime.now() )
    print(" EZFIO file = {}".format(EZFIO_file))
    print(" n_svd = {}\n".format(n_svd) )

    #VMCmat_list = [ 'SVD_QMCmat_FSVD_2b.txt'  , 'SVD_QMCmat_FSVD_3b.txt'  , 'SVD_QMCmat_FSVD_0.txt'    ]
    #VMCmat_list = [ 'SVD_QMCmat_FSVD_2b_1.txt', 'SVD_QMCmat_FSVD_3b_1.txt'
    #              , 'SVD_QMCmat_FSVD_4b_1.txt', 'SVD_QMCmat_FSVD_5b_1.txt'
    #              , 'SVD_QMCmat_FSVD_1.txt' ]
    #for VMCmat_file in VMCmat_list:
    #    print(" SVD matrices file = {}".format(VMCmat_file))
    #    E_svd, sigma_svd = get_Esvd()
    #    print(' svd enegry = {} \n'.format(E_svd) )


    VMCmat_file = 'SVD_QMCmat_FSVD_4b_1.txt'
    print(" SVD matrices file = {}".format(VMCmat_file))
    E_svd, sigma_svd = get_Esvd()
    sigma_svd = sigma_svd * np.sign(sigma_svd[0])
    print(' svd coeff  = {}'.format(sigma_svd))

    # ___________________________________________________________________________
    # ___________________________________________________________________________
    # 
    #save_to_EZFIO = 'svd_QMC'
    save_to_EZFIO = ''

    if( save_to_EZFIO == 'svd_QMC' ):
        ezfio.set_spindeterminants_psi_svd_coefs( sigma_svd )
        direc_svdcoeff = EZFIO_file + '/spindeterminants/psi_svd_coefs.gz'
        print(' {} is modified'.format(direc_svdcoeff) )
    else:
        pass
    # ___________________________________________________________________________
    # ___________________________________________________________________________

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

    # !!!
# !!!

~