diff --git a/src/main/java/fi/dy/masa/tweakeroo/config/Callbacks.java b/src/main/java/fi/dy/masa/tweakeroo/config/Callbacks.java
index a79589fa..971cb21d 100644
--- a/src/main/java/fi/dy/masa/tweakeroo/config/Callbacks.java
+++ b/src/main/java/fi/dy/masa/tweakeroo/config/Callbacks.java
@@ -407,7 +407,7 @@ else if (key == Hotkeys.OPEN_CONFIG_GUI.getKeybind())
             }
             else if (key == Hotkeys.SWAP_ELYTRA_CHESTPLATE.getKeybind())
             {
-                InventoryUtils.swapElytraWithChestPlate(this.mc.player);
+                InventoryUtils.swapElytraAndChestPlate(this.mc.player);
                 return true;
             }
             else if (key == Hotkeys.TOGGLE_GRAB_CURSOR.getKeybind())
diff --git a/src/main/java/fi/dy/masa/tweakeroo/mixin/MixinClientPlayerEntity.java b/src/main/java/fi/dy/masa/tweakeroo/mixin/MixinClientPlayerEntity.java
index 640c5f8f..2e33b249 100644
--- a/src/main/java/fi/dy/masa/tweakeroo/mixin/MixinClientPlayerEntity.java
+++ b/src/main/java/fi/dy/masa/tweakeroo/mixin/MixinClientPlayerEntity.java
@@ -129,22 +129,14 @@ private void onFallFlyingCheckChestSlot(CallbackInfo ci)
     {
         if (FeatureToggle.TWEAK_AUTO_SWITCH_ELYTRA.getBooleanValue())
         {
-            // Sakura's version calculating fall distance...
-            //if ((!this.getEquippedStack(EquipmentSlot.CHEST).isOf(Items.ELYTRA) && this.fallDistance > 20.0f)
-
-            // Auto switch if it is not elytra after falling, or is totally broken.
-            // This also shouldn't activate on the Ground if the Chest Equipment is EMPTY,
-            // or not an Elytra to be swapped back.
-            //
-            // !isOnGround(): Minecraft also check elytra even if the player is on the ground, skip it.
-            // !isSwimming(): Check if the Player is swimming so, it doesn't perform the Elytra Swap while underwater.
-            if (!this.isOnGround() && !this.isSwimming()
-                && (!this.getEquippedStack(EquipmentSlot.CHEST).isOf(Items.ELYTRA))
-                || (this.getEquippedStack(EquipmentSlot.CHEST).getDamage() > this.getEquippedStack(EquipmentSlot.CHEST).getMaxDamage() - 10)
-                && (!this.getEquippedStack(EquipmentSlot.CHEST).isEmpty() || this.autoSwitchElytraChestplate.isOf(Items.ELYTRA)))
+            // PlayerEntity#checkFallFlying
+            if (!this.isOnGround() && !this.isFallFlying() && !this.isInFluid() && !this.hasStatusEffect(StatusEffects.LEVITATION))
             {
-                this.autoSwitchElytraChestplate = this.getEquippedStack(EquipmentSlot.CHEST).copy();
-                InventoryUtils.swapElytraWithChestPlate(this);
+                if (!this.getEquippedStack(EquipmentSlot.CHEST).isOf(Items.ELYTRA) ||
+                    this.getEquippedStack(EquipmentSlot.CHEST).getDamage() > this.getEquippedStack(EquipmentSlot.CHEST).getMaxDamage() - 10)
+                {
+                    InventoryUtils.equipBestElytra(this);
+                }
             }
         }
         else
@@ -180,7 +172,7 @@ private void onStopFlying(TrackedData<?> data, CallbackInfo ci)
                     else
                     {
                         // if cached previous item is empty, try to swap back to the default chest plate.
-                        InventoryUtils.swapElytraWithChestPlate(this);
+                        InventoryUtils.swapElytraAndChestPlate(this);
                     }
                 }
             }
diff --git a/src/main/java/fi/dy/masa/tweakeroo/util/InventoryUtils.java b/src/main/java/fi/dy/masa/tweakeroo/util/InventoryUtils.java
index bf120f7c..d190371c 100644
--- a/src/main/java/fi/dy/masa/tweakeroo/util/InventoryUtils.java
+++ b/src/main/java/fi/dy/masa/tweakeroo/util/InventoryUtils.java
@@ -865,7 +865,7 @@ private static int findRepairableItemNotInRepairableSlot(Slot targetSlot, Player
         return -1;
     }
 
-    public static void swapElytraWithChestPlate(@Nullable PlayerEntity player)
+    public static void equipBestElytra(PlayerEntity player)
     {
         if (player == null || GuiUtils.getCurrentScreen() != null)
         {
@@ -873,50 +873,69 @@ public static void swapElytraWithChestPlate(@Nullable PlayerEntity player)
         }
 
         ScreenHandler container = player.currentScreenHandler;
-        ItemStack currentStack = player.getEquippedStack(EquipmentSlot.CHEST);
-
-        Predicate<ItemStack> stackFilterChestPlate = (s) -> s.getItem() instanceof ArmorItem && ((ArmorItem) s.getItem()).getSlotType() == EquipmentSlot.CHEST;
-        Predicate<ItemStack> stackFilterElytra = (s) -> s.getItem() instanceof ElytraItem && ElytraItem.isUsable(s);
-        boolean switchingToElytra = (currentStack.isEmpty() || stackFilterChestPlate.test(currentStack));
-        Predicate<ItemStack> stackFilter = switchingToElytra ? stackFilterElytra : stackFilterChestPlate;
-        Predicate<ItemStack> finalFilter = (s) -> s.isEmpty() == false && stackFilter.test(s) && s.getDamage() < s.getMaxDamage() - 10;
 
+        Predicate<ItemStack> filter = (s) ->  s.getItem() instanceof ElytraItem && ElytraItem.isUsable(s) && s.getDamage() < s.getMaxDamage() - 10;
         int targetSlot = findSlotWithBestItemMatch(container, (testedStack, previousBestMatch) -> {
-            if (!finalFilter.test(testedStack)) return false;
-            if (!finalFilter.test(previousBestMatch)) return true;
-            if (switchingToElytra)
+            if (!filter.test(testedStack)) return false;
+            if (!filter.test(previousBestMatch)) return true;
+            if (getEnchantmentLevel(testedStack, Enchantments.UNBREAKING) > getEnchantmentLevel(previousBestMatch, Enchantments.UNBREAKING))
             {
-                if (getEnchantmentLevel(testedStack, Enchantments.UNBREAKING) < getEnchantmentLevel(previousBestMatch, Enchantments.UNBREAKING))
-                {
-                    return false;
-                }
-                if (testedStack.getDamage() > previousBestMatch.getDamage())
-                {
-                     return false;
-                }
+                return true;
             }
-            else
+            if (getEnchantmentLevel(testedStack, Enchantments.UNBREAKING) < getEnchantmentLevel(previousBestMatch, Enchantments.UNBREAKING))
             {
-                if (getArmorAndArmorToughnessValue(previousBestMatch, 1, AttributeModifierSlot.CHEST) > getArmorAndArmorToughnessValue(testedStack, 1, AttributeModifierSlot.CHEST))
-                {
-                    return false;
-                }
-                if (getEnchantmentLevel(previousBestMatch, Enchantments.PROTECTION) > getEnchantmentLevel(testedStack, Enchantments.PROTECTION))
-                {
-                    return false;
-                }
+                return false;
             }
-
-            return true;
+            return testedStack.getDamage() <= previousBestMatch.getDamage();
         }, UniformIntProvider.create(9, container.slots.size() - 1));
 
         if (targetSlot >= 0)
         {
-            //targetSlots.sort();
             swapItemToEquipmentSlot(player, EquipmentSlot.CHEST, targetSlot);
         }
     }
 
+    public static void swapElytraAndChestPlate(@Nullable PlayerEntity player)
+    {
+        if (player == null || GuiUtils.getCurrentScreen() != null)
+        {
+            return;
+        }
+
+        ScreenHandler container = player.currentScreenHandler;
+        ItemStack currentStack = player.getEquippedStack(EquipmentSlot.CHEST);
+
+        Predicate<ItemStack> stackFilterChestPlate = (s) -> s.getItem() instanceof ArmorItem && ((ArmorItem) s.getItem()).getSlotType() == EquipmentSlot.CHEST;
+
+        if (currentStack.isEmpty() || stackFilterChestPlate.test(currentStack))
+        {
+            equipBestElytra(player);
+        }
+        else
+        {
+            Predicate<ItemStack> finalFilter = (s) -> stackFilterChestPlate.test(s) && s.getDamage() < s.getMaxDamage() - 10;
+
+            int targetSlot = findSlotWithBestItemMatch(container, (testedStack, previousBestMatch) -> {
+                if (!finalFilter.test(testedStack)) return false;
+                if (!finalFilter.test(previousBestMatch)) return true;
+                if (getArmorAndArmorToughnessValue(previousBestMatch, 1, AttributeModifierSlot.CHEST) < getArmorAndArmorToughnessValue(testedStack, 1, AttributeModifierSlot.CHEST))
+                {
+                    return true;
+                }
+                if (getArmorAndArmorToughnessValue(previousBestMatch, 1, AttributeModifierSlot.CHEST) > getArmorAndArmorToughnessValue(testedStack, 1, AttributeModifierSlot.CHEST))
+                {
+                    return false;
+                }
+                return getEnchantmentLevel(previousBestMatch, Enchantments.PROTECTION) <= getEnchantmentLevel(testedStack, Enchantments.PROTECTION);
+            }, UniformIntProvider.create(9, container.slots.size() - 1));
+
+            if (targetSlot >= 0)
+            {
+                swapItemToEquipmentSlot(player, EquipmentSlot.CHEST, targetSlot);
+            }
+        }
+    }
+
     private static double getArmorAndArmorToughnessValue(ItemStack stack, double base, AttributeModifierSlot slot)
     {
         final double[] total = {base};