#!/usr/bin/env python3

import scipy as sp
import scipy.optimize
import numpy as np
import sys
import os
import time
import subprocess

QMCCHEM_PATH=os.environ["QMCCHEM_PATH"]

sys.path.insert(0,QMCCHEM_PATH+"/EZFIO/Python/")

from ezfio import ezfio

# PARAMETERS
thresh = 1.e-3
block_time = 20

def main():
    if len(sys.argv) != 2:
      print("Usage: %s <EZFIO_DIRECTORY>"%sys.argv[0])
      sys.exti(1)

    filename = sys.argv[1]
    ezfio.set_file(filename)


    def make_atom_map():
        labels = {}
        dimension = 0
        for i,k in enumerate(ezfio.nuclei_nucl_label):
            if k in labels:
                labels[k].append(i)
            else:
                labels[k] = [dimension, i]
                dimension += 1
        atom_map = [[] for i in range(dimension)]
        for atom in labels.keys():
            l = labels[atom]
            atom_map[l[0]] = l[1:]
        return atom_map
    atom_map = make_atom_map()


    def get_params_pen():
        d = ezfio.jastrow_jast_pen
        return np.array([d[m[0]] for m in atom_map])


    def get_energy():
        buffer = subprocess.check_output(['qmcchem', 'result', '-e', 'e_loc', filename],
                                         encoding='UTF-8')
        if buffer.strip() != "":
            buffer = buffer.splitlines()[-1]
            _, energy, error = [float(x) for x in buffer.split()]
            return energy, error
        else:
            return None, None


    def set_params_pen(x):
        x = np.abs(x)
        y=list(ezfio.jastrow_jast_pen)
        for i,m in enumerate(atom_map):
            for j in m:
                y[j] = x[i]
        ezfio.set_jastrow_jast_pen(y)


    def run_qmc():
        subprocess.check_output(['qmcchem', 'run', filename])


    def stop_qmc():
        subprocess.check_output(['qmcchem', 'stop', filename])


    def set_vmc_params():
#        subprocess.check_output(['qmcchem', 'edit', '-c', '-j', 'Simple',
#                                 '-m', 'VMC',
#                                 '-l', str(block_time),
#                                 '--time-step=0.3',
#                                 '--stop-time=36000',
#                                 '--norm=1.e-5',
#                                 '-w', '10',
#                                 filename])
        subprocess.check_output(['qmcchem', 'edit', '-c', '-j', 'Simple',
                                 '-l', str(block_time),
                                 filename])


    memo_energy = {'fmin': 100000000.}
    def f(x):
        print ("x = %s"%str(x))
        h = str(x)
        if h in memo_energy:
            return memo_energy[h]
        set_params_pen(x)
        set_vmc_params()
        pid = os.fork()
        if pid == 0:
            run_qmc()
            os._exit(os.EX_OK)
        else:
            import atexit
            atexit.register(stop_qmc)

            err = thresh+1.
            time.sleep(block_time/2)
            local_thresh = thresh
            while err > local_thresh:
                time.sleep(block_time)
                energy, err = get_energy()
                if energy is not None:
                    print(" %f  %f"%(energy, err))
                    if (energy-2.*err) > memo_energy['fmin']+thresh:
                        local_thresh = 10.*thresh
                else:
                    err = local_thresh+1.

                # Check if PID is still running
                try: os.kill(pid,0)
                except OSError: err = 0.
            stop_qmc()
        os.wait()
        memo_energy[h] = energy
        memo_energy['fmin'] = min(energy, memo_energy['fmin'])
        return energy


    def run():
        x = get_params_pen()
        opt = sp.optimize.minimize(f,x,method="Powell",
                             options= {'disp':True, 'ftol':thresh,'xtol':0.02})
        print("x = "+str(opt))
        set_params_pen(opt['x'])

    run()

if __name__ == '__main__':
    main()