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, load_regression
import funxc

inp = load_input("Li")
rho, sigma, lapl, tau = inp["rho"], inp["sigma"], inp["lapl"], inp["tau"]
ref1 = load_regression("mgga_x", "mgga_x_scan", "Li", "pol", 1)

# installed library
from pylibxc import LibXCFunctional
lf = LibXCFunctional("MGGA_X_SCAN", "polarized")
out = 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_vrho = out["vrho"].reshape(-1,2)

f = funxc.functional("MGGA_X_SCAN", polarized=True)
fout = f.exc_vxc(rho, sigma, lapl, tau)
fx_vrho = np.asarray(fout["vrho"])

print("idx   golden_vrhoB      library_vrhoB     funxc_vrhoB")
for i in [3,5]:
    print(i, f"{ref1['vrho(b)'][i]: .10e}  {lib_vrho[i,1]: .10e}  {fx_vrho[i,1]: .10e}")
    print("   rho:", rho[i], "sigma:", sigma[i], "tau:", tau[i])