
import scipy
import openfermion
import openfermionpsi4
import os
import numpy as np
import copy
import random
import sys

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

import operator_pools
import vqe_methods

from tVQE import *
from orbital_opt import *
from openfermion import *


import pyscf_helper
import orbital_opt
def ooadapt(no,nia,nib,nva,nvb,hamiltonian_op, pool, reference_ket,
        adapt_conver    = 'norm',
        adapt_thresh    = 1e-3,
        theta_thresh    = 1e-7,
        adapt_maxiter   = 200,
        psi4_filename   = "psi4_%12.12f"%random.random()
        ):
# {{{

    hamiltonian = openfermion.transforms.get_sparse_operator(hamiltonian_op)
    ref_energy = reference_ket.T.conj().dot(hamiltonian.dot(reference_ket))[0,0].real
    print(" Reference Energy: %12.8f" %ref_energy)
    #Thetas
    parameters = []

    pool.generate_SparseMatrix()
    pool.gradient_print_thresh = theta_thresh

    ansatz_ops = []     #SQ operator strings in the ansatz
    ansatz_mat = []     #Sparse Matrices for operators in ansatz

    print(" Start ADAPT-VQE algorithm")
    op_indices = []
    parameters = []
    curr_state = 1.0*reference_ket
    singles =orbital_opt.ops_singlesGSD_unitary(no,nia,nib,nva,nvb)
    par_s=[]
    np.random.seed(42)
    for i in range(len(singles)):
        par_s.insert(0, np.random.random()/10)  #how to initialize the parameters?
    print(" Now start to grow the ansatz")
    for n_iter in range(0,adapt_maxiter):

        print("\n\n\n")
        print(" --------------------------------------------------------------------------")
        print("                         ADAPT-VQE iteration: ", n_iter)
        print(" --------------------------------------------------------------------------")
        next_index = None
        next_deriv = 0
        curr_norm = 0

        min_options = {'gtol': theta_thresh, 'disp':False}
       
        #create effective Hamiltonian
        hnew,par_s=orbopt(no,nia,nib,nva,nvb,curr_state,hamiltonian,min_options,singles,par_s)
        print(hnew.dot(curr_state), hamiltonian.dot(curr_state))
        exit()   


        print(" Check each new operator for coupling")
        next_term = []
        print(" Measure Operator Pool Gradients:")

        #sig = hamiltonian.dot(curr_state)
        
        #constuct sig using effective Hamiltonian
        sig = hnew.dot(curr_state)
        



        e_curr = curr_state.T.conj().dot(sig)[0,0]
        var = sig.T.conj().dot(sig)[0,0] - e_curr**2
        uncertainty = np.sqrt(var.real)
        assert(np.isclose(var.imag,0))
        print(" Variance:    %12.8f" %var.real)
        print(" Uncertainty: %12.8f" %uncertainty)
        for oi in range(pool.n_ops):

            gi = pool.compute_gradient_i(oi, curr_state, sig)

            curr_norm += gi*gi
            if abs(gi) > abs(next_deriv):
                next_deriv = gi
                next_index = oi

        curr_norm = np.sqrt(curr_norm)



        max_of_gi = next_deriv
        print(" Norm of <[H,A]> = %12.8f" %curr_norm)
        print(" Max  of <[H,A]> = %12.8f" %max_of_gi)
        converged = False
        if adapt_conver == "norm":
            if curr_norm < adapt_thresh:
                converged = True
        elif adapt_conver == "var":
            if abs(var) < adapt_thresh:
                #variance
                converged = True
        else:
            print(" FAIL: Convergence criterion not defined")
            exit()

        if converged:
            print(" Ansatz Growth Converged!")
            print(" Number of operators in ansatz: ", len(ansatz_ops))
            print(" *Finished: %20.12f" % trial_model.curr_energy)
            print(" -----------Final ansatz----------- ")
            print(" %4s %12s %18s" %("#","Coeff","Term"))
            for si in range(len(ansatz_ops)):
                opstring = pool.get_string_for_term(ansatz_ops[si])
                print(" %4i %12.8f %s" %(si, parameters[si], opstring) )
                #exit()
            break

        print(" Add operator %4i" %next_index)
        parameters.insert(0,0)
        ansatz_ops.insert(0,pool.fermi_ops[next_index])
        ansatz_mat.insert(0,pool.spmat_ops[next_index])
        #ansatz_mat gives matrix form of a^a operators acting on 4 qubits. par*ansatz_mat[i] 
        #print('ansatz_mat',ansatz_mat[0].toarray())
        #exit()

        #using new H
        #trial_model = tUCCSD(hamiltonian, ansatz_mat, reference_ket, parameters)
        trial_model = tUCCSD(hnew, ansatz_mat, reference_ket, parameters)


        opt_result = scipy.optimize.minimize(trial_model.energy, parameters, jac=trial_model.gradient,
                options = min_options, method = 'BFGS', callback=trial_model.callback)

        parameters = list(opt_result['x'])
        curr_state = trial_model.prepare_state(parameters)
        print(" Finished: %20.12f" % trial_model.curr_energy)
        print(" -----------New ansatz----------- ")
        print(" %4s %12s %18s" %("#","Coeff","Term"))
        for si in range(len(ansatz_ops)):
            opstring = pool.get_string_for_term(ansatz_ops[si])
            print(" %4i %12.8f %s" %(si, parameters[si], opstring) )
    return trial_model.curr_energy, curr_state, parameters,ansatz_mat

# }}}

