# !!!
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_modifsvd():

    Ci_h_matrix_modifsvd = np.zeros( (n_modifsvd,n_modifsvd) )
    SErr_H               = np.zeros( (n_modifsvd,n_modifsvd) )

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

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




# ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! ~ ! 
def get_Ci_overlap_matrix_modifsvd():

    Ci_overlap_matrix_modifsvd = np.zeros( (n_modifsvd,n_modifsvd) )
    SErr_O                     = np.zeros( (n_modifsvd,n_modifsvd) )

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

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



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



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

    Ci_h_matrix_modifsvd      , SErr_H = get_Ci_h_matrix_modifsvd()
    Ci_overlap_matrix_modifsvd, SErr_O = get_Ci_overlap_matrix_modifsvd()

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

    psi_tmp           = psi_svd_coefs[0,0:n_modifsvd]
    recouvre_modifsvd = np.abs(psi_tmp @ vr)
    ind_gssvd = np.argmax(recouvre_modifsvd)

    E_modifsvd = eigvals_modifsvd[ind_gssvd] + E_toadd

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




if __name__ == '__main__':

    t0 = time.time()

    # ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ #
    #EZFIO_file = "Cr2_compactSVD30"

    #EZFIO_file = "Cr2_compactSVD30_it1"
    EZFIO_file = "Cr2_compactSVD30_it1_1"

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

    ezfio.set_file(EZFIO_file)
    n_svd         = ezfio.get_spindeterminants_n_svd_coefs()
    n_modifsvd    = ezfio.get_spindeterminants_n_modifsvd_coefs()
    psi_svd_coefs = np.array(ezfio.get_spindeterminants_psi_svd_coefs())

    print(" ")
    print(" Today's date:", datetime.now() )
    print(" EZFIO file = {}".format(EZFIO_file))
    print(" n_svd      = {}".format(n_svd) )
    print(" n_modifsvd = {}\n".format(n_modifsvd) )
    #print(' initial svd coeff = {}\n'.format(psi_svd_coefs))

    compactVMCmat_list = [ 'compactSVD_QMCmat_1.txt' , 'compactSVD_QMCmat_2.txt'  ]
    for compactVMCmat_file in compactVMCmat_list:
        print(" compactSVD matrices file = {}".format(compactVMCmat_file))
        E_compactsvd, sigma_compactsvd = get_Esvd()
        print(' compact svd enegry = {}'.format(E_compactsvd) )

    print('')
    compactVMCmat_file = 'compactSVD_QMCmat_1.txt'
    print(" compactSVD matrices file = {}".format(compactVMCmat_file))
    E_compactsvd, sigma_compactsvd = get_Esvd()
    sigma_compactsvd = sigma_compactsvd * np.sign(sigma_compactsvd[0])
    print(' compact svd enegry = {}'.format(E_compactsvd) )
    print(' compact svd coeff  = {}\n'.format(sigma_compactsvd))

    print(' compact svd factor  = {}\n'.format(sigma_compactsvd[n_modifsvd-1]))

    # ___________________________________________________________________________
    # ___________________________________________________________________________
    # 
    #save_to_EZFIO = 'compactsvd_QMC'
    save_to_EZFIO = ''

    if( save_to_EZFIO == 'compactsvd_QMC' ):
        ezfio_sigma_svd = np.zeros((n_svd))

        for i in range(n_modifsvd-1):
            ezfio_sigma_svd[i] = sigma_compactsvd[i]
        for i in range(n_modifsvd-1, n_svd):
            ezfio_sigma_svd[i] = psi_svd_coefs[0,i] * sigma_compactsvd[n_modifsvd-1]

        #print(' new svd coeff  = {}\n'.format(ezfio_sigma_svd))

        ezfio.set_spindeterminants_psi_svd_coefs( ezfio_sigma_svd )
        direc_svdcoeff = EZFIO_file + '/spindeterminants/psi_svd_coefs.gz'
        print(' {} is modified'.format(direc_svdcoeff) )
    else:
        print(' no modifications has been done in {}'.format(EZFIO_file))
    # ___________________________________________________________________________
    # ___________________________________________________________________________

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

    # !!!
# !!!

~