from django.test import TestCase

from core.exchanges import convert
from core.models import RiskModel, User
from core.trading import risk


class RiskModelTestCase(TestCase):
    def setUp(self):
        self.user = User.objects.create_user(
            username="testuser", email="test@example.com", password="test"
        )

        self.risk_model = RiskModel.objects.create(
            user=self.user,
            name="Test Risk Model",
            max_loss_percent=50,
            max_risk_percent=10,
            max_open_trades=3,
            max_open_trades_per_symbol=2,
        )
        self.account_initial_balance = 100
        self.trade = {
            "symbol": "XXXYYY",
            "side": "BUY",
            # We already calculated the TP percent loss relative to the account size
            "take_profit_percent": 9,
            "take_profit_usd": 9,
            "stop_loss_percent": 9,
            "stop_loss_usd": 9,
            "trailing_stop_loss_percent": 9,
            "trailing_stop_loss_usd": 9,
        }

    def test_check_max_loss(self):
        """
        Check that we can open a trade within the max loss limit.
        """
        account_balance = 100  # We have lost no money
        allowed = risk.check_max_loss(
            self.risk_model, self.account_initial_balance, account_balance
        )
        self.assertTrue(allowed)

    def test_check_max_loss_fail_exact(self):
        """
        Check that we cannot open a trade outside the max loss limit where the amount
        lost is exactly the max loss limit.
        """
        account_balance = 50  # We have lost 50% exactly
        allowed = risk.check_max_loss(
            self.risk_model, self.account_initial_balance, account_balance
        )
        self.assertFalse(allowed)

    def test_check_max_loss_fail(self):
        """
        Check that we cannot open a trade outside the max loss limit.
        """
        account_balance = 49  # We have lost 51%
        allowed = risk.check_max_loss(
            self.risk_model, self.account_initial_balance, account_balance
        )
        self.assertFalse(allowed)

    def test_check_max_risk(self):
        """
        Check that we can open a trade within the max risk limit.
        """
        account_trades = [self.trade]
        allowed = risk.check_max_risk(
            self.risk_model, self.account_initial_balance, account_trades
        )
        self.assertTrue(allowed)

    def test_check_max_risk_multiple_trades(self):
        """
        Check that we can open a trade within the max risk limit where there are
        multiple trades.
        """
        trade = self.trade.copy()
        trade["stop_loss_percent"] = 1
        trade["trailing_stop_loss_percent"] = 1
        trade["stop_loss_usd"] = 1
        trade["trailing_stop_loss_usd"] = 1
        account_trades = [trade] * 9
        allowed = risk.check_max_risk(
            self.risk_model, self.account_initial_balance, account_trades
        )
        self.assertTrue(allowed)

    def test_check_max_risk_fail_exact(self):
        """
        Check that we cannot open a trade outside the max risk limit where the amount
        risked is exactly the max risk limit.
        """
        trade = self.trade.copy()
        trade["stop_loss_percent"] = 10
        trade["trailing_stop_loss_percent"] = 10
        trade["stop_loss_usd"] = 10
        trade["trailing_stop_loss_usd"] = 10
        account_trades = [trade]
        allowed = risk.check_max_risk(
            self.risk_model, self.account_initial_balance, account_trades
        )
        self.assertFalse(allowed)

    def test_check_max_risk_fail_exact_multiple_trades(self):
        """
        Check that we cannot open a trade outside the max risk limit where the amount
        risked is exactly the max risk limit with multiple trades.
        """
        trade = self.trade.copy()
        trade["stop_loss_percent"] = 1
        trade["trailing_stop_loss_percent"] = 1
        trade["stop_loss_usd"] = 1
        trade["trailing_stop_loss_usd"] = 1
        account_trades = [trade] * 10
        allowed = risk.check_max_risk(
            self.risk_model, self.account_initial_balance, account_trades
        )
        self.assertFalse(allowed)

    def test_check_max_open_trades(self):
        """
        Check that we can open a trade within the max open trades limit.
        """
        account_trades = [self.trade] * 2
        allowed = risk.check_max_open_trades(self.risk_model, account_trades)
        self.assertTrue(allowed)

    def test_check_max_open_trades_fail_exact(self):
        """
        Check that we cannot open a trade at the max open trades limit.
        """
        account_trades = [self.trade] * 3
        allowed = risk.check_max_open_trades(self.risk_model, account_trades)
        self.assertFalse(allowed)

    def test_check_max_open_trades_fail(self):
        """
        Check that we cannot open a trade outside the max open trades limit.
        """
        account_trades = [self.trade] * 4
        allowed = risk.check_max_open_trades(self.risk_model, account_trades)
        self.assertFalse(allowed)

    def test_check_max_open_trades_per_symbol(self):
        """
        Check that we can open a trade within the max open trades per symbol limit.
        """
        account_trades = [self.trade]
        allowed = risk.check_max_open_trades_per_symbol(self.risk_model, account_trades)
        self.assertTrue(allowed)

    def test_check_max_open_trades_per_symbol_fail_exact(self):
        """
        Check that we cannot open a trade at the max open trades per symbol limit.
        """
        account_trades = [self.trade] * 2
        allowed = risk.check_max_open_trades_per_symbol(self.risk_model, account_trades)
        self.assertFalse(allowed)

    def test_check_max_open_trades_per_symbol_fail(self):
        """
        Check that we cannot open a trade outside the max open trades per symbol limit.
        """
        account_trades = [self.trade] * 3
        allowed = risk.check_max_open_trades_per_symbol(self.risk_model, account_trades)
        self.assertFalse(allowed)

    def test_check_max_open_trades_per_symbol_different_symbols(self):
        """
        Check that we can open a trade within the max open trades per symbol limit with
        different symbols.
        """
        trade1 = self.trade.copy()
        trade2 = self.trade.copy()
        trade1["symbol"] = "ONE"
        trade2["symbol"] = "TWO"
        account_trades = [trade1, trade2]
        allowed = risk.check_max_open_trades_per_symbol(self.risk_model, account_trades)
        self.assertTrue(allowed)

    def test_check_max_open_trades_per_symbol_fail_exact_different_symbols(self):
        """
        Check that we cannot open a trade at the max open trades per symbol limit with
        different symbols.
        """
        trade1 = self.trade.copy()
        trade2 = self.trade.copy()
        trade1["symbol"] = "ONE"
        trade2["symbol"] = "TWO"
        account_trades = [trade1, trade2, trade1, trade2]
        allowed = risk.check_max_open_trades_per_symbol(self.risk_model, account_trades)
        self.assertFalse(allowed)

    def test_check_max_open_trades_per_symbol_fail_different_symbols(self):
        """
        Check that we cannot open a trade outside the max open trades per symbol limit
        with different symbols.
        """
        trade1 = self.trade.copy()
        trade2 = self.trade.copy()
        trade1["symbol"] = "ONE"
        trade2["symbol"] = "TWO"
        # Each one 3 times
        account_trades = [trade1, trade2, trade1, trade2, trade1, trade2]
        allowed = risk.check_max_open_trades_per_symbol(self.risk_model, account_trades)
        self.assertFalse(allowed)

    # Market data tests, account size: $1000
    def test_check_max_risk_market_data(self):
        """
        Check that we can open a trade within the max risk limit with market data.
        """
        trade = {
            "id": "abd123",
            "symbol": "EUR_USD",
            "currentUnits": 100,
            "side": "long",
            "state": "open",
            "price": 1.0,  # initial
            "unrealizedPL": 0,  # price == initial
            "stopLossOrder": {
                "price": 0.95,  # down by 5%, 5% risk
            },
            # Hardcoded prices to avoid calling market API here
            "stop_loss_usd": 50,  # 5% of $1000
        }
        converted = convert.convert_trades([trade])
        self.assertEqual(converted[0]["stop_loss_percent"], 5)
        self.assertEqual(converted[0]["stop_loss_usd"], 50)
        max_risk_check = risk.check_max_risk(self.risk_model, 1000, converted)
        self.assertTrue(max_risk_check)  # 5% risk is fine

    def test_check_max_risk_market_data_multiple(self):
        """
        Check that we can open a trade within the max risk limit with market data
        and multiple trades.
        """
        trade = {
            "id": "abd123",
            "symbol": "EUR_USD",
            "currentUnits": 100,
            "side": "long",
            "state": "open",
            "price": 1.0,  # initial
            "unrealizedPL": 0,  # price == initial
            "stopLossOrder": {
                "price": 0.96,  # down by 4%, 4% risk
            },
            # Hardcoded prices to avoid calling market API here
            "stop_loss_usd": 40,  # 4% of $1000
        }
        converted = convert.convert_trades([trade, trade])
        self.assertEqual(converted[0]["stop_loss_percent"], 4)
        self.assertEqual(converted[0]["stop_loss_usd"], 40)
        max_risk_check = risk.check_max_risk(self.risk_model, 1000, converted)
        self.assertTrue(max_risk_check)  # 8% risk is fine

    def test_check_max_risk_market_data_fail(self):
        """
        Check that we can not open a trade outside the max risk limit with market data.
        """
        trade = {
            "id": "abd123",
            "symbol": "EUR_USD",
            "currentUnits": 100,
            "side": "long",
            "state": "open",
            "price": 1.0,  # initial
            "unrealizedPL": 0,  # price == initial
            "stopLossOrder": {
                "price": 0.9,  # down by 10%, 10% risk
            },
            # Hardcoded prices to avoid calling market API here
            "stop_loss_usd": 100,  # 10% of $1000
        }
        converted = convert.convert_trades([trade])
        self.assertEqual(converted[0]["stop_loss_percent"], 10)
        self.assertEqual(converted[0]["stop_loss_usd"], 100)
        max_risk_check = risk.check_max_risk(self.risk_model, 1000, converted)
        self.assertFalse(max_risk_check)  # 10% risk is too much

    def test_check_max_risk_market_data_fail_multiple(self):
        """
        Check that we can not open a trade outside the max risk limit with market data
        and multiple trades.
        """
        trade = {
            "id": "abd123",
            "symbol": "EUR_USD",
            "currentUnits": 100,
            "side": "long",
            "state": "open",
            "price": 1.0,  # initial
            "unrealizedPL": 0,  # price == initial
            "stopLossOrder": {
                "price": 0.95,  # down by 5%, 5% risk
            },
            # Hardcoded prices to avoid calling market API here
            "stop_loss_usd": 50,  # 5% of $1000
        }
        converted = convert.convert_trades([trade, trade])
        self.assertEqual(converted[0]["stop_loss_percent"], 5)
        self.assertEqual(converted[0]["stop_loss_usd"], 50)
        self.assertEqual(converted[1]["stop_loss_percent"], 5)
        self.assertEqual(converted[1]["stop_loss_usd"], 50)
        max_risk_check = risk.check_max_risk(self.risk_model, 1000, converted)
        self.assertFalse(max_risk_check)  # 10% risk is too much

    def test_check_max_risk_market_data_fail_multiple_mixed(self):
        """
        Check that we can not open a trade outside the max risk limit with market data
        and multiple trades, mixing SL and TSL.
        """
        trade = {
            "id": "abd123",
            "symbol": "EUR_USD",
            "currentUnits": 100,
            "side": "long",
            "state": "open",
            "price": 1.0,  # initial
            "unrealizedPL": 0,  # price == initial
            "stopLossOrder": {
                "price": 0.95,  # down by 5%, 5% risk
            },
            # Hardcoded prices to avoid calling market API here
            "stop_loss_usd": 50,  # 5% of $1000
            "trailing_stop_loss_usd": 50,
        }
        trade2 = trade.copy()
        trade2["trailingStopLossOrder"] = {"price": 0.95}
        del trade2["stopLossOrder"]

        converted = convert.convert_trades([trade, trade2])
        self.assertEqual(converted[0]["stop_loss_percent"], 5)
        self.assertEqual(converted[0]["stop_loss_usd"], 50)
        self.assertEqual(converted[1]["trailing_stop_loss_percent"], 5)
        self.assertEqual(converted[1]["trailing_stop_loss_usd"], 50)
        max_risk_check = risk.check_max_risk(self.risk_model, 1000, converted)
        self.assertFalse(max_risk_check)  # 10% risk is too much

    def test_check_max_risk_market_data_fail_multiple_mixed_both(self):
        """
        Check that we can not open a trade outside the max risk limit with market data
        and multiple trades, mixing SL and TSL, where both are set.
        """
        trade = {
            "id": "abd123",
            "symbol": "EUR_USD",
            "currentUnits": 100,
            "side": "long",
            "state": "open",
            "price": 1.0,  # initial
            "unrealizedPL": 0,  # price == initial
            "stopLossOrder": {
                "price": 0.95,  # down by 5%, 5% risk
            },
            # Hardcoded prices to avoid calling market API here
            "stop_loss_usd": 50,  # 5% of $1000
            "trailing_stop_loss_usd": 49,
        }
        trade2 = trade.copy()
        trade2["trailingStopLossOrder"] = {"price": 0.951}

        converted = convert.convert_trades([trade, trade2])
        self.assertEqual(converted[0]["stop_loss_percent"], 5)
        self.assertEqual(converted[0]["stop_loss_usd"], 50)
        self.assertEqual(float(converted[1]["trailing_stop_loss_percent"]), 4.9)
        self.assertEqual(converted[1]["trailing_stop_loss_usd"], 49)
        max_risk_check = risk.check_max_risk(self.risk_model, 1000, converted)
        self.assertFalse(max_risk_check)  # 10% risk is too much