diff --git a/src/main/java/no/ntnu/idatt1002/demo/controller/BudgetController.java b/src/main/java/no/ntnu/idatt1002/demo/controller/BudgetController.java
index 0144b420673cc2db9517d6bdc45ed1aeada7215f..ada81e37cbba8617ae16a07fb1ebab31c2c7bbcb 100644
--- a/src/main/java/no/ntnu/idatt1002/demo/controller/BudgetController.java
+++ b/src/main/java/no/ntnu/idatt1002/demo/controller/BudgetController.java
@@ -98,75 +98,51 @@ public class BudgetController implements FinanceController {
         amountCol.setCellValueFactory(new PropertyValueFactory<BudgetItem, Double>("budgetAmount"));
         descriptionCol.setCellValueFactory(new PropertyValueFactory<BudgetItem, String>("budgetDescription"));
 
-        /*FileHandlingBudget fileHandlingBudget = new FileHandlingBudget();
         try {
-            if (fileHandlingBudget.isEmpty("Budget")) {
-
+            general = loadBudgetDataFromFile("Budget");
+            budgetList = FXCollections.observableArrayList(general.getBudgetItems());
+            budgetTableView.setItems(budgetList);
+            if (FileHandlingBudget.isNewBudget("Budget")) {
                 returnBtn.setOpacity(0);
+                returnBtn.setDisable(true);
             } else {
-                addBtn.setDisable(true);
-                editBtn.setDisable(true);
-                deleteBtn.setDisable(true);
+                refreshPieChart();
+                //addBtn.setDisable(true);
+                //editBtn.setDisable(true);
+                //deleteBtn.setDisable(true);
+                backBtn.setDisable(true);
+                continueBtn.setDisable(true);
                 backBtn.setOpacity(0);
                 continueBtn.setOpacity(0);
             }
-
         } catch(IOException ioe) {
             showErrorDialogBox("File reading error", "Error in reading file", "Could not"
                 + "read from the Budget file");
-        }*/
-
-        try { //Initialize registers, tableview and pie charts
-            //if (FileHandlingBudget.hasBudgetSet - from line 4/5, so after the general amount and such) {
-            //add, edit and delete buttons are removed.
-
-            //hvis linje fire er tom, instansier et nytt register, hvis ikke, last ned fra fil
-            general = loadBudgetDataFromFile("Budget");
-            budgetList = FXCollections.observableArrayList(general.getBudgetItems());
-            budgetTableView.setItems(budgetList);
-            refreshPieChart();
-        } catch(IOException ioe) {
-            showErrorDialogBox("File reading error", "Error in reading from file", "Could not read"
-                + " registers from file");
         }
-
-        setButtons();
         formatDatePicker();
-
-        //Initialize sum field under the tableview
-        //sum.setText(String.valueOf(general.totalSum()));
     }
 
-    /**
-     * Method that disables or removes button, based on where the budget window
-     * is being used.
-     */
-    private void setButtons() {
-        try {
-            if (FileHandlingBudget.isEmpty("Budget")) {
-                returnBtn.setOpacity(0);
-            } else {
-                addBtn.setDisable(true);
-                editBtn.setDisable(true);
-                deleteBtn.setDisable(true);
-                backBtn.setOpacity(0);
-                continueBtn.setOpacity(0);
-            }
+    private ObservableList<PieChart.Data> createBudgetPieChart() throws IllegalArgumentException { //
+        ObservableList<PieChart.Data> budgetData = FXCollections.observableArrayList();
+        List<ExpenseCategory> chosenCategories = general.getChosenBudgetCategories();
 
-        } catch(IOException ioe) {
-            showErrorDialogBox("File reading error", "Error in reading file", "Could not"
-                + "read from the Budget file");
+        //Only adds the budget data for chosen categories to the pie chart
+        for (ExpenseCategory category : chosenCategories) {
+            budgetData.add(new Data(category.toString().substring(0, 1).toUpperCase().
+                concat(category.toString().substring(1)),
+                general.getBudgetItem(category).getBudgetAmount()));
         }
+        return budgetData;
     }
 
-    private ObservableList<PieChart.Data> createBudgetPieChart() { //
+    /*private ObservableList<PieChart.Data> createBudgetPieChart() {
         return FXCollections.observableArrayList(
             new Data("Food", general.getBudgetItem(ExpenseCategory.FOOD).getBudgetAmount()),
             new Data("Books", general.getBudgetItem(ExpenseCategory.BOOKS).getBudgetAmount()),
             new Data("Clothes", general.getBudgetItem(ExpenseCategory.CLOTHES).getBudgetAmount()),
             new Data("Other", general.getBudgetItem(ExpenseCategory.OTHER).getBudgetAmount())
         );
-    }
+    }*/
 
     private void refreshPieChart() {
         this.budgetPieChart.setData(createBudgetPieChart());
@@ -238,10 +214,7 @@ public class BudgetController implements FinanceController {
             try {
                 general.addToBudgetBudgetItem(item);
             } catch(IllegalArgumentException e) {
-                String title = "Budget amount exceeded/Category already exists";
-                String header = "Your budget exceeds the max limit OR a budget item of the same category already exists in the table";
-                String content = "The total budget sum must be below " + general.getMaxAmount() + " OR Each category can only have one entry in the budget table";
-                showErrorDialogBox(title, header, content);
+                showErrorDialogBox(e.getMessage(), e.getMessage(), e.getMessage());
             }
         }
         //Updates the tableview using the register
diff --git a/src/main/java/no/ntnu/idatt1002/demo/data/Budget/GeneralBudget.java b/src/main/java/no/ntnu/idatt1002/demo/data/Budget/GeneralBudget.java
index 94d2a707c59d04361979f6669949034503d5a407..c93ec12c05e4e238bfd28980ba9f9558b252bf90 100644
--- a/src/main/java/no/ntnu/idatt1002/demo/data/Budget/GeneralBudget.java
+++ b/src/main/java/no/ntnu/idatt1002/demo/data/Budget/GeneralBudget.java
@@ -194,6 +194,17 @@ public class GeneralBudget {
         return budgetItems.stream().anyMatch(budgetItem -> budgetItem.getBudgetCategory().equals(category));
     }
 
+    /**
+     * Method that gets the chosen budget categories. Returns the chosen categories in a list.
+
+     * @return A list of the chosen budget categories.
+     */
+    public List<ExpenseCategory> getChosenBudgetCategories() {
+        return Arrays.stream(ExpenseCategory.values()).toList().
+            stream().filter(this::hasBudgetCategory).toList();
+    }
+
+
     /**
      * This method checks if there are unused categories in the budget.