import sys, numpy as np
import jax; jax.config.update("jax_enable_x64", True)
sys.path.insert(0, "tests"); sys.path.insert(0, "libxc")
from libxc_reference import load_input
import funxc
from pylibxc import LibXCFunctional

inp = load_input("Li")
rho, sigma, lapl, tau = inp["rho"], inp["sigma"], inp["lapl"], inp["tau"]
lf = LibXCFunctional("MGGA_X_SCAN", "polarized")
o = lf.compute({"rho":np.ascontiguousarray(rho.reshape(-1)),
  "sigma":np.ascontiguousarray(sigma.reshape(-1)),
  "lapl":np.ascontiguousarray(lapl.reshape(-1)),
  "tau":np.ascontiguousarray(tau.reshape(-1))})
lib_zk = o["zk"].reshape(-1)
f = funxc.functional("MGGA_X_SCAN", polarized=True)
fx = f.exc_vxc(rho, sigma, lapl, tau)
fx_zk = np.asarray(fx["zk"]).reshape(-1)
for i in [3,5]:
    print(i, f"lib_zk={lib_zk[i]:.16e} fx_zk={fx_zk[i]:.16e} reldiff={(fx_zk[i]-lib_zk[i])/lib_zk[i]:.2e}")

# Also finite-difference libxc zk*rho wrt rho_dn to get vrho(b) and compare
def e_of(rd, i):
    r = rho.copy(); r[i,1]=rd
    oo = lf.compute({"rho":np.ascontiguousarray(r.reshape(-1)),
      "sigma":np.ascontiguousarray(sigma.reshape(-1)),
      "lapl":np.ascontiguousarray(lapl.reshape(-1)),
      "tau":np.ascontiguousarray(tau.reshape(-1))})
    n = r[i,0]+r[i,1]
    return oo["zk"].reshape(-1)[i]*n
for i in [3,5]:
    rd=rho[i,1]; h=rd*1e-6
    fd=(e_of(rd+h,i)-e_of(rd-h,i))/(2*h)
    print(i,"lib_FD_vrhoB=",fd)