import ffsim
import numpy as np
import pyscf
import pyscf.mcscf
import scipy.sparse.linalg
from qiskit_addon_sqd.fermion import rotate_integrals

# Build N2 molecule
mol = pyscf.gto.Mole()
mol.build(
    atom=[["N", (0, 0, 0)], ["N", (1.0, 0, 0)]],
    basis="sto-6g",
    symmetry="Dooh",
)

# Define active space
n_frozen = 4
active_space = range(n_frozen, mol.nao_nr())

# Get molecular integrals
scf = pyscf.scf.RHF(mol).run()
norb = len(active_space)
n_electrons = int(sum(scf.mo_occ[active_space]))
num_elec_a = (n_electrons + mol.spin) // 2
num_elec_b = (n_electrons - mol.spin) // 2
cas = pyscf.mcscf.CASCI(scf, norb, (num_elec_a, num_elec_b))
mo = cas.sort_mo(active_space, base=0)
hcore, nuclear_repulsion_energy = cas.get_h1cas(mo)
eri = pyscf.ao2mo.restore(1, cas.get_h2cas(mo), norb)
nelec = (num_elec_a, num_elec_b)

# Compute exact energy
exact_energy = cas.run().e_tot

# Rotate our integrals out of MO basis
rng = np.random.default_rng(24)
num_params = (norb**2 - norb) // 2  # antisymmetric, specified by upper triangle
k_rot = (rng.random(num_params) - 0.5) * 0.5
hcore_rot, eri_rot = rotate_integrals(hcore, eri, k_rot)

# Compute eigenvalue in original basis
mol_ham = ffsim.MolecularHamiltonian(
    one_body_tensor=hcore, two_body_tensor=eri, constant=nuclear_repulsion_energy
)
linop = ffsim.linear_operator(mol_ham, norb, nelec)
(eig,), _ = scipy.sparse.linalg.eigsh(linop, k=1, which="SA")

# Compute eigenvalue in rotated basis
mol_ham = ffsim.MolecularHamiltonian(
    one_body_tensor=hcore_rot, two_body_tensor=eri_rot, constant=nuclear_repulsion_energy
)
linop = ffsim.linear_operator(mol_ham, norb, nelec)
(eig_rot,), _ = scipy.sparse.linalg.eigsh(linop, k=1, which="SA")

np.testing.assert_allclose(eig, eig_rot, atol=1e-8)