from time import sleep

def V_idg(N, p, kT, cost=0):
    sleep(cost)
    return N * kT / p

"""
Explanation: 1.3. A Basic Workflow
This part of the tutorial requires numpy.
Operations
For this part of the tutorial we will imagine that we are still not convinced of the pressure-volume relations that we just "discovered" and that calculating the volume is actually a very expensive procedure, such as a many particle simulation.
We emulate this by adding an optional cost argument to our volume calculation function:
End of explanation
"""


def compute_volume(job):
    print('compute volume', job)
    V = V_idg(cost=1, **job.statepoint())
    job.document['V'] = V
    with open(job.fn('V.txt'), 'w') as file:
        file.write(str(V) + '\n')

"""
Explanation: It is useful to think of each modification of the workspace, that includes addition, modification, and removal of data, in terms of an operation.
An operation should take only one(!) argument: the job handle.
Any additional arguments may represent hidden state point parameters which would lead to a loss of provenance and possibly render our data space inconsistent.
The following function is an example for an operation:
End of explanation
"""


import signac

project = signac.get_project('projects/tutorial')

for job in project:
    compute_volume(job)

"""
Explanation: This operation computes the volume solely based on the state point parameters and stores the results such that they are clearly associated with the job, i.e., in the job document and in a file within the job's workspace.
Please note, that the only reason for storing the the same result in two different ways is for demonstration purposes.
Execution
To execute our first data space operation, we simply loop through our project's data space:
End of explanation
"""


import signac
import numpy as np

project = signac.get_project(root='projects/tutorial')

def init_statepoints(n):
    for p in np.linspace(0.1, 10.0, n):
        sp = {'p': p, 'kT': 1.0, 'N': 1000}
        job = project.open_job(sp)
        job.init()
        print('initialize', job)
        
init_statepoints(5)

"""
Explanation: Data Space Initialization
Since our operation is now more expensive, it is a good idea to split initialization and execution.
Let's initialize a few more state points in one go:
End of explanation
"""


for job in project:
    if 'V' not in job.document:
        compute_volume(job)

"""
Explanation: We see that initializing more jobs and even reinitializing old jobs is no problem.
However, since our calculation will be "expensive", we would want to skip the computation whenever the result is already available.
One possibility is to add a simple check before executing the computation:
End of explanation
"""


init_statepoints(10)

"""
Explanation: Classification
It would be even better, if we could get an overview of which state points have been calculated and which not.
We call this a project's status.
Before we continue, let's initialize a few more state points.
End of explanation
"""


def classify(job):
    yield 'init'
    if 'V' in job.document and job.isfile('V.txt'):
        yield 'volume-computed'

"""
Explanation: Next, we implement a classify() generator function, which labels a job based on certain conditions:
End of explanation
"""


print('Status: {}'.format(project))
for job in project:
    labels = ', '.join(classify(job))
    p = round(job.sp.p, 1)
    print(job, p, labels)

"""
Explanation: Our classifier will always yield the init label, but the volume-computed label is only yielded if the result has been computed and stored both in the job document and as a text file.
We can then use this function to get an overview of our project's status.
End of explanation
"""


for job in project:
    labels = classify(job)
    if 'volume-computed' not in labels:
        compute_volume(job)

"""
Explanation: Using only simple classification functions, we already get a very good grasp on our project's overall status.
Furthermore, we can use the classification labels for controling the execution of operations:
End of explanation
"""


list(map(compute_volume, project))
print('Done.')

"""
Explanation: Parallelization
So far, we have executed all operations in serial using a simple for-loop.
We will now learn how to easily parallelize the execution!
Instead of using a for-loop, we can also take advantage of Python's built-in map-operator:
End of explanation
"""


from multiprocessing import Pool

with Pool() as pool:
    pool.map(compute_volume, project)

"""
Explanation: Using the map() expression makes it trivial to implement parallelization patterns, for example, using a process Pool:
End of explanation
"""


from multiprocessing.pool import ThreadPool

with ThreadPool() as pool:
    pool.map(compute_volume, project)

"""
Explanation: Or a ThreadPool:
End of explanation
"""


# % rm -r projects/tutorial/workspace

"""
Explanation: Uncomment and execute the following line if you want to remove all data and start over.
End of explanation
"""


