import scipy
import vqe_methods
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 = 1.60
    for rrr in range(0,23):
        if rrr < 17:
            r1 = 1.6 + rrr*0.05
        else:
            r1 = 2.4 + (rrr-16)*0.1

        geometry = '''
        H 0 0   0
        H 0 0  {0}
        H 0 {1}  0
        H 0 {1} {0}
        '''.format(r0,r1)

        geometry = [('H', (0,0,0)), ('H', ( 0, 0, r0)), ('H', (0, r1, 0)),('H',(0,r1,r0))]

        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=4)


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

        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,17,which='SA')
        for ei in range(len(e)):
            S2 = v[:,ei].conj().T.dot(s2.dot(v[:,ei]))
            print(" State %4i: %12.8f au  GAP:%16.8f <S2>: %12.8f" %(ei,e[ei]+E_nuc,e[ei]-e[0],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_GSD()
        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]))
        [e,v,params,ansatz_mat] = vqe_methods.adapt_vqe(fermi_ham, pool, reference_ket, theta_thresh=1e-9,adapt_thresh=1e-3)


        #vqe_methods.ucc(fermi_ham,pool,reference_ket)

        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_ipfull(n_orb,n_a,n_b,n_orb-n_a,n_orb-n_b,reference_ket)
        #print('op[0] is',op[0])
        #exit()

        #transform H with e^{sigma}
        barH=qeom.barH(params, ansatz_mat, hamiltonian)
        barS2=qeom.barH(params, ansatz_mat, s2)

        qeom.sceom(op,barH,barS2,reference_ket,e)
        qeom.qse(op,hamiltonian,reference_ket,e,v,E_nuc)
        qeom.ibm(op,hamiltonian,reference_ket,e,v)

        print("*rad r0:%4.2f"%r1)

if __name__== "__main__":
    test()
