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

# spread of points: vary rho, sigma, tau to hit small/moderate/large s and alpha
rho = np.array([[0.5,0.5],[0.3,0.2],[1.0,0.8],[0.1,0.1],[2.0,1.0],[0.05,0.04]])
sigma = np.array([[0.1,0.05,0.1],[0.02,0.01,0.02],[0.5,0.3,0.4],
                  [0.001,0.0005,0.001],[3.0,2.0,2.5],[1e-4,5e-5,1e-4]])
tau = np.array([[0.3,0.3],[0.2,0.15],[0.9,0.7],[0.05,0.05],[1.5,0.8],[0.02,0.015]])
lapl = np.zeros_like(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); lib_vrho=o["vrho"].reshape(-1,2)
f=funxc.functional("MGGA_X_SCAN",polarized=True)
fx=f.exc_vxc(rho,sigma,lapl,tau)
fz=np.asarray(fx["zk"]).reshape(-1); fv=np.asarray(fx["vrho"])
print("pt   lib_zk         fx_zk          relerr_zk    relerr_vrhoA  relerr_vrhoB")
for i in range(len(rho)):
    rz=(fz[i]-lib_zk[i])/lib_zk[i]
    ra=(fv[i,0]-lib_vrho[i,0])/lib_vrho[i,0]
    rb=(fv[i,1]-lib_vrho[i,1])/lib_vrho[i,1]
    print(f"{i}  {lib_zk[i]: .6e}  {fz[i]: .6e}  {rz: .2e}  {ra: .2e}  {rb: .2e}")