import scipy
import vqe_methods 
import vqe_methods_edit
import operator_pools
import pyscf_helper 

import pyscf
from pyscf import lib
from pyscf import gto, scf, mcscf, fci, ao2mo, lo,  cc
from pyscf.cc import ccsd

import openfermion 
from openfermion import *
from tVQE import *

import qeom
from scipy.sparse.linalg import eigs

def test():
    r0 = 1.5
    r1 = 2.0
    r2 = 100.0
    r3 = 102.0
    geometry = '''
    H 0 0   0
    H 0 0  {0}
    H 0 {1}  0
    H 0 {1} {0}
    H 0 {2}   0
    H 0 {2}  {0}
    H 0 {3}  0
    H 0 {3} {0}
    '''.format(r0,r1,r2,r3)

    charge = 0
    spin = 0
    basis = 'sto3g'
    [n_orb, n_a, n_b, h, g, mol, E_nuc, E_scf, C, S] = pyscf_helper.init(geometry,charge,spin,basis,n_frzn_occ=0,n_act=8)

    print(" n_orb: %4i" %n_orb)
    print(" n_a  : %4i" %n_a)
    print(" n_b  : %4i" %n_b)

    sq_ham = pyscf_helper.SQ_Hamiltonian()

    sq_ham.init(h, g, C, S)
    print(" HF Energy: %12.8f" %(E_nuc + sq_ham.energy_of_determinant(range(n_a),range(n_b))))
    #ehf = E_nuc + sq_ham.energy_of_determinant(range(n_a),range(n_b))
    #assert(abs(ehf - -1.82913741) < 1e-7)
    fermi_ham  = sq_ham.export_FermionOperator()
    hamiltonian = openfermion.transforms.get_sparse_operator(fermi_ham)
    #print(hamiltonian.toarray())
    #exit()

    s2 = vqe_methods.Make_S2(n_orb)

    #build reference configuration
    occupied_list = []
    for i in range(n_a):
        occupied_list.append(i*2)
    for i in range(n_b):
        occupied_list.append(i*2+1)

    print(" Build reference state with %4i alpha and %4i beta electrons" %(n_a,n_b), occupied_list)

    reference_ket = scipy.sparse.csc_matrix(openfermion.jw_configuration_state(occupied_list, 2*n_orb)).transpose()

    [e,v] = scipy.sparse.linalg.eigsh(hamiltonian.real,7,which='SA',v0=reference_ket.todense())
    for ei in range(len(e)):
        S2 = v[:,ei].conj().T.dot(s2.dot(v[:,ei]))
        print(" State %4i: %12.8f au  <S2>: %12.8f" %(ei,e[ei]+E_nuc,S2))
    
    fermi_ham += FermionOperator((),E_nuc)
    #pyscf.molden.from_mo(mol, "full.molden", sq_ham.C)

    #   Francesco, change this to singlet_GSD() if you want generalized singles and doubles
    pool = operator_pools.singlet_SD()
    pool.init(n_orb, n_occ_a=n_a, n_occ_b=n_b, n_vir_a=n_orb-n_a, n_vir_b=n_orb-n_b)
    for i in range(pool.n_ops):
        print(pool.get_string_for_term(pool.fermi_ops[i]))
    print("starting adapt vqe function")
    #[e,v,params,ansatz_mat] = vqe_methods.adapt_vqe(fermi_ham, pool, reference_ket, theta_thresh=1e-9)
    [e,v,params,ansatz_mat] = vqe_methods_edit.adapt_vqe_edit(fermi_ham, pool, reference_ket, theta_thresh=1e-9,det_limit=4)

    print(" Final ADAPT-VQE energy: %12.8f" %e)
    print(" <S^2> of final state  : %12.8f" %(v.conj().T.dot(s2.dot(v))[0,0].real))
    #exit()
    #store FCI values for checking
    #a,b=eigs(hamiltonian)
    #fci_levels=a+E_nuc

    #create operators single and double for each excitation
    op=qeom.createops_eefullwithI(n_orb,n_a,n_b,n_orb-n_a,n_orb-n_b,reference_ket)
    print("testing gr state energy")
    #transform 
    left=qeom.wfgr(params,ansatz_mat,reference_ket)
    print('barH based energy diff=0?',qeom.expvalue(left.transpose().conj(),hamiltonian,left)[0,0].real-e+E_nuc)    

    #create ex operator

    M=np.zeros((len(op),len(op)))
    V=np.zeros((len(op),len(op)))
    print("start creating M and S")
    for i in range(len(op)):
        for j in range(len(op)):
            #left=qeom.wfgr(params, ansatz_mat,op[i],reference_ket)
            #mat=op[i].transpose().conj().dot(barH.dot(op[j]))
            mat1=qeom.expvalue(op[i].transpose().conj(),hamiltonian,op[j])
            mat1=scipy.sparse.csr_matrix(mat1)
            M[i,j]=qeom.expvalue(v.transpose().conj(),mat1,v)[0,0]
            mat3=op[i].transpose().conj().dot(op[j])
            V[i,j]=qeom.expvalue(v.transpose().conj(),mat3,v)[0,0]
    print("start diagonalizing")
    eig,aval=scipy.linalg.eig(M,V)
    #print('V',V)
    eigprint=np.sort(eig.real)
    print('final excited state energies',(np.sort(eig.real)+E_nuc))
    print('final excitation energies',(np.sort(eig.real)-eigprint[0]))
    #print('final excitation energies',(np.sort(eig.real)+E_nuc-e))
    #print('eigenvector 1st',aval[0])
    #idx = eig.argsort()
    #eig = eig[idx]
    #aval = aval[:, idx]

    #for ei in range(len(eig)):
    #    print(" State %4i: %12.8f au Gap:%12.8f" %(ei,eig[ei]-e,eig[ei]))
if __name__== "__main__":
    test()
