"""Functions from PySCF that need to be tweeked to allow our hack of adding Vemb to Hcore."""

import logging

import numpy as np
from numpy.typing import NDArray
from pyscf import ao2mo  # type:ignore

logger = logging.getLogger(__name__)


def energy_elec(mf, dm=None, h1e=None, vhf=None) -> tuple[float, float]:
    """Electronic energy of Unrestricted Hartree-Fock.

    Note this function has side effects which cause mf.scf_summary updated.

    Args:
        mf (pyscf.scf.hf.HF): Hartree-Fock object
        dm (np.ndarray): Density matrix
        h1e (np.ndarray): Core Hamiltonian
        vhf (np.ndarray): 2-electron contribution to effective potential

    Returns:
        e_elec (np.ndarray): Hartree-Fock electronic energy
        e_coul (np.ndarray): 2-electron contribution to electronic energy
    """
    logger.debug("Running spin-aware energy_elec")
    if dm is None:
        dm = mf.make_rdm1()
    if h1e is None:
        h1e = mf.get_hcore()
    if isinstance(dm, np.ndarray) and dm.ndim == 2:
        dm = np.array((dm * 0.5, dm * 0.5))
    if vhf is None:
        vhf = mf.get_veff(mf.mol, dm)
    logger.debug(f"{h1e.shape=}")
    logger.debug(f"{dm.shape=}")
    logger.debug(f"{vhf.shape=}")
    e1 = np.einsum("ij,ji->", h1e[0], dm[0])
    e1 += np.einsum("ij,ji->", h1e[1], dm[1])
    e_coul = (
        np.einsum("ij,ji->", vhf[0], dm[0]) + np.einsum("ij,ji->", vhf[1], dm[1])
    ) * 0.5
    e_elec = (e1 + e_coul).real
    mf.scf_summary["e1"] = e1.real
    mf.scf_summary["e2"] = e_coul.real
    return e_elec, e_coul


def _absorb_h1e(h1e, eri, norb, nelec, fac=1):
    if not isinstance(nelec, (int, np.number)):
        nelec = sum(nelec)
    h1e_a, h1e_b = h1e
    print("h1e_a", h1e_a.shape)
    print("h1e_b", h1e_b.shape)
    h1e_a = np.einsum("jik->jk", h1e_a)
    h1e_b = np.einsum("jik->jk", h1e_b)
    print("h1e_a", h1e_a.shape)
    print("h1e_b", h1e_b.shape)
    h2e_aa: NDArray = ao2mo.restore(1, eri[0], norb).copy()  # type: ignore
    h2e_ab: NDArray = ao2mo.restore(1, eri[1], norb).copy()  # type: ignore
    h2e_bb: NDArray = ao2mo.restore(1, eri[2], norb).copy()  # type: ignore
    print("x1", h2e_aa.shape)
    print("y2", h2e_bb.shape)
    x = np.einsum("jiik->jk", h2e_aa) * 0.5
    y = np.einsum("jiik->jk", h2e_bb) * 0.5
    print("x", x.shape)
    print("y", y.shape)
    f1e_a = h1e_a - np.einsum("jiik->jk", h2e_aa) * 0.5
    f1e_b = h1e_b - np.einsum("jiik->jk", h2e_bb) * 0.5
    f1e_a *= 1.0 / (nelec + 1e-100)
    f1e_b *= 1.0 / (nelec + 1e-100)
    for k in range(norb):
        h2e_aa[:, :, k, k] += f1e_a
        h2e_aa[k, k, :, :] += f1e_a
        h2e_ab[:, :, k, k] += f1e_a
        h2e_ab[k, k, :, :] += f1e_b
        h2e_bb[:, :, k, k] += f1e_b
        h2e_bb[k, k, :, :] += f1e_b
    return (
        ao2mo.restore(4, h2e_aa, norb) * fac,  # type: ignore
        ao2mo.restore(4, h2e_ab, norb) * fac,  # type: ignore
        ao2mo.restore(4, h2e_bb, norb) * fac,  # type: ignore
    )
