diff --git a/model_data/optimiser/Optimiser.py b/model_data/optimiser/Optimiser.py
index 5adce91d..e6766172 100644
--- a/model_data/optimiser/Optimiser.py
+++ b/model_data/optimiser/Optimiser.py
@@ -33,41 +33,65 @@ C = 4000
 # group all the parts
 groups = [wall, floor, roof]
 
-# Initialize Model
-m = Model("knapsack")
-
-# Create variables
-vars = [[m.add_var(var_type=BINARY, name=str(component["id"])) for component in group] for group in groups]
-
-# Set objective
-# This objective is the sum
-# gain_ig * x_ig, where gain_ig represents the gain for ith part in group g
-# and x_ig is the binary decision variable for the ith part in group g
-m.objective = maximize(
-    xsum(
-        component['gain'] * var for group, group_vars in zip(groups, vars) for component, var in zip(group, group_vars)
-    )
-)
-
-# Add constraints
-# This constrain ensures that sum of cost_ig * x_ig <= C, where cost_ig represents the cost for the ith component
-# in group g, and x_ig is the binary decision variable for the ith component in group g
-m += xsum(item['cost'] * var for group, group_vars in zip(groups, vars) for item, var in zip(group, group_vars)) <= C
-
-# At most one item from each group
-# This constraint ensures that at most one item from each group is selected
-# This is expressed by summing up the decision variables for each group and ensuring that the sum is <= 1
-for group_vars in vars:
-    m += xsum(var for var in group_vars) <= 1
-
-# Solve the problem
-m.optimize()
-
 # Get the selected items
-selected_items = [
-    item for group, group_vars in zip(groups, vars) for item, var in zip(group, group_vars) if var.x >= 0.99
-]
+
 
 total_gain = m.objective.x
 actual_cost = sum([component['cost'] for component in selected_items])
 print("Selected items:", selected_items)
+
+
+class GainOptimiser:
+    """
+    This class is used maximise gain, given a constrained cost
+    """
+
+    def __init__(self, components):
+        self.components = components
+        self.m = None
+        self.variables = []
+        self.solution = []
+
+    def setup(self):
+        # Initialize Model
+        self.m = Model("knapsack")
+
+        # Create variables
+        self.variables = [
+            [self.m.add_var(var_type=BINARY, name=str(component["id"])) for component in group] for group in groups
+        ]
+
+        # Set objective
+        # This objective is the sum
+        # gain_ig * x_ig, where gain_ig represents the gain for ith part in group g
+        # and x_ig is the binary decision variable for the ith part in group g
+        self.m.objective = maximize(
+            xsum(
+                component['gain'] * var for group, group_vars in zip(groups, self.variables) for component, var in
+                zip(group, group_vars)
+            )
+        )
+
+        # Add constraints
+        # This constrain ensures that sum of cost_ig * x_ig <= C, where cost_ig represents the cost for the ith
+        # component
+        # in group g, and x_ig is the binary decision variable for the ith component in group g
+        self.m += xsum(
+            item['cost'] * var for group, group_vars in zip(groups, self.variables) for item, var in
+            zip(group, group_vars)
+        ) <= C
+
+        # At most one item from each group
+        # This constraint ensures that at most one item from each group is selected
+        # This is expressed by summing up the decision variables for each group and ensuring that the sum is <= 1
+        for group_vars in self.variables:
+            self.m += xsum(var for var in group_vars) <= 1
+
+    def solve(self):
+        # Solve the problem
+        self.m.optimize()
+
+        self.solution = [
+            item for group, group_vars in zip(groups, self.variables) for item, var in zip(group, group_vars) if
+            var.x >= 0.99
+        ]