From e080084830e77af87891da1f4ff56dcbcf988fd4 Mon Sep 17 00:00:00 2001
From: Sollace <sollacea@gmail.com>
Date: Sun, 8 Aug 2021 19:21:14 +0200
Subject: [PATCH] Fix some axolotl angles

---
 .../entity/behaviour/AxolotlBehaviour.java    | 19 +++++++++++++++++++
 .../entity/behaviour/EntityBehaviour.java     |  1 +
 2 files changed, 20 insertions(+)
 create mode 100644 src/main/java/com/minelittlepony/unicopia/entity/behaviour/AxolotlBehaviour.java

diff --git a/src/main/java/com/minelittlepony/unicopia/entity/behaviour/AxolotlBehaviour.java b/src/main/java/com/minelittlepony/unicopia/entity/behaviour/AxolotlBehaviour.java
new file mode 100644
index 00000000..797dd888
--- /dev/null
+++ b/src/main/java/com/minelittlepony/unicopia/entity/behaviour/AxolotlBehaviour.java
@@ -0,0 +1,19 @@
+package com.minelittlepony.unicopia.entity.behaviour;
+
+import com.minelittlepony.unicopia.ability.magic.Caster;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
+import net.minecraft.entity.passive.AxolotlEntity;
+import net.minecraft.util.math.Vec3f;
+
+public class AxolotlBehaviour extends EntityBehaviour<AxolotlEntity> {
+    @Override
+    public void update(Caster<?> source, AxolotlEntity entity, DisguiseSpell spell) {
+        float toRad = 0.017453292F;
+        Vec3f current = entity.getModelAngles().get("body");
+        entity.getModelAngles().put("body", new Vec3f(
+               source.getEntity().isSubmergedInWater() ? source.getEntity().getPitch() * toRad : 0,
+               0,
+               current == null ? 0 : current.getZ()
+       ));
+    }
+}
diff --git a/src/main/java/com/minelittlepony/unicopia/entity/behaviour/EntityBehaviour.java b/src/main/java/com/minelittlepony/unicopia/entity/behaviour/EntityBehaviour.java
index acd0f31f..8a7007d6 100644
--- a/src/main/java/com/minelittlepony/unicopia/entity/behaviour/EntityBehaviour.java
+++ b/src/main/java/com/minelittlepony/unicopia/entity/behaviour/EntityBehaviour.java
@@ -274,6 +274,7 @@ public class EntityBehaviour<T extends Entity> {
         register(SheepBehaviour::new, EntityType.SHEEP);
         register(BeeBehaviour::new, EntityType.BEE);
         register(GhastBehaviour::new, EntityType.GHAST);
+        register(AxolotlBehaviour::new, EntityType.AXOLOTL);
         register(EndermanBehaviour::new, EntityType.ENDERMAN);
         EntityBehaviour.<LlamaEntity>register(() -> new RangedAttackBehaviour<>(SoundEvents.ENTITY_LLAMA_SPIT, LlamaSpitEntity::new), EntityType.LLAMA, EntityType.TRADER_LLAMA);
         EntityBehaviour.<SnowGolemEntity>register(() -> new RangedAttackBehaviour<>(SoundEvents.ENTITY_SNOW_GOLEM_SHOOT, SnowballEntity::new), EntityType.SNOW_GOLEM);