#!/usr/bin/env python3

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



QMCCHEM_PATH=os.environ["QMCCHEM_PATH"]

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

from ezfio import ezfio

# PARAMETERS
thresh = 1.e-2
block_time = 6

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

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

    jast_type = ezfio.jastrow_jast_type
    print (jast_type)

    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_b():
        return ezfio.jastrow_mu_erf

    def set_params_b(x):
        x = np.abs(x)
        ezfio.set_jastrow_mu_erf(x)

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

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

    def get_norm():
        return 1.0, 0.0
        buffer = subprocess.check_output(['qmcchem', 'result', '-e', 'psi_norm', 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 get_energy():
        buffer = subprocess.check_output(['qmcchem', 'result', '-e', 'Ci_dress_mu_opt_qmcvar', 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 get_variance():
        buffer = subprocess.check_output(['qmcchem', 'result', '-e',
                                          'Ci_dress_mu_opt_qmcvar', filename],
                                         encoding='UTF-8')
        if buffer.strip() != "":
            buffer = buffer.splitlines()[-1]
            _, variance, error = [float(x) for x in buffer.split()]
            return variance, error
        else:
            return None, None



    def run_qmc():
        return 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', jast_type,
#                                 '-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', jast_type,
                                 '-l', str(block_time),
                                 filename])


    memo_energy = {'fmin': 100000000.}
    def f(x):
        print ("x = %s"%str(x))
        sys.stdout.flush()
        h = str(x)
        if h in memo_energy:
            return memo_energy[h]
        set_params_b(x[0])
        set_params_pen(x[1:])
        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(3.*block_time/4.)
            local_thresh = thresh
            while err > local_thresh:
                time.sleep(block_time)
                e, e_err = get_energy()
                variance, v_err = get_variance() 
                if e is None or variance is None:
                    continue
                norm, _ = get_norm() 
                e, e_err = e/norm, e_err/norm
                variance, v_err =variance/norm, v_err/norm
                energy = e #+ variance
                err = e_err #sqrt(e_err*e_err+v_err*v_err)
                print(" %f  %f    %f   %f"%(e, err, local_thresh, memo_energy['fmin']))
                sys.stdout.flush()
                if (energy-2.*err) > memo_energy['fmin']+local_thresh:
                   local_thresh = 10.*local_thresh
                elif (energy+2.*err) < memo_energy['fmin']-local_thresh:
                   local_thresh = 10.*local_thresh

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


    def run():
        x = np.array([ get_params_b() ] + list(get_params_pen()))
        if sum(x) == 0.:
          jast_a_up_dn = ezfio.jastrow_jast_a_up_dn
          x += jast_a_up_dn
        opt = sp.optimize.minimize(f,x,method="Powell",
                             options= {'disp':True, 'ftol':thresh,'xtol':0.02})
#        opt = sp.optimize.minimize(f,x,method="Nelder-Mead",
#                             options= {'disp':True, 'ftol':thresh,'xtol':0.02})
        print("x = "+str(opt))
        sys.stdout.flush()
        set_params_b(opt['x'][0])
        set_params_pen(opt['x'][1:])

    run()

if __name__ == '__main__':
    main()