import sys, os, time

"""
Explanation: AdaptiveMD
Example 1 - Setup
0. Imports
End of explanation
"""


# verbose = os.environ.get('RADICAL_PILOT_VERBOSE', 'REPORT')
os.environ['RADICAL_PILOT_VERBOSE'] = 'ERROR'

"""
Explanation: We want to stop RP from reporting all sorts of stuff for this example so we set a specific environment variable to tell RP to do so. If you want to see what RP reports change it to REPORT.
End of explanation
"""


from adaptivemd import Project

from adaptivemd.engine.openmm import OpenMMEngine
from adaptivemd.analysis.pyemma import PyEMMAAnalysis

from adaptivemd import File, Directory, WorkerScheduler

"""
Explanation: We will import the appropriate parts from AdaptiveMD as we go along so it is clear what it needed at what stage. Usually you will have the block of imports at the beginning of your script or notebook as suggested in PEP8.
End of explanation
"""


# Project.delete('test')

project = Project('testcase-worker')

"""
Explanation: Let's open a project with a UNIQUE name. This will be the name used in the DB so make sure it is new and not too short. Opening a project will always create a non-existing project and reopen an exising one. You cannot chose between opening types as you would with a file. This is a precaution to not accidentally delete your project.
End of explanation
"""


print len([t for t in project.trajectories if t.exists])

for w in project.workers:
    print w.hostname, w.state

from uuid import UUID

import datetime
datetime.datetime.fromtimestamp(modeller.__time__).strftime("%Y-%m-%d %H:%M:%S")

project.generators.add(engine)
project.generators.add(modeller)

sc = WorkerScheduler(project.resource)
sc.enter(project)

t = engine.task_run_trajectory(project.new_trajectory(pdb_file, 100, restart=True)). extend(50).extend(100)

sc(t)

sc.advance()

print project.generators

t1 = engine.task_run_trajectory(project.new_trajectory(pdb_file, 100, restart=True))
t2 = t1.extend(100)

t2.trajectory.restart

project.tasks.add(t2)

for f in project.trajectories:
    print f.drive, f.basename, len(f), f.created, f.__time__, f.exists, hex(f.__uuid__)

for f in project.files:
    print f.drive, f.path, f.created, f.__time__, f.exists, hex(f.__uuid__)

w = project.workers.last
print w.state
print w.command

for t in project.tasks:
    print t.state, t.worker.hostname if t.worker else 'None'

sc.advance()

t1 = engine.task_run_trajectory(project.new_trajectory(pdb_file, 100))
t2 = t1.extend(100)

project.tasks.add(t2)

# from adaptivemd.engine import Trajectory
# t3 = engine.task_run_trajectory(Trajectory('staging:///trajs/0.dcd', pdb_file, 100)).extend(100)
# t3.dependencies = []

# def get_created_files(t, s):
#     if t.is_done():
#         print 'done', s
#         return s - set(t.added_files)
#     else:
#         adds = set(t.added_files)
#         rems = set(s.required[0] for s in t._pre_stage)
#         print '+', adds
#         print '-', rems
#         q = set(s) - adds | rems 
        
#         if t.dependencies is not None:
#             for d in t.dependencies:                
#                 q = get_created_files(d, q)

#         return q
    
# get_created_files(t3, {})

for w in project.workers:
    print w.hostname, w.state

w = project.workers.last
print w.state
print w.command

w.command = 'shutdown'

for t in project.tasks:
    print t.state, t.worker.hostname if t.worker else 'None'

for f in project.trajectories:
    print f.drive, f.basename, len(f), f.created, f.__time__, f.exists, hex(f.__uuid__)

project.trajectories.one[0]

t = engine.task_run_trajectory(project.new_trajectory(project.trajectories.one[0], 100))

project.tasks.add(t)

print project.files
print project.tasks

t = modeller.execute(list(project.trajectories))

project.tasks.add(t)

from uuid import UUID

project.storage.tasks._document.find_one({'_dict': {'generator' : { '_dict': }}})

genlist = ['openmm']


scheduler = sc
prefetch = 1

while True:
    scheduler.advance()
    if scheduler.is_idle:
        for _ in range(prefetch):
            tasklist = scheduler(project.storage.tasks.consume_one())

        if len(tasklist) == 0:
            break

    time.sleep(2.0)

"""
Explanation: Now we have a handle for our project. First thing is to set it up to work on a resource.
End of explanation
"""


scheduler = project.get_scheduler(cores=1)

"""
Explanation: Note, that you cannot add the same engine twice. But if you create a new engine it will be considered different and hence you can store it again. 
Create one intial trajectory
Finally we are ready to run a first trajectory that we will store as a point of reference in the project. Also it is nice to see how it works in general.
1. Open a scheduler
a job on the cluster to execute tasks
the .get_scheduler function delegates to the resource and uses the get_scheduler functions from there. This is merely a convenience since a Scheduler has the responsibility to open queues on the resource for you. 
You have the same options as the queue has in the resource. This is often the number of cores and walltime, but can be additional ones, too. 
Let's open the default queue and use a single core for it since we only want to run one simulation.
End of explanation
"""


trajectory = project.new_trajectory(engine['pdb_file'], 100)
trajectory

"""
Explanation: Next we create the parameter for the engine to run the simulation. Since it seemed appropriate we use a Trajectory object (a special File with initial frame and length) as the input. You could of course pass these things separately, but this way, we can actualy reference the no yet existing trajectory and do stuff with it.
A Trajectory should have a unique name and so there is a project function to get you one. It uses numbers and makes sure that this number has not been used yet in the project.
End of explanation
"""


task = engine.task_run_trajectory(trajectory)

"""
Explanation: This says, initial is alanine.pdb run for 100 frames and is named xxxxxxxx.dcd.
Now, we want that this trajectory actually exists so we have to make it (on the cluster which is waiting for things to do). So we need a Task object to run a simulation. Since Task objects are very flexible there are helper functions to get them to do, what you want, like the ones we already created just before. Let's use the openmm engine to create an openmm task
End of explanation
"""


scheduler(task)

"""
Explanation: That's it, just that a trajectory description and turn it into a task that contains the shell commands and needed files, etc. 
Last step is to really run the task. You can just use a scheduler as a function or call the .submit() method.
End of explanation
"""


scheduler.is_idle

print scheduler.generators

"""
Explanation: Now we have to wait. To see, if we are done, you can check the scheduler if it is still running tasks.
End of explanation
"""


# scheduler.wait()

"""
Explanation: or you wait until it becomes idle using .wait()
End of explanation
"""


print project.files
print project.trajectories

"""
Explanation: If all went as expected we will now have our first trajectory.
End of explanation
"""


scheduler.exit()

"""
Explanation: Excellent, so cleanup and close our queue
End of explanation
"""


project.close()

"""
Explanation: and close the project.
End of explanation
"""


