#!/usr/bin/env python3
"""
Test diversification integration with TradeManager.

Usage:
    python test_diversification_integration.py
"""

import sys
import logging

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)

sys.path.insert(0, '/home/ai/projects/AI_Strategy')

print("="*80)
print("🧪 Testing Diversification Integration")
print("="*80)
print()

try:
    from trade.manager import TradeManager
    from trade.diversification_manager import DiversificationManager

    print("✅ Imports successful")
    print()

    # Test 1: DiversificationManager initialization
    print("Test 1: DiversificationManager initialization")
    div_mgr = DiversificationManager()
    print(f"  ✅ Initialized with config:")
    print(f"     - Max positions: {div_mgr.config['max_open_positions']}")
    print(f"     - Per ticker: {div_mgr.config['max_trades_per_ticker']}")
    print(f"     - Per direction: {div_mgr.config['max_trades_per_direction']}")
    print()

    # Test 2: Simulate trades
    print("Test 2: Simulating trades")
    signals = [
        {'ticker': 'SBER', 'direction': 'LONG', 'atr_entry': 0.003, 'entry_price': 310.0},
        {'ticker': 'GAZP', 'direction': 'LONG', 'atr_entry': 0.002, 'entry_price': 72.0},
        {'ticker': 'X5', 'direction': 'LONG', 'atr_entry': 0.0025, 'entry_price': 250.0},
        {'ticker': 'MOEX', 'direction': 'LONG', 'atr_entry': 0.0015, 'entry_price': 210.0},
    ]

    for i, sig in enumerate(signals):
        result = div_mgr.can_open_position(
            signal=sig,
            current_positions=[],
            portfolio_value=100000
        )

        can_open = result['can_open']
        reason = result.get('reason', 'N/A')
        score = result.get('score', 0)

        status = "✅" if can_open else "❌"
        score_display = f"{score:.2f}" if isinstance(score, (int, float)) else str(score)
        print(f"  {status} Trade {i+1}: {sig['ticker']} ({sig['direction']})")
        print(f"      Score: {score_display}, Reason: {reason}")

        if can_open:
            # Open with custom ID based on position
            position_id = i
            div_mgr.open_position(position_id, sig, int(__import__('datetime').datetime.now().timestamp()))
            print(f"      Position {position_id} opened in manager")

    print()

    # Test 3: Check positions
    print("Test 3: Checking open positions")
    open_positions = div_mgr.get_open_positions()
    print(f"  Total open positions: {len(open_positions)}")

    for pos in open_positions:
        print(f"    - {pos['ticker']} ({pos['direction']})")

    print()

    # Test 4: Test diversification blocking
    print("Test 4: Testing diversification blocking")
    additional_signals = [
        {'ticker': 'SBER', 'direction': 'LONG', 'atr_entry': 0.002, 'entry_price': 315.0},
        {'ticker': 'GAZP', 'direction': 'SHORT', 'atr_entry': 0.003, 'entry_price': 75.0},
    ]

    for idx, sig in enumerate(additional_signals):
        result = div_mgr.can_open_position(
            signal=sig,
            current_positions=open_positions,
            portfolio_value=100000
        )

        can_open = result['can_open']
        reason = result.get('reason', 'N/A')

        status = "✅" if can_open else "❌"
        print(f"  {status} Trade: {sig['ticker']} ({sig['direction']})")
        print(f"      Blocked: {not can_open}")
        print(f"      Reason: {reason}")

        if can_open:
            div_mgr.open_position(100, sig, int(__import__('datetime').datetime.now().timestamp()))
        else:
            print(f"      ⚠️  Blocked as expected: {reason}")

    print()

    # Test 5: Test position age
    print("Test 5: Testing position age")
    print("  Positions should block if min_age_hours = 12")
    print("  (For test, this may pass immediately)")

    # Test 6: Test cooldown after loss
    print("Test 6: Testing cooldown after loss")
    print("  Close position and test if cooldown applies")

    if open_positions:
        # Close first position (ID 0)
        first_id = 0
        div_mgr.close_position(
            first_id,
            open_positions[0].get('entry_price', 100),
            -100,
            int(__import__('datetime').datetime.now().timestamp())
        )

        print(f"  Closed first position {first_id}")
        print(f"  Re-checking diversification...")

        result = div_mgr.can_open_position(
            signal=additional_signals[0],
            current_positions=div_mgr.get_open_positions(),
            portfolio_value=100000
        )

        can_open = result['can_open']
        reason = result.get('reason', 'N/A')

        print(f"  Blocked: {not can_open}")
        print(f"  Reason: {reason}")

    print()

    # Test 7: Diversification report
    print("Test 7: Diversification report")
    from trade.diversification_manager import create_diversification_report
    report = create_diversification_report(
        current_positions=div_mgr.get_open_positions(),
        open_count=len(div_mgr.get_open_positions()),
        max_positions=div_mgr.config['max_open_positions']
    )
    print(report)

    print()
    print("="*80)
    print("✅ All tests passed!")
    print("="*80)

    print()
    print("📊 Integration summary:")
    print(f"  - DiversificationManager: {len(open_positions)} open positions")
    print(f"  - Max positions: {div_mgr.config['max_open_positions']}")
    print(f"  - Per ticker: {div_mgr.config['max_trades_per_ticker']}")
    print(f"  - Per direction: {div_mgr.config['max_trades_per_direction']}")
    print(f"  - Cooldown: {div_mgr.config['cooldown_hours_after_loss']}h")

except ImportError as e:
    print(f"❌ Import error: {e}")
    print()
    print("Make sure files exist:")
    print("  - trade/manager.py")
    print("  - trade/diversification_manager.py")

except Exception as e:
    print(f"❌ Error: {e}")
    import traceback
    traceback.print_exc()