<gh_stars>0
from sacred import Experiment
import os

try:
    import oyaml as yaml
except ImportError:
    import yaml

from shutil import rmtree

from schnetpack.sacred.calculator_ingredients import (calculator_ingradient,
                                                      build_calculator)
from schnetpack.sacred.simulator_ingredients import (simulator_ingredient,
                                                     build_simulator)
from schnetpack.sacred.integrator_ingredients import (integrator_ingredient,
                                                      build_integrator)
from schnetpack.sacred.system_ingredients import (system_ingredient,
                                                  build_system)
from schnetpack.sacred.thermostat_ingredients import thermostat_ingredient, \
    build_thermostat

md = Experiment('md', ingredients=[simulator_ingredient, calculator_ingradient,
                                   integrator_ingredient, system_ingredient,
                                   thermostat_ingredient])

SETUP_STRING_WIDTH = 30
SETUP_STRING = "\n\n{:s}\n{:s}\n{:s}".format(SETUP_STRING_WIDTH * "=",
                                             f'{{:^{SETUP_STRING_WIDTH}s}}',
                                             SETUP_STRING_WIDTH * '=')


@md.config
def config():
    """configuration for the simulation experiment"""
    experiment_dir = 'experiment'
    simulation_steps = 1000
    device = 'cpu'
    overwrite = True


@md.capture
def save_system_config(_config, experiment_dir):
    """
    Save the configuration to the model directory.

    Args:
        _config (dict): configuration of the experiment
        simulation_dir (str): path to the simulation directory

    """
    with open(os.path.join(experiment_dir, 'config.yaml'), 'w') as f:
        yaml.dump(_config, f, default_flow_style=False)


@md.capture
def setup_simulation(_log, simulation_dir, device):
    _log.info(SETUP_STRING.format('CALCULATOR SETUP'))
    calculator = build_calculator(device=device)
    _log.info(SETUP_STRING.format('SYSTEM SETUP'))
    system = build_system(device=device)
    _log.info(SETUP_STRING.format('INTEGRATOR SETUP'))
    integrator = build_integrator(n_beads=system.n_replicas, device=device)
    _log.info(SETUP_STRING.format('THERMOSTAT SETUP'))
    thermostat = build_thermostat()
    _log.info(SETUP_STRING.format('SIMULATOR SETUP'))
    simulator = build_simulator(system=system,
                                integrator_object=integrator,
                                calculator_object=calculator,
                                simulation_dir=simulation_dir,
                                thermostat_object=thermostat)
    return simulator


@md.capture
def create_dirs(_log, experiment_dir, overwrite):
    """
    Create the directory for the experiment.

    Args:
        _log:
        experiment_dir (str): path to the experiment directory
        overwrite (bool): overwrites the model directory if True

    """

    simulation_dir = os.path.join(experiment_dir, 'simulation')

    _log.info("Create model directory")
    if simulation_dir is None:
        raise ValueError('Config `experiment_dir` has to be set!')

    if os.path.exists(simulation_dir) and not overwrite:
        raise ValueError(
            'Model directory already exists (set overwrite flag?):',
            simulation_dir)

    if os.path.exists(simulation_dir) and overwrite:
        rmtree(simulation_dir)

    if not os.path.exists(simulation_dir):
        os.makedirs(simulation_dir)


@md.command
def simulate(experiment_dir, simulation_steps):
    simulation_dir = os.path.join(experiment_dir, 'simulation')
    create_dirs()
    save_system_config()
    simulator = setup_simulation(simulation_dir=simulation_dir)
    simulator.simulate(simulation_steps)


@md.command
def save_config(_log, _config, experiment_dir):
    file_name = f"{experiment_dir}_config.yaml"
    with open(file_name, 'w') as f:
        yaml.dump(_config, f, default_flow_style=False)
    _log.info(f'Stored config to {file_name}')


@md.automain
def main(_log):
    save_config()
    _log.info('To run simulation call script with "simulate with <config file>"')

from sacred import Experiment
import os

try:
    import oyaml as yaml
except ImportError:
    import yaml

from shutil import rmtree

from schnetpack.sacred.calculator_ingredients import (calculator_ingradient,
                                                      build_calculator)
from schnetpack.sacred.simulator_ingredients import (simulator_ingredient,
                                                     build_simulator)
from schnetpack.sacred.integrator_ingredients import (integrator_ingredient,
                                                      build_integrator)
from schnetpack.sacred.system_ingredients import (system_ingredient,
                                                  build_system)
from schnetpack.sacred.thermostat_ingredients import thermostat_ingredient, \
    build_thermostat

md = Experiment('md', ingredients=[simulator_ingredient, calculator_ingradient,
                                   integrator_ingredient, system_ingredient,
                                   thermostat_ingredient])

SETUP_STRING_WIDTH = 30
SETUP_STRING = "\n\n{:s}\n{:s}\n{:s}".format(SETUP_STRING_WIDTH * "=",
                                             f'{{:^{SETUP_STRING_WIDTH}s}}',
                                             SETUP_STRING_WIDTH * '=')


@md.config
def config():
    """configuration for the simulation experiment"""
    experiment_dir = 'experiment'
    simulation_steps = 1000
    device = 'cpu'
    overwrite = True


@md.capture
def save_system_config(_config, experiment_dir):
    """
    Save the configuration to the model directory.

    Args:
        _config (dict): configuration of the experiment
        simulation_dir (str): path to the simulation directory

    """
    with open(os.path.join(experiment_dir, 'config.yaml'), 'w') as f:
        yaml.dump(_config, f, default_flow_style=False)


@md.capture
def setup_simulation(_log, simulation_dir, device):
    _log.info(SETUP_STRING.format('CALCULATOR SETUP'))
    calculator = build_calculator(device=device)
    _log.info(SETUP_STRING.format('SYSTEM SETUP'))
    system = build_system(device=device)
    _log.info(SETUP_STRING.format('INTEGRATOR SETUP'))
    integrator = build_integrator(n_beads=system.n_replicas, device=device)
    _log.info(SETUP_STRING.format('THERMOSTAT SETUP'))
    thermostat = build_thermostat()
    _log.info(SETUP_STRING.format('SIMULATOR SETUP'))
    simulator = build_simulator(system=system,
                                integrator_object=integrator,
                                calculator_object=calculator,
                                simulation_dir=simulation_dir,
                                thermostat_object=thermostat)
    return simulator


@md.capture
def create_dirs(_log, experiment_dir, overwrite):
    """
    Create the directory for the experiment.

    Args:
        _log:
        experiment_dir (str): path to the experiment directory
        overwrite (bool): overwrites the model directory if True

    """

    simulation_dir = os.path.join(experiment_dir, 'simulation')

    _log.info("Create model directory")
    if simulation_dir is None:
        raise ValueError('Config `experiment_dir` has to be set!')

    if os.path.exists(simulation_dir) and not overwrite:
        raise ValueError(
            'Model directory already exists (set overwrite flag?):',
            simulation_dir)

    if os.path.exists(simulation_dir) and overwrite:
        rmtree(simulation_dir)

    if not os.path.exists(simulation_dir):
        os.makedirs(simulation_dir)


@md.command
def simulate(experiment_dir, simulation_steps):
    simulation_dir = os.path.join(experiment_dir, 'simulation')
    create_dirs()
    save_system_config()
    simulator = setup_simulation(simulation_dir=simulation_dir)
    simulator.simulate(simulation_steps)


@md.command
def save_config(_log, _config, experiment_dir):
    file_name = f"{experiment_dir}_config.yaml"
    with open(file_name, 'w') as f:
        yaml.dump(_config, f, default_flow_style=False)
    _log.info(f'Stored config to {file_name}')


@md.automain
def main(_log):
    save_config()
    _log.info('To run simulation call script with "simulate with <config file>"')
configuration for the simulation experiment Save the configuration to the model directory.      Args:         _config (dict): configuration of the experiment         simulation_dir (str): path to the simulation directory Create the directory for the experiment.      Args:         _log:         experiment_dir (str): path to the experiment directory         overwrite (bool): overwrites the model directory if True