Implement the collaborator enchantment

src/main/java/com/minelittlepony/unicopia/item/enchantment/AttributedEnchantment.java

@@ -2,8 +2,6 @@ package com.minelittlepony.unicopia.item.enchantment;
 
 import java.util.HashMap;
 import java.util.Map;
-import java.util.function.IntFunction;
-
 import com.minelittlepony.unicopia.entity.Living;
 
 import net.minecraft.enchantment.EnchantmentTarget;
@@ -15,7 +13,7 @@ import net.minecraft.entity.attribute.EntityAttributeModifier;
 
 public class AttributedEnchantment extends SimpleEnchantment {
 
-    private final Map<EntityAttribute, IntFunction<EntityAttributeModifier>> modifiers = new HashMap<>();
+    private final Map<EntityAttribute, ModifierFactory> modifiers = new HashMap<>();
 
     protected AttributedEnchantment(Rarity rarity, EnchantmentTarget target, boolean cursed, int maxLevel, EquipmentSlot... slots) {
         super(rarity, target, cursed, maxLevel, slots);
@@ -25,7 +23,7 @@ public class AttributedEnchantment extends SimpleEnchantment {
         super(rarity, cursed, maxLevel, slots);
     }
 
-    public AttributedEnchantment addModifier(EntityAttribute attribute, IntFunction<EntityAttributeModifier> modifierSupplier) {
+    public AttributedEnchantment addModifier(EntityAttribute attribute, ModifierFactory modifierSupplier) {
         modifiers.put(attribute, modifierSupplier);
         return this;
     }
@@ -37,7 +35,7 @@ public class AttributedEnchantment extends SimpleEnchantment {
             modifiers.forEach((attr, modifierSupplier) -> {
                 EntityAttributeInstance instance = entity.getAttributeInstance(attr);
 
-                EntityAttributeModifier modifier = modifierSupplier.apply(level);
+                EntityAttributeModifier modifier = modifierSupplier.get(user, level);
 
                 instance.removeModifier(modifier.getId());
                 instance.addPersistentModifier(modifier);
@@ -51,7 +49,7 @@ public class AttributedEnchantment extends SimpleEnchantment {
         modifiers.forEach((attr, modifierSupplier) -> {
             EntityAttributeInstance instance = entity.getAttributeInstance(attr);
 
-            instance.tryRemoveModifier(modifierSupplier.apply(1).getId());
+            instance.tryRemoveModifier(modifierSupplier.get(user, 1).getId());
         });
         user.getEnchants().remove(this);
     }
@@ -59,4 +57,8 @@ public class AttributedEnchantment extends SimpleEnchantment {
     protected boolean shouldChangeModifiers(Living<?> user, int level) {
         return user.getEnchants().computeIfAbsent(this, Data::new).update(level);
     }
+
+    interface ModifierFactory {
+        EntityAttributeModifier get(Living<?> user, int level);
+    }
 }

src/main/java/com/minelittlepony/unicopia/item/enchantment/CollaboratorEnchantment.java

@@ -0,0 +1,34 @@
+package com.minelittlepony.unicopia.item.enchantment;
+
+import java.util.UUID;
+import com.minelittlepony.unicopia.entity.Living;
+import net.minecraft.enchantment.EnchantmentHelper;
+import net.minecraft.enchantment.EnchantmentTarget;
+import net.minecraft.entity.EquipmentSlot;
+import net.minecraft.entity.LivingEntity;
+import net.minecraft.entity.attribute.EntityAttributeModifier;
+import net.minecraft.entity.attribute.EntityAttributes;
+import net.minecraft.entity.attribute.EntityAttributeModifier.Operation;
+
+public class CollaboratorEnchantment extends AttributedEnchantment {
+    protected CollaboratorEnchantment() {
+        super(Rarity.COMMON, EnchantmentTarget.WEAPON, false, 3, EquipmentSlot.MAINHAND);
+        this.addModifier(EntityAttributes.GENERIC_ATTACK_DAMAGE, this::getModifier);
+    }
+
+    @Override
+    protected boolean shouldChangeModifiers(Living<?> user, int level) {
+        return super.shouldChangeModifiers(user, getTeamCollectiveLevel(user, 2 + (level * 2)));
+    }
+
+    private EntityAttributeModifier getModifier(Living<?> user, int level) {
+        return new EntityAttributeModifier(UUID.fromString("5f08c02d-d959-4763-ac84-16e2acfd4b62"), "Team Strength", user.getEnchants().computeIfAbsent(this, Data::new).level / 2, Operation.ADDITION);
+    }
+
+    private int getTeamCollectiveLevel(Living<?> user, int radius) {
+        return user.findAllEntitiesInRange(radius, e -> e instanceof LivingEntity)
+                .mapToInt(e -> EnchantmentHelper.getEquipmentLevel(this, (LivingEntity)e))
+                .reduce((a, b) -> a + b)
+                .orElse(0);
+    }
+}

src/main/java/com/minelittlepony/unicopia/item/enchantment/UEnchantments.java

@@ -32,14 +32,9 @@ public interface UEnchantments {
 
     /**
      * Heavy players move more slowly but are less likely to be flung around wildly.
-     *
-     * TODO:
      */
     Enchantment HEAVY = register("heavy", new AttributedEnchantment(Rarity.COMMON, EnchantmentTarget.ARMOR_FEET, false, 4, EquipmentSlot.FEET))
-            .addModifier(EntityAttributes.GENERIC_MOVEMENT_SPEED, level -> {
+            .addModifier(EntityAttributes.GENERIC_MOVEMENT_SPEED, (user, level) -> {
-                // 1 -> 0.9
-                // 2 -> 0.8
-                // 3 -> 0.7
                 return new EntityAttributeModifier(UUID.fromString("a3d5a94f-4c40-48f6-a343-558502a13e10"), "Heavyness", (1 - level/(float)10) - 1, Operation.MULTIPLY_TOTAL);
             });
 
@@ -47,10 +42,8 @@ public interface UEnchantments {
      * It's dangerous to go alone, take this!
      *
      * Weapons will become stronger the more allies you have around.
-     *
-     * TODO:
      */
-    Enchantment COLLABORATOR = register("collaborator", new SimpleEnchantment(Rarity.COMMON, EnchantmentTarget.WEAPON, false, 1, EquipmentSlot.MAINHAND));
+    Enchantment COLLABORATOR = register("collaborator", new CollaboratorEnchantment());
 
     /**
      * I want it, I neeeed it!