diff --git a/bo_ptycho.py b/bo_ptycho.py
deleted file mode 100644
index d1555e3..0000000
--- a/bo_ptycho.py
+++ /dev/null
@@ -1,94 +0,0 @@
-import os
-import sys
-import yaml
-
-def ptycho_run(config):
-    ptycho_engine = config['ptycho']['engine']
-
-    if ptycho_engine == 'fold_slice':
-        from ptycho.fold_slice import run
-        run(config)
-
-    # Add more ptycho engines here as needed
-    # write their respective run functions and import them above
-
-def ptycho_error(config):
-    ptycho_engine = config['ptycho']['engine']
-
-    if ptycho_engine == 'fold_slice':
-        from ptycho.fold_slice import error
-        return error(config)
-
-def bo_initialize(config):
-    bo_algorithm = config['bo']['algorithm']
-
-    if bo_algorithm == 'ucb':
-        from bo.ucb import initialize
-        bo_state = initialize(config)
-    
-    else:
-        from bo.random import initialize
-        bo_state = initialize(config)
-    
-    return bo_state
-
-def bo_ask(config, bo_state):
-    bo_engine = config['bo']['algorithm']
-
-    if bo_engine == 'ucb':
-        from bo.ucb import ask
-        next_config = ask(config, bo_state)
-    
-    else:
-        from bo.random import ask
-        next_config = ask(config, bo_state)
-
-    # Add more BO engines here as needed
-    # write their respective ask functions and import them above
-
-    return next_config
-
-def bo_tell(config, job_config, bo_state, y_value):
-    bo_engine = config['bo']['algorithm']
-
-    if bo_engine == 'ucb':
-        from bo.ucb import tell
-        bo_state = tell(job_config, bo_state, y_value)
-    
-    else:
-        from bo.random import tell
-        bo_state = tell(config, job_config, bo_state, y_value)
-    
-    return bo_state
-
-def main(config):
-
-    bo_state = bo_initialize(config)
-
-    max_iterations = config['bo']['max_iterations']
-
-    for _ in range(max_iterations):
-        job_config = bo_ask(config, bo_state)
-
-        ptycho_run(job_config)
-        y_value = ptycho_error(job_config)
-
-        bo_state = bo_tell(config, job_config, bo_state, y_value)
-
-        print(bo_state['train_x'])
-        print(bo_state['train_y'])
-
-    return bo_state
-
-
-if __name__ == "__main__":
-    if len(sys.argv) != 2:
-        print("Usage: python bo-ptycho.py <config_yaml>")
-        sys.exit(1)
-
-    config_yaml = sys.argv[1]
-    with open(config_yaml, 'r') as f:
-        config = yaml.safe_load(f)
-    
-    main(config)
-    
diff --git a/main.py b/main.py
new file mode 100644
index 0000000..f7a9f5a
--- /dev/null
+++ b/main.py
@@ -0,0 +1,36 @@
+import os
+import sys
+import yaml
+
+import bo
+import ptycho
+
+def main(config):
+
+    bo_engine = bo.RandomBOEngine(config)
+
+    max_iterations = config['bo']['max_iterations']
+
+    for _ in range(max_iterations):
+        job_config = bo_engine.ask()
+
+        ptycho_engine = ptycho.FoldSlicePtychoEngine(job_config)
+        ptycho_engine.run()
+        y_value = ptycho_engine.metric()
+
+        bo_engine.tell(job_config, y_value)
+
+        print(bo_engine.state['train_x'])
+        print(bo_engine.state['train_y'])
+
+
+if __name__ == "__main__":
+    if len(sys.argv) != 2:
+        print("Usage: python bo-ptycho.py <config_yaml>")
+        sys.exit(1)
+
+    config_yaml = sys.argv[1]
+    with open(config_yaml, 'r') as f:
+        config = yaml.safe_load(f)
+    
+    main(config)