#!/usr/bin/env python

from __future__ import print_function
import sys
import os
import errno
import atexit
import random

# Comment either this line or the following to attain different communication schemes:
# to start pyro
try:
    from IMP.isd.PyroGrid import PyroGrid as Grid
#   from IMP.isd.FileBasedGrid import FileBasedGrid as Grid
except ImportError:
    print("This example needs the Python Pyro module", file=sys.stderr)
    sys.exit(0)
from IMP.isd.hosts import create_host_list
import IMP.atom
import IMP.container
import IMP.isd
from IMP.isd.Replica import ReplicaTracker

import _shared_functions as sf

IMP.set_log_level(IMP.NONE)

# simulation settings
# where to output files
outfolder = os.path.join(os.getcwd(), 'results')
# temp dir
tmpdir = os.path.join(os.getcwd(), 'tmp')
# number of replicas / hosts
nreps = 8
# lambda scaling distribution
kB = (1.381 * 6.02214) / 4184.0
lambda_1 = 1.0
lambda_N = 0.8
lambdas = [lambda_N * (lambda_1 / lambda_N) ** ((float(nreps) - k) / (nreps - 1))
           for k in xrange(1, nreps + 1)]
# thermostat coupling constant (berendsen, in fs)
tau = [500.0] * nreps
# stat_rate is the rate at which to print out traj statistics
#(in units of gibbs sampling steps)
stat_rate = [1] * nreps
# list of files relative to the current dir to copy over to all nodes
initpdb = IMP.isd.get_example_path("generated2.pdb")
charmmtop = IMP.isd.get_example_path("top.lib")
charmmpar = IMP.isd.get_example_path("par.lib")
restraints = IMP.isd.get_example_path("NOE_HN-full_7A_sparse100.tbl")
sequence = IMP.isd.get_example_path('sequence.dat')
# export the files in a local tmp directory
# add whatever you want
filelist = [initpdb, charmmtop, charmmpar, sequence,
            restraints, IMP.isd.get_example_path('_shared_functions.py')]
# prefix of output files
nums = [os.path.join(outfolder, 'r%02d' % (i + 1)) for i in xrange(nreps)]
# thermalization (mc parameters stay fixed)
n_therm = 100  # number of loops, where temperatures are scaled to target value
n_hmc_therm = 10  # number of steps per loop
# number of gibbs sampling steps in  the first temperature relaxation
n_gibbs1 = 100
# number of gibbs sampling steps in  the second temperature relaxation
n_gibbs2 = 100
# number of gibbs sampling steps in  the production temperature relaxation
n_gibbs3 = 10000
# number of md steps
n_md = 1000
# number of mc steps (not < 50 because adaptive)
n_mc = 100
# where to run sims
hostlist = ['localhost'] * nreps
#qsub or ssh
grid_method = 'ssh'
# "-now no -cwd -q -j y -N 'slave' -S /bin/bash"
qsub_config = "-inherit $HOSTNAME"
# replica exchange scheme
rex_scheme = 'standard'
# replica exchange exchange method
rex_xchg = 'random'
# whether to use TuneRex to tune temperatures
tune_temps = True
tune_data = {'rate': 100,  # temp optimization rate, in rex steps.
             'method': 'ar',
             # either "ar" acceptance ratio optimization, or "flux"
             # optimization
             'alpha': 0.05}  # type I error on the estimates
templog = os.path.join(outfolder, 'temps.txt')
rexlog = os.path.join(outfolder, 'replicanums.txt')

# misc

imppy = os.path.abspath(
    os.path.join(os.getenv('IMP_ISD_DATA'), '../../tools/imppy.sh'))
src_path = os.path.abspath(
    os.path.join(os.getenv('IMP_ISD_DATA'), '../lib/IMP/isd'))
showX11 = False
grid_debug = False
grid_verbose = False
X11_delay = 1.0
window_size = '80x25'
# pyroGrid
shared_temp_path = True
terminate_during_publish = False
nshost = None


def mkdir_p(path):
    "mkdir -p, taken from stackoverflow"
    try:
        os.makedirs(path)
    except OSError as exc:
        if exc.errno == errno.EEXIST:
            pass
        else:
            raise


def launch_grid():
    hosts = create_host_list(hostlist, tmpdir)
    for host in hosts:
        # ugly hack
        host.init_cmd = imppy + ' !'

    # pyro grid
    grid = Grid(hosts, src_path, showX11, X11_delay, grid_debug,
                grid_verbose, shared_temp_path, nshost, terminate_during_publish,
                method=grid_method, qsub_config=qsub_config)

    # uncomment the following lines and comments the two previous ones to use file based grid
    # file based grid
    # grid = Grid(hosts, src_path, showX11, X11_delay, grid_debug,
    #        grid_verbose)
    #grid.shared_temp_path = shared_temp_path
    if showX11:
        grid.window_size = window_size
    grid.copy_files('./', filelist)
    # grid.copy_files(src_path,["_shared_functions.py"])

    # start grid on all nodes
    grid.start()
    return grid


def main():

    # launch grid
    print("launching grid")
    grid = launch_grid()

    # publish the shared object
    print("publishing sfo")
    sfo = sf.sfo()
    sfo_id = grid.publish(sfo)

    print("communication test")
    # get a worker to do a task
    #proxy = grid.acquire_service(sfo_id)
    # print proxy.hello().get()
    # grid.release_service(proxy)

    print("broadcast test")
    print(grid.gather(grid.broadcast(sfo_id, 'hello')))

    # call init on all nodes
    print("initializing model")
    mkdir_p(tmpdir)
    mkdir_p(outfolder)
    requests = grid.broadcast(
        sfo_id,
        'init_model',
        tmpdir,
        sequence,
        initpdb,
        restraints)
    # wait til init is done
    results = grid.gather(requests)

    # turn off verbose noise (works because IMP.NONE is picklable, being
    # an int.
    grid.gather(grid.broadcast(sfo_id, 'set_checklevel', IMP.NONE))
    grid.gather(grid.broadcast(sfo_id, 'set_loglevel', IMP.NONE))

    # evaluate the score of the whole system (without derivatives, False flag)
    print("initial energy")
    grid.gather(grid.broadcast(sfo_id, 'm', 'evaluate', False))

    # berendsen 300K tau=0.5ps
    # perform two independent MC moves for sigma and gamma
    print("initializing simulation and statistics")
    grid.gather(grid.scatter(sfo_id, 'init_simulation',
                             zip(lambdas, tau)))
    grid.gather(grid.scatter(sfo_id, 'init_stats', zip(nums, stat_rate)))
    replica = ReplicaTracker(nreps, lambdas, grid, sfo_id,
                             rexlog=rexlog,
                             scheme=rex_scheme, xchg=rex_xchg,
                             tune_temps=tune_temps, tune_data=tune_data, templog=templog)

    print("thermalization")
    for i in range(n_therm):
        print("\rgibbs step %d" % i, end=' ')
        sys.stdout.flush()
        grid.gather(grid.scatter(sfo_id, 'set_inv_temp',
                                 [n_therm / float(i + 1) * l for l in lambdas]))
        grid.gather(grid.broadcast(sfo_id, 'do_md', n_hmc_therm))
        grid.gather(grid.broadcast(sfo_id, 'write_stats'))

    print("start gibbs sampling loop: first relaxation")
    for i in range(n_gibbs1):
        print("\rgibbs step %d" % i, end=' ')
        sys.stdout.flush()
        # print " md"
        grid.gather(grid.broadcast(sfo_id, 'do_md', n_md))
        # print " mc"
        grid.gather(grid.broadcast(sfo_id, 'do_mc', n_mc))
        grid.gather(grid.broadcast(sfo_id, 'write_stats'))
        # print " swaps"
        replica.tune_data['dumb_scale'] = 0.5
        replica.replica_exchange()
        # print " stats"
        replica.write_rex_stats()

    print("start gibbs sampling loop: second relaxation")
    for i in range(n_gibbs2):
        print("\rgibbs step %d" % i, end=' ')
        sys.stdout.flush()
        # print " md"
        grid.gather(grid.broadcast(sfo_id, 'do_md', n_md))
        # print " mc"
        grid.gather(grid.broadcast(sfo_id, 'do_mc', n_mc))
        grid.gather(grid.broadcast(sfo_id, 'write_stats'))
        # print " swaps"
        replica.tune_data['dumb_scale'] = 0.2
        replica.replica_exchange()
        # print " stats"
        replica.write_rex_stats()

    print("start gibbs sampling loop: production")
    for i in range(n_gibbs3):
        print("\rgibbs step %d" % i, end=' ')
        sys.stdout.flush()
        # print " md"
        grid.gather(grid.broadcast(sfo_id, 'do_md', n_md))
        # print " mc"
        grid.gather(grid.broadcast(sfo_id, 'do_mc', n_mc))
        grid.gather(grid.broadcast(sfo_id, 'write_stats'))
        # print " swaps"
        replica.tune_temps = False
        replica.replica_exchange()
        # print " stats"
        replica.write_rex_stats()

    print("terminating grid")
    grid.terminate()
    print("done.")


if __name__ == '__main__':

    main()
