Flat correct when optimizing

cockpit

@@ -46,14 +46,8 @@ def extract_motor_positions(id19_header):
     return {k: float(v) for (k, v) in zip(names, values)}
 
 
-def have_flats(path):
-    large = os.path.join(path, 'fc')
-    small = os.path.join(path, 'fc-small')
-
-    return (os.path.exists(large) and
-            os.path.exists(small) and
-            glob.glob(os.path.join(large, '*.tif')) and
-            glob.glob(os.path.join(small, '*.tif')))
+def have_data(path):
+    return os.path.exists(path) and glob.glob(os.path.join(path, '*.tif'))
 
 
 class Configuration(object):
@@ -216,7 +210,6 @@ class StateMachine(object):
     QUIT = 1
     SYNC = 2
     CLEAN = 3
-    FLATCORRECT = 4
     OPTIMIZE = 5
     QUICK_OPTIMIZE = 6
     RECONSTRUCT = 7
@@ -227,7 +220,6 @@ class StateMachine(object):
         self.transitions = {
             StateMachine.START: collections.OrderedDict(),
             StateMachine.CLEAN: collections.OrderedDict(),
-            StateMachine.FLATCORRECT: collections.OrderedDict(),
             StateMachine.QUIT: collections.OrderedDict(),
             StateMachine.SYNC: collections.OrderedDict(),
             StateMachine.OPTIMIZE: collections.OrderedDict(),
@@ -265,6 +257,9 @@ class Dataset(object):
     def __repr__(self):
         return '<Dataset(prefix={})>'.format(self.prefix)
 
+    def join_path(self, child_path):
+        return os.path.join(self.path, child_path)
+
 
 class Application(object):
     def __init__(self, config):
@@ -439,13 +434,14 @@ class Application(object):
         self.log.highlight("Cleaning {}".format(self.config.destination))
         return True
 
-    def on_flat_correct(self):
-        self.log.highlight("Flat field correction ...")
+    def flat_correct(self, output_path, append=''):
+        self.log.highlight("Generate flat field corrected projections ...")
 
         try:
             info = self.read_info()
             path = self.current.path
-            data = dict(path=path, prefix=self.current.prefix, num=info['TOMO_N'], step=1)
+            data = dict(path=path, prefix=self.current.prefix,
+                        num=info['TOMO_N'], step=1, output=output_path)
 
             cmd = ('tofu flatcorrect --verbose'
                    ' --reduction-mode median'
@@ -453,19 +449,14 @@ class Application(object):
                    ' --darks {path}/darkend0000.edf'
                    ' --flats {path}/ref*_0000.edf'
                    ' --flats2 {path}/ref*_{num}.edf'
+                   ' --output {output}/fc-%04i.tif'
                    ' --number {num}'
                    ' --step {step}'
                    ' --absorptivity'
-                   ' --fix-nan-and-inf'.format(**data))
-
-            large_cmd = cmd + ' --output {path}/fc/fc-%04i.tif'.format(path=path)
-            small_cmd = cmd + ' --resize 2 --output {path}/fc-small/fc-%04i.tif'.format(path=path)
+                   ' --fix-nan-and-inf '.format(**data))
 
-            if not self.run_command(small_cmd):
-                self.log.error("Could not create small flat field corrected projections")
-                return result
-
-            return self.run_command(large_cmd)
+            cmd += append
+            return self.run_command(cmd)
         except Exception as e:
             self.log.error("Error: {}".format(e))
 
@@ -474,6 +465,12 @@ class Application(object):
     def on_quick_optimize(self):
         self.log.highlight("Quick optimization ...")
 
+        path = self.current.join_path('fc-small')
+
+        if not have_data(path):
+            if not self.flat_correct(path, append='--resize 2'):
+                return False
+
         try:
             if self.optimize(True):
                 with open(self.params_name) as f:
@@ -495,6 +492,12 @@ class Application(object):
     def on_optimize(self):
         self.log.highlight("Optimizing ...")
 
+        path = self.current.join_path('fc')
+
+        if not have_data(path):
+            if not self.flat_correct(path):
+                return False
+
         try:
             if self.optimize(False):
                 params = self.read_optimal_params()
@@ -544,7 +547,6 @@ class Application(object):
         quick_optimize = Action('q', 'Quick optimization', self.on_quick_optimize, machine.QUICK_OPTIMIZE)
         optimize = Action('o', 'Optimization', self.on_optimize, machine.OPTIMIZE)
         reconstruct = Action('r', 'Reconstruct', self.on_reconstruct, machine.RECONSTRUCT)
-        flatcorrect = Action('f', 'Flat correct', self.on_flat_correct, machine.FLATCORRECT)
         next_dataset = Action('n', 'Next dataset', self.on_next_dataset, machine.START)
         previous_dataset = Action('p', 'Previous dataset', self.on_previous_dataset, machine.START)
 
@@ -572,12 +574,6 @@ class Application(object):
         machine.add_action(machine.CLEAN, sync)
         machine.add_action(machine.CLEAN, quit)
 
-        machine.add_action(machine.FLATCORRECT, quick_optimize)
-        machine.add_action(machine.FLATCORRECT, optimize)
-        machine.add_action(machine.FLATCORRECT, next_dataset)
-        machine.add_action(machine.FLATCORRECT, previous_dataset)
-        machine.add_action(machine.FLATCORRECT, quit)
-
         self.log = LogList(right_pane, colors)
         self.log.info('Source dir set to {}'.format(self.config.source))
         self.log.info('Destination dir set to {}'.format(self.config.destination))
@@ -588,10 +584,6 @@ class Application(object):
             self.index = 0
             self.current = self.datasets[0]
 
-        if not have_flats(self.current.path):
-            machine.add_action(machine.START, flatcorrect)
-            machine.add_action(machine.SYNC, flatcorrect)
-
         if os.path.exists(self.params_name):
             machine.add_action(machine.START, reconstruct)