tests/test_pnl.py

"""
Tests for P&L calculation functions.
"""

import datetime
import unittest
from unittest.mock import patch

import numpy as np

from src.folio.data_model import OptionPosition, StockPosition
from src.folio.pnl import (
    calculate_breakeven_points,
    calculate_max_profit_loss,
    calculate_position_pnl,
    calculate_strategy_pnl,
    determine_price_range,
    summarize_strategy_pnl,
)


class TestPnLCalculations(unittest.TestCase):
    """Test cases for P&L calculation functions."""

    def setUp(self):
        """Set up test fixtures."""
        # Create a sample stock position
        self.stock_position = StockPosition(
            ticker="SPY",
            quantity=100,
            beta=1.0,
            beta_adjusted_exposure=45000.0,
            market_exposure=45000.0,  # 100 shares * $450
            price=450.0,  # $450 per share
            cost_basis=400.0,  # $400 per share cost basis
        )

        # Create a sample call option position
        expiry_date = datetime.datetime.now() + datetime.timedelta(days=30)
        expiry_str = expiry_date.strftime("%Y-%m-%d")
        self.call_option = OptionPosition(
            ticker="SPY",
            position_type="option",
            quantity=1,
            beta=1.0,
            beta_adjusted_exposure=1000.0,
            market_exposure=1000.0,  # 1 contract * $10 * 100 shares
            strike=460.0,
            expiry=expiry_str,
            option_type="CALL",
            delta=0.5,
            delta_exposure=2250.0,  # 0.5 * 100 * $450 * 1
            notional_value=45000.0,  # 100 * $450 * 1
            underlying_beta=1.0,
            price=10.0,  # $10 per contract
            cost_basis=8.0,  # $8 per contract cost basis
        )

        # Create a sample put option position
        self.put_option = OptionPosition(
            ticker="SPY",
            position_type="option",
            quantity=-2,  # Short 2 contracts
            beta=1.0,
            beta_adjusted_exposure=-1000.0,
            market_exposure=-1000.0,  # -2 contracts * $5 * 100 shares
            strike=440.0,
            expiry=expiry_str,
            option_type="PUT",
            delta=-0.4,
            delta_exposure=3600.0,  # -0.4 * 100 * $450 * -2
            notional_value=90000.0,  # 100 * $450 * 2
            underlying_beta=1.0,
            price=5.0,  # $5 per contract
            cost_basis=6.0,  # $6 per contract cost basis
        )

    def test_determine_price_range(self):
        """Test price range determination."""
        # Test with stock only
        price_range = determine_price_range([self.stock_position], 450.0)
        self.assertEqual(len(price_range), 2)
        self.assertLess(price_range[0], 450.0)
        self.assertGreater(price_range[1], 450.0)

        # Test with options
        price_range = determine_price_range(
            [self.stock_position, self.call_option, self.put_option], 450.0
        )
        self.assertEqual(len(price_range), 2)
        # Should include strike prices with margin
        self.assertLessEqual(price_range[0], 440.0 * 0.8)
        self.assertGreaterEqual(price_range[1], 460.0 * 1.2)

    @patch("src.folio.pnl.calculate_bs_price")
    def test_calculate_position_pnl_stock(self, mock_calculate_bs_price):
        """Test P&L calculation for a stock position."""
        # Calculate P&L for stock position using current price as entry price (default)
        pnl_data = calculate_position_pnl(
            self.stock_position,
            price_range=(400.0, 500.0),
            num_points=11,  # 400, 410, 420, ..., 500
            use_cost_basis=False,  # Use current price as entry price
        )

        # Verify the structure of the result
        self.assertIn("price_points", pnl_data)
        self.assertIn("pnl_values", pnl_data)
        self.assertEqual(len(pnl_data["price_points"]), 11)
        self.assertEqual(len(pnl_data["pnl_values"]), 11)

        # Verify P&L calculations for stock
        # P&L = (price - entry_price) * quantity
        # Entry price is $450 per share (current price)
        expected_pnls = [
            (price - 450.0) * 100 for price in np.linspace(400.0, 500.0, 11)
        ]
        for i, expected_pnl in enumerate(expected_pnls):
            self.assertAlmostEqual(pnl_data["pnl_values"][i], expected_pnl, places=2)

        # Verify mock wasn't called for stock position
        mock_calculate_bs_price.assert_not_called()

        # Reset mock for the next test
        mock_calculate_bs_price.reset_mock()

        # Calculate P&L for stock position using cost basis as entry price
        pnl_data_cost_basis = calculate_position_pnl(
            self.stock_position,
            price_range=(400.0, 500.0),
            num_points=11,  # 400, 410, 420, ..., 500
            use_cost_basis=True,  # Use cost basis as entry price
        )

        # Verify P&L calculations for stock using cost basis
        # P&L = (price - cost_basis) * quantity
        # Cost basis is $400 per share
        expected_pnls_cost_basis = [
            (price - 400.0) * 100 for price in np.linspace(400.0, 500.0, 11)
        ]
        for i, expected_pnl in enumerate(expected_pnls_cost_basis):
            self.assertAlmostEqual(
                pnl_data_cost_basis["pnl_values"][i], expected_pnl, places=2
            )

        # Verify mock wasn't called for stock position
        mock_calculate_bs_price.assert_not_called()

    @patch("src.folio.pnl.calculate_bs_price")
    def test_calculate_position_pnl_option(self, mock_calculate_bs_price):
        """Test P&L calculation for an option position."""
        # Mock the option pricing function for default mode
        mock_calculate_bs_price.side_effect = [5.0, 10.0, 15.0, 20.0, 25.0]

        # Calculate P&L for call option position using current price as entry price (default)
        pnl_data = calculate_position_pnl(
            self.call_option,
            price_range=(440.0, 480.0),
            num_points=5,  # 440, 450, 460, 470, 480
            use_cost_basis=False,  # Use current price as entry price
        )

        # Verify the structure of the result
        self.assertIn("price_points", pnl_data)
        self.assertIn("pnl_values", pnl_data)
        self.assertEqual(len(pnl_data["price_points"]), 5)
        self.assertEqual(len(pnl_data["pnl_values"]), 5)

        # Verify P&L calculations for option
        # P&L = (theo_price - entry_price) * quantity * contract_multiplier
        # Entry price is $10 per contract (current price), quantity is 1
        # Contract multiplier is 100 (each contract controls 100 shares)
        contract_multiplier = 100
        expected_pnls = [
            (price - 10.0) * 1 * contract_multiplier
            for price in [5.0, 10.0, 15.0, 20.0, 25.0]
        ]
        for i, expected_pnl in enumerate(expected_pnls):
            self.assertAlmostEqual(pnl_data["pnl_values"][i], expected_pnl, places=2)

        # Verify mock was called for option position
        self.assertEqual(mock_calculate_bs_price.call_count, 5)

        # Reset mock and set new side effect for cost basis mode
        mock_calculate_bs_price.reset_mock()
        mock_calculate_bs_price.side_effect = [5.0, 10.0, 15.0, 20.0, 25.0]

        # Calculate P&L for call option position using cost basis as entry price
        pnl_data_cost_basis = calculate_position_pnl(
            self.call_option,
            price_range=(440.0, 480.0),
            num_points=5,  # 440, 450, 460, 470, 480
            use_cost_basis=True,  # Use cost basis as entry price
        )

        # Verify P&L calculations for option using cost basis
        # P&L = (theo_price - cost_basis) * quantity * contract_multiplier
        # Cost basis is $8 per contract, quantity is 1
        expected_pnls_cost_basis = [
            (price - 8.0) * 1 * contract_multiplier
            for price in [5.0, 10.0, 15.0, 20.0, 25.0]
        ]
        for i, expected_pnl in enumerate(expected_pnls_cost_basis):
            self.assertAlmostEqual(
                pnl_data_cost_basis["pnl_values"][i], expected_pnl, places=2
            )

        # Verify mock was called for option position
        self.assertEqual(mock_calculate_bs_price.call_count, 5)

    @patch("src.folio.pnl.calculate_position_pnl")
    def test_calculate_strategy_pnl(self, mock_calculate_position_pnl):
        """Test P&L calculation for a strategy (multiple positions)."""
        # Mock the position P&L calculations for default mode
        mock_calculate_position_pnl.side_effect = [
            {
                "price_points": [400.0, 450.0, 500.0],
                "pnl_values": [-4000.0, 1000.0, 6000.0],
                "position": {},
            },
            {
                "price_points": [400.0, 450.0, 500.0],
                "pnl_values": [500.0, 200.0, -100.0],
                "position": {},
            },
            {
                "price_points": [400.0, 450.0, 500.0],
                "pnl_values": [1000.0, 0.0, -1000.0],
                "position": {},
            },
        ]

        # Calculate P&L for a strategy with all positions using current price as entry price (default)
        positions = [self.stock_position, self.call_option, self.put_option]
        pnl_data = calculate_strategy_pnl(
            positions, price_range=(400.0, 500.0), num_points=3, use_cost_basis=False
        )

        # Verify the structure of the result
        self.assertIn("price_points", pnl_data)
        self.assertIn("pnl_values", pnl_data)
        self.assertIn("individual_pnls", pnl_data)
        self.assertEqual(len(pnl_data["price_points"]), 3)
        self.assertEqual(len(pnl_data["pnl_values"]), 3)
        self.assertEqual(len(pnl_data["individual_pnls"]), 3)

        # Verify combined P&L calculations
        # Combined P&L = sum of individual P&Ls
        expected_combined_pnls = [-2500.0, 1200.0, 4900.0]
        for i, expected_pnl in enumerate(expected_combined_pnls):
            self.assertAlmostEqual(pnl_data["pnl_values"][i], expected_pnl, places=2)

        # Verify mock was called for each position
        self.assertEqual(mock_calculate_position_pnl.call_count, 3)

        # Reset mock for cost basis mode
        mock_calculate_position_pnl.reset_mock()

        # Mock the position P&L calculations for cost basis mode
        mock_calculate_position_pnl.side_effect = [
            {
                "price_points": [400.0, 450.0, 500.0],
                "pnl_values": [
                    -3000.0,
                    2000.0,
                    7000.0,
                ],  # Different values for cost basis
                "position": {},
            },
            {
                "price_points": [400.0, 450.0, 500.0],
                "pnl_values": [700.0, 400.0, 100.0],  # Different values for cost basis
                "position": {},
            },
            {
                "price_points": [400.0, 450.0, 500.0],
                "pnl_values": [
                    800.0,
                    -200.0,
                    -1200.0,
                ],  # Different values for cost basis
                "position": {},
            },
        ]

        # Calculate P&L for a strategy with all positions using cost basis as entry price
        pnl_data_cost_basis = calculate_strategy_pnl(
            positions, price_range=(400.0, 500.0), num_points=3, use_cost_basis=True
        )

        # Verify combined P&L calculations for cost basis mode
        expected_combined_pnls_cost_basis = [
            -1500.0,
            2200.0,
            5900.0,
        ]  # Different values for cost basis
        for i, expected_pnl in enumerate(expected_combined_pnls_cost_basis):
            self.assertAlmostEqual(
                pnl_data_cost_basis["pnl_values"][i], expected_pnl, places=2
            )

        # Verify mock was called for each position
        self.assertEqual(mock_calculate_position_pnl.call_count, 3)

    def test_calculate_breakeven_points(self):
        """Test calculation of breakeven points."""
        # Create sample P&L data with a zero crossing
        pnl_data = {
            "price_points": [400.0, 425.0, 450.0, 475.0, 500.0],
            "pnl_values": [-1000.0, -500.0, 0.0, 500.0, 1000.0],
        }

        # Calculate breakeven points
        breakeven_points = calculate_breakeven_points(pnl_data)

        # Verify the result - should find 2 breakeven points due to numerical precision
        self.assertEqual(len(breakeven_points), 2)
        # Both should be close to 450.0
        for bp in breakeven_points:
            self.assertAlmostEqual(bp, 450.0, places=1)

        # Test with multiple zero crossings
        pnl_data = {
            "price_points": [400.0, 425.0, 450.0, 475.0, 500.0],
            "pnl_values": [500.0, -500.0, 0.0, -500.0, 500.0],
        }

        breakeven_points = calculate_breakeven_points(pnl_data)

        # Should find 4 breakeven points due to numerical precision
        self.assertEqual(len(breakeven_points), 4)

    def test_calculate_max_profit_loss(self):
        """Test calculation of maximum profit and loss."""
        # Create sample P&L data
        pnl_data = {
            "price_points": [400.0, 425.0, 450.0, 475.0, 500.0],
            "pnl_values": [-1000.0, -500.0, 0.0, 1500.0, 1000.0],
        }

        # Calculate max profit/loss
        max_pl = calculate_max_profit_loss(pnl_data)

        # Verify the result
        self.assertEqual(max_pl["max_profit"], 1500.0)
        self.assertEqual(max_pl["max_profit_price"], 475.0)
        self.assertEqual(max_pl["max_loss"], -1000.0)
        self.assertEqual(max_pl["max_loss_price"], 400.0)

    def test_summarize_strategy_pnl(self):
        """Test strategy P&L summary generation."""
        # Create sample P&L data
        pnl_data = {
            "price_points": [400.0, 425.0, 450.0, 475.0, 500.0],
            "pnl_values": [-1000.0, -500.0, 0.0, 1500.0, 1000.0],
        }

        # Generate summary
        summary = summarize_strategy_pnl(pnl_data, 450.0)

        # Verify the structure of the result
        self.assertIn("breakeven_points", summary)
        self.assertIn("max_profit", summary)
        self.assertIn("max_loss", summary)
        self.assertIn("current_pnl", summary)
        self.assertIn("profitable_ranges", summary)

        # Verify specific values
        self.assertAlmostEqual(summary["max_profit"], 1500.0, places=2)
        self.assertAlmostEqual(summary["max_loss"], -1000.0, places=2)
        self.assertAlmostEqual(summary["current_pnl"], 0.0, places=2)

        # Should have two breakeven points due to numerical precision
        self.assertEqual(len(summary["breakeven_points"]), 2)
        # Both should be close to 450.0
        for bp in summary["breakeven_points"]:
            self.assertAlmostEqual(bp, 450.0, places=1)

        # Should have one profitable range
        self.assertEqual(len(summary["profitable_ranges"]), 1)
        start, end = summary["profitable_ranges"][0]
        # The profitable range starts at 475.0 in our implementation
        self.assertAlmostEqual(start, 475.0, places=2)
        self.assertAlmostEqual(end, 500.0, places=2)


if __name__ == "__main__":
    unittest.main()