diff --git a/src/main/java/com/kasetoatz/elytramace/ElytraMace.java b/src/main/java/com/kasetoatz/elytramace/ElytraMace.java
index 2b4faf5..1550884 100644
--- a/src/main/java/com/kasetoatz/elytramace/ElytraMace.java
+++ b/src/main/java/com/kasetoatz/elytramace/ElytraMace.java
@@ -7,14 +7,13 @@ import net.minecraft.entity.attribute.EntityAttributes;
 import net.minecraft.entity.player.PlayerInventory;
 import net.minecraft.item.*;
 import net.minecraft.registry.tag.ItemTags;
-import net.minecraft.screen.slot.SlotActionType;
+import net.minecraft.registry.tag.TagKey;
+import net.minecraft.util.Hand;
 
 import java.util.Comparator;
 
 public class ElytraMace implements ClientModInitializer
 {
-    private boolean unequipped = false;
-
     private int getItem(PlayerInventory inventory, Item item)
     {
         for (int i = 0; i < 9; i++)
@@ -27,7 +26,7 @@ public class ElytraMace implements ClientModInitializer
         return -1;
     }
 
-    private int getAxe(PlayerInventory inventory)
+    private int getBestItem(PlayerInventory inventory, TagKey<Item> tag)
     {
         int index = -1;
         int durability = -1;
@@ -35,7 +34,7 @@ public class ElytraMace implements ClientModInitializer
         {
             ItemStack stack = inventory.getStack(i);
             int maxDamage = stack.getMaxDamage();
-            if (stack.isIn(ItemTags.AXES) && maxDamage > durability)
+            if (stack.isIn(tag) && maxDamage > durability)
             {
                 index = i;
                 durability = maxDamage;
@@ -53,31 +52,45 @@ public class ElytraMace implements ClientModInitializer
                 return;
             }
             LivingEntity target = client.world.getEntitiesByClass(LivingEntity.class, client.player.getBoundingBox().expand(client.player.getAttributeValue(EntityAttributes.ENTITY_INTERACTION_RANGE) + client.player.getVelocity().length() * 2), e -> e != client.player).stream().min(Comparator.comparingDouble(client.player::squaredDistanceTo)).orElse(null);
-            if (unequipped && !client.player.isGliding())
+            if (target == null || client.player.fallDistance < 1.5)
             {
-                client.interactionManager.clickSlot(client.player.playerScreenHandler.syncId, 6, 0, SlotActionType.PICKUP, client.player);
-                unequipped = false;
-                if (target != null)
+                return;
+            }
+            PlayerInventory inventory = client.player.getInventory();
+            if (client.player.isGliding())
+            {
+                int chestplate = getBestItem(inventory, ItemTags.CHEST_ARMOR);
+                if (chestplate == -1)
                 {
-                    PlayerInventory inventory = client.player.getInventory();
-                    int axe = getAxe(inventory);
-                    int mace = getItem(inventory, Items.MACE);
-                    int fireworks = getItem(inventory, Items.FIREWORK_ROCKET);
-                    if (target.isBlocking() && axe > -1)
-                    {
-                        inventory.setSelectedSlot(axe);
-                        client.interactionManager.attackEntity(client.player, target);
-                    }
-                    inventory.setSelectedSlot(mace);
+                    return;
+                }
+                inventory.setSelectedSlot(chestplate);
+                client.interactionManager.interactItem(client.player, Hand.MAIN_HAND);
+            }
+            else
+            {
+                int mace = getItem(inventory, Items.MACE);
+                int elytra = getItem(inventory, Items.ELYTRA);
+                if (mace == -1 || elytra == -1)
+                {
+                    return;
+                }
+                inventory.setSelectedSlot(elytra);
+                client.interactionManager.interactItem(client.player, Hand.MAIN_HAND);
+                int axe = getBestItem(inventory, ItemTags.AXES);
+                if (target.isBlocking() && axe > -1)
+                {
+                    inventory.setSelectedSlot(axe);
                     client.interactionManager.attackEntity(client.player, target);
+                }
+                inventory.setSelectedSlot(mace);
+                client.interactionManager.attackEntity(client.player, target);
+                int fireworks = getItem(inventory, Items.FIREWORK_ROCKET);
+                if (fireworks != -1)
+                {
                     inventory.setSelectedSlot(fireworks);
                 }
             }
-            if (!unequipped && target != null && client.player.isGliding())
-            {
-                client.interactionManager.clickSlot(client.player.playerScreenHandler.syncId, 6, 0, SlotActionType.PICKUP, client.player);
-                unequipped = true;
-            }
         });
     }
 }