#!/usr/bin/env python
# coding: utf-8

# # Implementing a new simulation code
# Appendix C in the pyiron paper
# https://www.sciencedirect.com/science/article/pii/S0927025618304786

# In[1]:


from os.path import join


# In[2]:


from pyiron_base import GenericJob, GenericParameters


# In[3]:


class ToyJob(GenericJob):
    def __init__(self, project, job_name):
        super().__init__(project, job_name) 
        self.input = GenericParameters(table_name="input")
        self.input['input_energy'] = 100
        self.executable = "cat input > output"

    def write_input(self): 
        self.input.write_file( 
            file_name="input",
            cwd=self.working_directory
        )

    def collect_output(self):
        file = join(self.working_directory, "output") 
        with open(file) as f:
            line = f.readlines()[0]
        energy = float(line.split()[1]) 
        with self.project_hdf5.open("output/generic") as h5out: 
            h5out["energy_tot"] = energy

    def to_hdf(self, hdf=None, group_name=None): 
        super().to_hdf(
            hdf=hdf,
            group_name=group_name
        )
        with self.project_hdf5.open("input") as h5in:
            self.input.to_hdf(h5in)

    def from_hdf(self, hdf=None, group_name=None): 
        super().from_hdf(
            hdf=hdf,
            group_name=group_name
        )
        with self.project_hdf5.open("input") as h5in:
            self.input.from_hdf(h5in)


# In[4]:


from pyiron_base import Project
pr = Project('test')
pr.remove_jobs_silently(recursive=True)
job = pr.create_job(job_type=ToyJob, job_name="toy")
print(job.input)
job.run()
job['output/generic/energy_tot']