From 84a73242f4aa2376239eb47cb090b9ad60959db9 Mon Sep 17 00:00:00 2001
From: Tom Robiquet <tom-robiquet@hotmail.co.uk>
Date: Wed, 8 May 2024 16:49:38 +0100
Subject: [PATCH] add price after trade calcs

---
 lib/util/amm2.spec.ts | 27 +++++++++++++++++++++++++++
 lib/util/amm2.ts      | 24 ++++++++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/lib/util/amm2.spec.ts b/lib/util/amm2.spec.ts
index 3e6704f6f..8f4bc5c06 100644
--- a/lib/util/amm2.spec.ts
+++ b/lib/util/amm2.spec.ts
@@ -9,6 +9,8 @@ import {
   isValidBuyAmount,
   isValidSellAmount,
   calculateReserveAfterSell,
+  calculateSpotPriceAfterBuy,
+  calculateSpotPriceAfterSell,
 } from "./amm2";
 import { ZTG } from "@zeitgeistpm/sdk";
 
@@ -83,6 +85,31 @@ describe("amm2", () => {
     });
   });
 
+  describe("calculateSpotPriceAfterBuy", () => {
+    test("should work", () => {
+      const spotPrice = calculateSpotPriceAfterBuy(
+        new Decimal(59.00001516623987),
+        new Decimal(144.00003701590745),
+        new Decimal(543.3339883933237),
+        new Decimal(486),
+      );
+
+      expect(spotPrice.toNumber()).toEqual(0.98849704337919602199);
+    });
+  });
+  describe("calculateSpotPriceAfterSell", () => {
+    test("should work", () => {
+      const spotPrice = calculateSpotPriceAfterSell(
+        new Decimal(99.0),
+        new Decimal(108.04431012579187),
+        new Decimal(40),
+        new Decimal(14.27511827415865),
+      );
+
+      expect(spotPrice.toFixed(8)).toEqual("0.31525092");
+    });
+  });
+
   describe("approximateMaxAmountInForBuy", () => {
     //seems like correct number would be 41
     test("should work", () => {
diff --git a/lib/util/amm2.ts b/lib/util/amm2.ts
index d948b4664..f69e7e179 100644
--- a/lib/util/amm2.ts
+++ b/lib/util/amm2.ts
@@ -59,6 +59,30 @@ export const calculateSpotPrice = (
   return new Decimal(0).minus(reserve).div(liquidity).exp();
 };
 
+export const calculateSpotPriceAfterBuy = (
+  initialReserve: Decimal, // amount of asset in the pool
+  liquidity: Decimal, // liqudity parameter of the pool
+  outcomeAssetOut: Decimal,
+  baseAssetAmountIn: Decimal,
+) => {
+  const newReserve = initialReserve.minus(
+    outcomeAssetOut.minus(baseAssetAmountIn),
+  );
+  return calculateSpotPrice(newReserve, liquidity);
+};
+
+export const calculateSpotPriceAfterSell = (
+  initialReserve: Decimal, // amount of asset in the pool
+  liquidity: Decimal, // liqudity parameter of the pool
+  outcomeAssetIn: Decimal,
+  baseAssetAmountOut: Decimal,
+) => {
+  const newReserve = initialReserve.plus(
+    outcomeAssetIn.minus(baseAssetAmountOut),
+  );
+  return calculateSpotPrice(newReserve, liquidity);
+};
+
 export const approximateMaxAmountInForBuy = (
   reserve: Decimal, // amount of asset in the pool
   liquidity: Decimal, // liqudity parameter of the pool