"""Tests for N-dimensional ellipsoidal bonding surface."""

import numpy as np
import pytest
from src.primitives.n_dimensional_surface import (
    NDSurfaceParams,
    NDSurfaceState,
    create_params,
    compute_invariant,
    calc_out_given_in,
    calc_in_given_out,
    mint,
    redeem,
    spot_prices,
)
from src.utils.linear_algebra import random_rotation_matrix


class TestInvariant:
    def test_2d_positive(self):
        """2D surface should have positive invariant."""
        params = create_params(2)
        balances = np.array([1000.0, 1000.0])
        r = compute_invariant(balances, params)
        assert r > 0

    def test_3d_positive(self):
        """3D surface should have positive invariant."""
        params = create_params(3)
        balances = np.array([1000.0, 1000.0, 1000.0])
        r = compute_invariant(balances, params)
        assert r > 0

    def test_5d_positive(self):
        """5D surface should work."""
        params = create_params(5)
        balances = np.array([100.0, 200.0, 300.0, 400.0, 500.0])
        r = compute_invariant(balances, params)
        assert r > 0

    def test_invariant_increases_with_reserves(self):
        """More reserves → larger invariant."""
        params = create_params(3)
        r1 = compute_invariant(np.array([100.0, 100.0, 100.0]), params)
        r2 = compute_invariant(np.array([200.0, 200.0, 200.0]), params)
        assert r2 > r1

    def test_stretched_surface(self):
        """Surface with stretch factors should compute invariant."""
        params = create_params(3, lambdas=np.array([1.0, 5.0, 10.0]))
        balances = np.array([1000.0, 1000.0, 1000.0])
        r = compute_invariant(balances, params)
        assert r > 0


class TestSwaps:
    def test_swap_preserves_invariant(self):
        """Swap should preserve invariant."""
        params = create_params(3)
        balances = np.array([1000.0, 1000.0, 1000.0])
        r_before = compute_invariant(balances, params)

        dy = calc_out_given_in(balances, params, 0, 1, 50.0)
        new_balances = balances.copy()
        new_balances[0] += 50.0
        new_balances[1] -= dy
        r_after = compute_invariant(new_balances, params)

        assert abs(r_after - r_before) / r_before < 1e-4

    def test_round_trip(self):
        """calc_in_given_out inverts calc_out_given_in."""
        params = create_params(3)
        balances = np.array([1000.0, 1000.0, 1000.0])

        amount_in = 30.0
        amount_out = calc_out_given_in(balances, params, 0, 2, amount_in)
        recovered = calc_in_given_out(balances, params, 0, 2, amount_out)
        assert abs(recovered - amount_in) < 0.1  # Slightly looser tolerance for N-D

    def test_swap_positive_output(self):
        """Swap should produce positive output."""
        params = create_params(4)
        balances = np.array([1000.0, 1000.0, 1000.0, 1000.0])
        dy = calc_out_given_in(balances, params, 0, 3, 100.0)
        assert dy > 0

    def test_all_pairs(self):
        """Should be able to swap between any pair."""
        params = create_params(3)
        balances = np.array([1000.0, 1000.0, 1000.0])
        for i in range(3):
            for j in range(3):
                if i != j:
                    dy = calc_out_given_in(balances, params, i, j, 10.0)
                    assert dy > 0


class TestMintRedeem:
    def test_mint_increases_supply(self):
        """Minting should increase supply."""
        params = create_params(3)
        state = NDSurfaceState(
            balances=np.array([1000.0, 1000.0, 1000.0]),
            invariant=compute_invariant(np.array([1000.0, 1000.0, 1000.0]), params),
            supply=1000.0,
        )
        amounts_in = np.array([100.0, 100.0, 100.0])
        new_state, minted = mint(state, params, amounts_in)

        assert minted > 0
        assert new_state.supply > state.supply
        assert np.all(new_state.balances > state.balances)

    def test_redeem_proportional(self):
        """Proportional redeem should return proportional assets."""
        params = create_params(3)
        balances = np.array([1000.0, 2000.0, 3000.0])
        state = NDSurfaceState(
            balances=balances,
            invariant=compute_invariant(balances, params),
            supply=1000.0,
        )

        new_state, amounts_out = redeem(state, params, 100.0)

        # Should get 10% of each asset
        assert abs(amounts_out[0] - 100.0) < 1e-6
        assert abs(amounts_out[1] - 200.0) < 1e-6
        assert abs(amounts_out[2] - 300.0) < 1e-6
        assert abs(new_state.supply - 900.0) < 1e-6

    def test_initial_mint(self):
        """First mint with zero supply should bootstrap."""
        params = create_params(2)
        state = NDSurfaceState(
            balances=np.array([0.0, 0.0]),
            invariant=0.0,
            supply=0.0,
        )
        # Bootstrap with initial reserves
        amounts_in = np.array([1000.0, 1000.0])
        # Need to handle zero invariant case
        state.balances = amounts_in
        state.invariant = compute_invariant(amounts_in, params)
        state.supply = state.invariant

        assert state.supply > 0

    def test_mint_then_redeem_roundtrip(self):
        """Mint then full redeem should return ~original reserves."""
        params = create_params(2)
        initial_balances = np.array([1000.0, 1000.0])
        state = NDSurfaceState(
            balances=initial_balances.copy(),
            invariant=compute_invariant(initial_balances, params),
            supply=1000.0,
        )

        # Mint 10%
        new_state, minted = mint(state, params, np.array([100.0, 100.0]))
        # Redeem all minted
        final_state, out = redeem(new_state, params, minted)

        # Should be close to original + deposit - withdrawal ≈ original
        np.testing.assert_allclose(final_state.balances, initial_balances, rtol=1e-4)


class TestSpotPrices:
    def test_symmetric_prices(self):
        """Balanced pool with identity Q should have ~equal prices."""
        params = create_params(3)
        balances = np.array([1000.0, 1000.0, 1000.0])
        prices = spot_prices(balances, params, numeraire=0)
        assert abs(prices[0] - 1.0) < 1e-10
        assert abs(prices[1] - 1.0) < 0.1
        assert abs(prices[2] - 1.0) < 0.1