#!/usr/bin/env python

import numpy as np
from pyscf.pbc import gto, scf
import pyqmc
from pyqmc.slaterpbc import PySCFSlaterPBC, get_supercell
from pyqmc.dasktools import distvmc as vmc
from dask.distributed import Client, LocalCluster


def run_scf(nk):
    cell = gto.Cell()
    cell.atom = """
    He 0.000000000000   0.000000000000   0.000000000000
    """
    cell.basis = "bfd-vdz"
    cell.ecp = "bfd"
    cell.exp_to_discard = 1.0
    cell.a = """
    5.61, 0.00, 0.00
    0.00, 5.61, 0.00
    0.00, 0.00, 5.61"""
    cell.unit = "B"
    cell.verbose = 7
    cell.build()

    kpts = cell.make_kpts([nk, nk, nk])
    kmf = scf.KRHF(cell, exxdiv="ewald").density_fit()
    kmf.kpts = kpts
    ehf = kmf.kernel() 
    return cell, kmf


if __name__=="__main__":
    nk = 2
    cell, kmf = run_scf(nk)

    S = np.eye(3) * nk
    supercell = get_supercell(cell, S)
    wf = PySCFSlaterPBC(supercell, kmf)
    enacc = pyqmc.EnergyAccumulator(supercell)

    ncore = 20 
    nconfig = ncore * 300
    configs = pyqmc.initial_guess(supercell, nconfig, r=6.0) 

    cluster = LocalCluster(n_workers=ncore, threads_per_worker=1)
    client = Client(cluster)
    
    # Run VMC
    hdf_file = "bfd_hfvmc_nk{0}.hdf".format(nk)
    df, configs = vmc(
        wf, configs, nsteps=4000, accumulators={"energy": enacc}, hdf_file=hdf_file, verbose=True, 
        client=client, nsteps_per=50, 
    )