diff --git a/python/src/pools/stable_math.py b/python/src/pools/stable_math.py
index 6f302f9..79fb4c8 100644
--- a/python/src/pools/stable_math.py
+++ b/python/src/pools/stable_math.py
@@ -163,8 +163,8 @@ def compute_balance(
     # Calculate inv2 and c
     inv2 = invariant * invariant
 
-    c = div_up(inv2, amp_times_total * p_d) * AMP_PRECISION * balances[token_index]
-    b = sum_balances + (invariant // amp_times_total) * AMP_PRECISION
+    c = div_up(inv2 * AMP_PRECISION, amp_times_total * p_d) * balances[token_index]
+    b = sum_balances + (invariant * AMP_PRECISION) // amp_times_total
     # Initial approximation of tokenBalance
     prev_token_balance = int(0)
     token_balance = div_up(inv2 + c, invariant + b)
diff --git a/typescript/package.json b/typescript/package.json
index 58d626f..44f9864 100644
--- a/typescript/package.json
+++ b/typescript/package.json
@@ -9,7 +9,7 @@
     "publishConfig": {
         "access": "public"
     },
-    "version": "0.0.14",
+    "version": "0.0.15",
     "main": "dist/index.js",
     "module": "dist/index.mjs",
     "types": "dist/index.d.ts",
diff --git a/typescript/src/hooks/directionalFeeHook.ts b/typescript/src/hooks/directionalFeeHook.ts
new file mode 100644
index 0000000..0a4df58
--- /dev/null
+++ b/typescript/src/hooks/directionalFeeHook.ts
@@ -0,0 +1,111 @@
+import { HookBase } from './types';
+import { MathSol } from '../utils/math';
+import { SwapInput } from '@/vault/types';
+
+interface MinimalToken {
+    address: string;
+    decimals: bigint;
+    index: bigint;
+}
+
+export type HookStateDirectionalFee = {
+    tokens: MinimalToken[];
+    balancesLiveScaled18: bigint[];
+};
+
+export class DirectionalFeeHook implements HookBase {
+    public shouldCallComputeDynamicSwapFee = true;
+    public shouldCallBeforeSwap = false;
+    public shouldCallAfterSwap = false;
+    public shouldCallBeforeAddLiquidity = false;
+    public shouldCallAfterAddLiquidity = false;
+    public shouldCallBeforeRemoveLiquidity = false;
+    public shouldCallAfterRemoveLiquidity = false;
+    public enableHookAdjustedAmounts = false;
+
+    onComputeDynamicSwapFee(
+        params: SwapInput,
+        staticSwapFeePercentage: bigint,
+        hookState: HookStateDirectionalFee,
+    ): { success: boolean; dynamicSwapFee: bigint } {
+        const lastBalancesLiveScaled18 = hookState.balancesLiveScaled18;
+
+        const tokenInIndex = hookState.tokens
+            .map((token) => token.address.toLowerCase())
+            .indexOf(params.tokenIn.toLowerCase());
+        const tokenOutIndex = hookState.tokens
+            .map((token) => token.address.toLowerCase())
+            .indexOf(params.tokenOut.toLowerCase());
+
+        const decimalsToScaleWith =
+            params.swapKind === 0
+                ? 18n - hookState.tokens[tokenInIndex].decimals
+                : 18n - hookState.tokens[tokenOutIndex].decimals;
+        const amountsScaled18 = params.amountRaw * 10n ** decimalsToScaleWith;
+
+        const calculatedSwapFeePercentage =
+            this.calculateExpectedSwapFeePercentage(
+                lastBalancesLiveScaled18,
+                amountsScaled18,
+                tokenInIndex,
+                tokenOutIndex,
+            );
+
+        // Charge the static or calculated fee, whichever is greater.
+        const dynamicSwapFee =
+            calculatedSwapFeePercentage > staticSwapFeePercentage
+                ? calculatedSwapFeePercentage
+                : staticSwapFeePercentage;
+
+        return {
+            success: true,
+            dynamicSwapFee: dynamicSwapFee,
+        };
+    }
+
+    // the bigger the swap ( relative to pool size ) the bigger the fee
+    private calculateExpectedSwapFeePercentage(
+        poolBalances: bigint[],
+        swapAmount: bigint,
+        indexIn: number,
+        indexOut: number,
+    ): bigint {
+        const finalBalanceTokenIn = poolBalances[indexIn] + swapAmount;
+        const finalBalanceTokenOut = poolBalances[indexOut] - swapAmount;
+
+        if (finalBalanceTokenIn > finalBalanceTokenOut) {
+            const diff = finalBalanceTokenIn - finalBalanceTokenOut;
+            const totalLiquidity = finalBalanceTokenIn + finalBalanceTokenOut;
+
+            return MathSol.divDownFixed(diff, totalLiquidity);
+        }
+
+        return 0n;
+    }
+
+    onBeforeAddLiquidity() {
+        return { success: false, hookAdjustedBalancesScaled18: [] };
+    }
+
+    onAfterAddLiquidity() {
+        return { success: false, hookAdjustedAmountsInRaw: [] };
+    }
+
+    onBeforeRemoveLiquidity() {
+        return { success: false, hookAdjustedBalancesScaled18: [] };
+    }
+
+    onAfterRemoveLiquidity() {
+        return { success: false, hookAdjustedAmountsOutRaw: [] };
+    }
+
+    onBeforeSwap() {
+        return { success: false, hookAdjustedBalancesScaled18: [] };
+    }
+
+    onAfterSwap() {
+        return { success: false, hookAdjustedAmountCalculatedRaw: 0n };
+    }
+}
+
+export default DirectionalFeeHook;
diff --git a/typescript/src/hooks/types.ts b/typescript/src/hooks/types.ts
index e3c5d62..e9bbee9 100644
--- a/typescript/src/hooks/types.ts
+++ b/typescript/src/hooks/types.ts
@@ -1,7 +1,8 @@
 import { AddKind, RemoveKind, SwapInput, SwapKind } from '@/vault/types';
 import { HookStateExitFee } from './exitFeeHook';
+import { HookStateDirectionalFee } from './directionalFeeHook';
 
-export type HookState = HookStateExitFee;
+export type HookState = HookStateExitFee | HookStateDirectionalFee;
 
 export type AfterSwapParams = {
     kind: SwapKind;
diff --git a/typescript/src/stable/stableMath.ts b/typescript/src/stable/stableMath.ts
index e541b2f..66cd39d 100644
--- a/typescript/src/stable/stableMath.ts
+++ b/typescript/src/stable/stableMath.ts
@@ -181,10 +181,10 @@ export function _computeBalance(
     const inv2 = invariant * invariant;
     // We remove the balance from c by multiplying it.
     const c =
-        MathSol.divUp(inv2, ampTimesTotal * P_D) *
-        AMP_PRECISION *
+        MathSol.divUp(inv2 * AMP_PRECISION, ampTimesTotal * P_D) *
         balances[tokenIndex];
-    const b = sum + (invariant / ampTimesTotal) * AMP_PRECISION;
+
+    const b = sum + (invariant * AMP_PRECISION) / ampTimesTotal;
     // We iterate to find the balance.
     let prevTokenBalance = 0n;
     // We multiply the first iteration outside the loop with the invariant to set the value of the
diff --git a/typescript/src/vault/vault.ts b/typescript/src/vault/vault.ts
index 459036a..0cd6e73 100644
--- a/typescript/src/vault/vault.ts
+++ b/typescript/src/vault/vault.ts
@@ -26,6 +26,7 @@ import {
 import { HookBase, HookClassConstructor, HookState } from '../hooks/types';
 import { defaultHook } from '../hooks/constants';
 import { ExitFeeHook } from '../hooks/exitFeeHook';
+import { DirectionalFeeHook } from '../hooks/directionalFeeHook';
 
 const _MINIMUM_TRADE_AMOUNT = 1e6;
 // const _MINIMUM_WRAP_AMOUNT = 1e3;
@@ -53,6 +54,7 @@ export class Vault {
         };
         this.hookClasses = {
             ExitFee: ExitFeeHook,
+            DirectionalFee: DirectionalFeeHook,
             // custom hooks take precedence over base types
             ...hookClasses,
         };
diff --git a/typescript/test/hooks/directionalFee.test.ts b/typescript/test/hooks/directionalFee.test.ts
new file mode 100644
index 0000000..306c704
--- /dev/null
+++ b/typescript/test/hooks/directionalFee.test.ts
@@ -0,0 +1,188 @@
+// pnpm test -- directionalFee.test.ts
+import { describe, expect, test } from 'vitest';
+import { SwapKind, Vault, SwapInput, PoolState } from '../../src';
+import {
+    DirectionalFeeHook,
+    HookStateDirectionalFee,
+} from '../../src/hooks/directionalFeeHook';
+
+const poolBalancesScaled18 = [
+    20000000000000000000000n,
+    20000000000000000000000n,
+];
+const swapAmountRaw = 100000000n;
+const staticSwapFeePercentage = 1000000000000000n; // 0.1%
+const poolTokens = [
+    {
+        address: '0xaA8E23Fb1079EA71e0a56F48a2aA51851D8433D0',
+        decimals: 6n,
+        index: 0n,
+    },
+    {
+        address: '0xFF34B3d4Aee8ddCd6F9AFFFB6Fe49bD371b8a357',
+        decimals: 18n,
+        index: 1n,
+    },
+];
+
+const scalingFactors = [1000000000000n, 1n];
+const totalSupply = 40000000000000000000000n;
+
+const hookState: HookStateDirectionalFee = {
+    tokens: poolTokens,
+    balancesLiveScaled18: poolBalancesScaled18,
+};
+
+const stablePoolStateWithHook: PoolState = {
+    poolType: 'STABLE',
+    hookType: 'DirectionalFee',
+    tokens: poolTokens.map((token) => token.address),
+    scalingFactors: scalingFactors,
+    tokenRates: [1000000000000000000n, 1000000000000000000n],
+    balancesLiveScaled18: poolBalancesScaled18,
+    swapFee: staticSwapFeePercentage,
+    aggregateSwapFee: 0n,
+    totalSupply: totalSupply,
+    amp: 1000000n,
+    hookState: hookState,
+};
+
+const stablePoolStateWithoutHook: PoolState = {
+    poolType: 'STABLE',
+    tokens: poolTokens.map((token) => token.address),
+    scalingFactors: scalingFactors,
+    tokenRates: [1000000000000000000n, 1000000000000000000n],
+    balancesLiveScaled18: poolBalancesScaled18,
+    swapFee: staticSwapFeePercentage,
+    aggregateSwapFee: 0n,
+    totalSupply: totalSupply,
+    amp: 1000000n,
+    hookState: hookState,
+};
+
+const swapParams: SwapInput = {
+    swapKind: SwapKind.GivenIn,
+    amountRaw: swapAmountRaw,
+    tokenIn: poolTokens[0].address,
+    tokenOut: poolTokens[1].address,
+};
+
+describe('hook - directionalFee', () => {
+    const vault = new Vault();
+    const hook = new DirectionalFeeHook();
+
+    test('computes directional fee', () => {
+        // change poolStateWithHook to have no swap fees.
+        const { success, dynamicSwapFee } = hook.onComputeDynamicSwapFee(
+            swapParams,
+            0n,
+            hookState,
+        );
+        expect(success).toBe(true);
+        expect(dynamicSwapFee).toBeGreaterThan(0n);
+    });
+    test('uses static swap fee when directional fee is not applicable', () => {
+        // set swap amount so that dynamic swap fee is lower and static one gets used instead
+        const newSwapParams = {
+            ...swapParams,
+            amountRaw: swapAmountRaw / 10000000000000n,
+        };
+        const { success, dynamicSwapFee } = hook.onComputeDynamicSwapFee(
+            newSwapParams,
+            staticSwapFeePercentage,
+            hookState,
+        );
+        expect(success).toBe(true);
+        expect(dynamicSwapFee).toEqual(staticSwapFeePercentage);
+    });
+    test('it uses dynamic swap fee with high enough swap amount - given out', () => {
+        // swap amount is big enough to trigger dynamic swap fee
+        const { success, dynamicSwapFee } = hook.onComputeDynamicSwapFee(
+            swapParams,
+            staticSwapFeePercentage,
+            hookState,
+        );
+        expect(success).toBe(true);
+        expect(dynamicSwapFee).toBeGreaterThan(staticSwapFeePercentage);
+        // based on this simulation https://dashboard.tenderly.co/mcquardt/project/simulator/3e0f9953-f1f7-4936-b083-cbd2958bd801?trace=0.1.0.2.2.0.3.18
+        expect(dynamicSwapFee).toEqual(5000000000000000n);
+    });
+
+    test('directional fee higher than static swap fee - given in', () => {
+        // based on this simulation https://dashboard.tenderly.co/mcquardt/project/simulator/3e0f9953-f1f7-4936-b083-cbd2958bd801?trace=0.1.0.2.2.0.3.18
+        const outputAmountWithHook = vault.swap(
+            swapParams,
+            stablePoolStateWithHook,
+            hookState,
+        );
+        expect(outputAmountWithHook).toEqual(99499505472260433154n);
+
+        // since no hook is part of the pool state, the vault should
+        // not compute the onDymamicSwapFee logic and should go with the
+        // static swap fee percentage.
+        const outputAmountWithoutHook = vault.swap(
+            swapParams,
+            stablePoolStateWithoutHook,
+            hookState,
+        );
+        // This must always hold, otherwise the hook is wrongly implemented
+        expect(outputAmountWithHook).toBeLessThan(outputAmountWithoutHook);
+    });
+    test('directional fee lower than static swap fee - given in ', () => {
+        // due to swap amount, the dynamic swap fee is lower than the static swap fee
+        // so the vault will use the static swap fee percentage to calculate the swap
+        // https://dashboard.tenderly.co/mcquardt/project/simulator/a4b17794-6eef-4f21-8eee-1297bc280ddd?trace=0.1.0.2.2.0.3.21
+
+        const swapParamsWithLowerAmountIn = {
+            ...swapParams,
+            amountRaw: 1000000n,
+        };
+
+        const { success, dynamicSwapFee } = hook.onComputeDynamicSwapFee(
+            swapParamsWithLowerAmountIn,
+            staticSwapFeePercentage,
+            hookState,
+        );
+        expect(success).toBe(true);
+        expect(dynamicSwapFee).toEqual(1000000000000000n);
+
+        const outputAmountWithHook = vault.swap(
+            swapParamsWithLowerAmountIn,
+            stablePoolStateWithHook,
+            hookState,
+        );
+        expect(outputAmountWithHook).toEqual(998999950149802562n);
+    });
+    test('directional fee lower than static swap fee - given out',  () => {
+        // with 1 DAI out, the dynamic swap fee hook does not get triggered
+        // therefore the direction fee is lower than the static swap fee
+        // and the static swap fee is being charged by the vault
+        // sim: https://dashboard.tenderly.co/mcquardt/project/simulator/06bc1601-1987-4a9a-99cb-e565bb57218d
+        const swapParamsGivenOut = {
+            ...swapParams,
+            amountRaw: 1000000000000000000n,
+            swapKind: SwapKind.GivenOut,
+        };
+        const amountIn = vault.swap(
+            swapParamsGivenOut,
+            stablePoolStateWithHook,
+            hookState,
+        );
+        expect(amountIn).toEqual(1001002n);
+    });
+    test('directional fee higher than static swap fee - given out', () => {
+        // with 100 DAI out, the dynamic swap fee hook does get triggered
+        // sim: https://dashboard.tenderly.co/mcquardt/project/simulator/571c3a06-d969-43fd-b4b8-08833b9c0997?trace=0.1.0.2.2.0.3.21.0
+        const swapParamsGivenOut = {
+            ...swapParams,
+            amountRaw: 100000000000000000000n,
+            swapKind: SwapKind.GivenOut,
+        };
+        const amountIn = vault.swap(
+            swapParamsGivenOut,
+            stablePoolStateWithHook,
+            hookState,
+        );
+        expect(amountIn).toEqual(100503015n);
+    })
+});
diff --git a/typescript/test/stableMath.test.ts b/typescript/test/stableMath.test.ts
new file mode 100644
index 0000000..0d3c2a6
--- /dev/null
+++ b/typescript/test/stableMath.test.ts
@@ -0,0 +1,27 @@
+// pnpm test ./test/stableMath.test.ts
+
+import { _computeBalance, _computeInvariant } from 'src/stable/stableMath';
+import { describe, expect, test } from 'vitest';
+
+describe('test stableMath', () => {
+    test('_computeBalance', () => {
+        // based on this sim
+        // https://dashboard.tenderly.co/mcquardt/project/simulator/f174cf82-3525-4376-b13d-9e61bad1649c?trace=0.4.0
+        const finalBalances = _computeBalance(
+            1000000n,
+            [20099500000000000000000n, 20000000000000000000000n],
+            40000000000000000000000n,
+            1,
+        );
+        expect(finalBalances).toEqual(19900500494527739566845n);
+    });
+    test('_computeInvariant', () => {
+        // based on this sim
+        // https://dashboard.tenderly.co/mcquardt/project/simulator/f174cf82-3525-4376-b13d-9e61bad1649c?trace=0.4.0
+        const invariant = _computeInvariant(1000000n, [
+            20000000000000000000000n,
+            20000000000000000000000n,
+        ]);
+        expect(invariant).toEqual(40000000000000000000000n);
+    });
+});
diff --git a/typescript/test/stablePool.test.ts b/typescript/test/stablePool.test.ts
index deba318..f7b41e2 100644
--- a/typescript/test/stablePool.test.ts
+++ b/typescript/test/stablePool.test.ts
@@ -1,5 +1,6 @@
+// pnpm test -- stablePool.test.ts
 import { describe, expect, test } from 'vitest';
-import { SwapKind } from '../src/index';
+import { SwapKind, SwapParams } from '../src/index';
 import { Stable } from '../src/stable';
 
 describe('stable pool', () => {
@@ -88,4 +89,26 @@ describe('stable pool', () => {
             });
         });
     });
+    describe('onSwap', () => {
+        test('matches onchain results', () => {
+            // sim https://dashboard.tenderly.co/mcquardt/project/simulator/f174cf82-3525-4376-b13d-9e61bad1649c?trace=0
+            const tempPool = new Stable({
+                amp: 1000000n,
+            });
+
+            const swapParams: SwapParams = {
+                swapKind: SwapKind.GivenIn,
+                amountGivenScaled18: 99500000000000000000n,
+                balancesLiveScaled18: [
+                    20000000000000000000000n,
+                    20000000000000000000000n,
+                ],
+                indexIn: 0,
+                indexOut: 1,
+            };
+
+            const amountOut = tempPool.onSwap(swapParams);
+            expect(amountOut).toEqual(99499505472260433154n);
+        });
+    });
 });