#!/usr/bin/env python

import numpy as np
import subprocess
import sys, os
import time
from datetime import datetime

QP_PATH=os.environ["QP_ROOT"]+"/external/ezfio/Python/"
sys.path.insert(0,QP_PATH)

from ezfio import ezfio


#____________________________________________________________________________________
#
def get_energy():
    energy_buffer = subprocess.check_output(
                  ['qmcchem', 'result', '-e', 'e_loc', EZFIO_file]
                  ,  encoding='UTF-8')

    if energy_buffer.strip() != "":
        energy_buffer = energy_buffer.splitlines()[-1]
        _, energy, error = [float(x) for x in energy_buffer.split()]
        return energy, error
    else:
      return None, None
#____________________________________________________________________________________



#____________________________________________________________________________________
#
def read_scalar(scalar):

    scal_txt = scalar + ' : '
    beg_read = results.find(scal_txt) + len(scal_txt)
    scal_buf = results[beg_read:beg_read+100].split('\n')[0].split()

    scal_val , scal_err = float( scal_buf[0] ) , float( scal_buf[2] )

    return(scal_val,scal_err)
#____________________________________________________________________________________




#____________________________________________________________________________________
#
def read_ci_dress_svd():

    Ci_dress_svd_val = np.zeros((n_svd))
    Ci_dress_svd_err = np.zeros((n_svd))

    beg_read = results.find('Ci_dress_svd : [ ') + len( 'Ci_dress_svd : [ ' )
    end_read = len(results)
    vect_buf = results[beg_read:end_read]
    vect_buf = vect_buf.split( '\n' )

    for iline in range(1, n_svd+1):
        line = vect_buf[iline].split()

        # check
        indc = int( line[0] )
        if( indc != iline ):
            print(' Error in reading Ci_dress_svd')
            print(' indc  = {}'.format(indc) )
            print(' iline = {}'.format(iline) )
            break

        else:
            Ci_dress_svd_val[indc-1] = float( line[2] )
            Ci_dress_svd_err[indc-1] = float( line[4] )
            # !!!
        # !!!
    # !!!

    return(Ci_dress_svd_val,Ci_dress_svd_err)
#____________________________________________________________________________________


#____________________________________________________________________________________
#
def get_ci_dress():

    U_svd = np.array(ezfio.get_spindeterminants_psi_svd_alpha())
    V_svd = np.array(ezfio.get_spindeterminants_psi_svd_beta())
    U_svd = U_svd[0,:,:].T
    V_svd = V_svd[0,:,:].T

    ci_dress_mat_val = np.dot(U_svd,np.dot(np.diagflat(Ci_dress_svd_val),V_svd.T)) 
    ci_dress_mat_err = np.dot(U_svd,np.dot(np.diagflat(Ci_dress_svd_err),V_svd.T)) 

    n_det = ezfio.get_spindeterminants_n_det()
    i_row = np.array(ezfio.get_spindeterminants_psi_coef_matrix_rows())
    j_col = np.array(ezfio.get_spindeterminants_psi_coef_matrix_columns())
    
    Ci_dress_val = np.zeros((n_det))
    Ci_dress_err = np.zeros((n_det))
    for k in range(n_det):
        i = i_row[k] - 1
        j = j_col[k] - 1
        Ci_dress_val[k] = ci_dress_mat_val[i,j] 
        Ci_dress_err[k] = ci_dress_mat_err[i,j] 

    Ci_dress_val = Ci_dress_val / psi_norm
    Ci_dress_err = Ci_dress_err / psi_norm

    return(Ci_dress_val,Ci_dress_err)
#____________________________________________________________________________________




#____________________________________________________________________________________
#
def save_to_file(v_val,v_err, name_f):

    with open(name_f, 'w') as ff:
        for i in range(len(v_val)):
            ff.write(' {}     {}     {}\n'.format(i, v_val[i], v_err[i]))
#____________________________________________________________________________________



if __name__ == '__main__':
    t0 = time.time()

    # _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
    #
    #EZFIO_file = sys.argv[1]
    #EZFIO_file = "/home/aammar/qp2/src/svdwf/cr2_work/cr2_1p6_dress"
    EZFIO_file = "/home/aammar/qp2/src/svdwf/cr2_work/cr2_1p6_dress_0"
    # _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

    ezfio.set_file(EZFIO_file)

    n_svd = ezfio.get_spindeterminants_n_svd_coefs()

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

    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. ))

    # read QMC results
    E_loc, ErrE_loc = get_energy()
    print(' E_loc = {} +/- {}'.format(E_loc, ErrE_loc))
    print(' {} <= E_loc <= {}\n'.format(E_loc-ErrE_loc,E_loc+ErrE_loc))

    E_loc_qmcvar, ErrE_loc_qmcvar = read_scalar('E_loc_qmcvar')
    print(' E_loc_qmcvar = {} +/- {} \n'.format(E_loc_qmcvar,ErrE_loc_qmcvar))

    psi_norm, Errpsi_norm = read_scalar('Psi_norm')
    print(' psi_norm = {} +/- {}\n'.format(psi_norm,Errpsi_norm))

    psi_norm_qmcvar, Errpsi_norm_qmcvar = read_scalar('Psi_norm_qmcvar')
    print(' Psi_norm_qmcvar = {} +/- {} \n'.format(psi_norm_qmcvar, Errpsi_norm_qmcvar))

    Ci_dress_svd_val, Ci_dress_svd_err = read_ci_dress_svd()
    save_to_file(Ci_dress_svd_val,Ci_dress_svd_err, 'ci_dress_svd.txt')

    Ci_dress_val, Ci_dress_err = get_ci_dress()
    save_to_file(Ci_dress_val,Ci_dress_err, 'ci_dress.txt')


    # _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
    #
    #                       SAVE TO EZFIO
    #
    #ezfio.set_dmc_dress_dmc_delta_h(Ci_dress_val)
    # _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
    #


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