#!/usr/bin/env python3
"""
Kelly Criterion Position Sizing

Implements optimal position sizing based on 0xRicker's strategy.
The Kelly criterion tells us the optimal fraction of capital to risk
on each trade to maximize long-term growth while avoiding ruin.

Key Insights from 0xRicker's $1.3M Profit:
- Optimal Kelly cap: 71% (never risk more than this)
- Use half-Kelly for safety margin
- Adjust position size based on win rate and risk/reward ratio

Formula:
    f* = (p * (b + 1) - 1) / b
    
Where:
    f* = Optimal fraction of capital to risk
    p  = Probability of winning
    b  = Net odds received (win_size / loss_size)

For trading:
    b = expected_return / stop_loss
    p = win_rate (from historical performance or model confidence)
"""

import logging
from typing import Optional, Dict, Any
from dataclasses import dataclass
from datetime import datetime, timedelta
import numpy as np

logger = logging.getLogger(__name__)


@dataclass
class KellyPosition:
    """Represents a calculated position size."""
    symbol: str
    fraction: float  # Fraction of portfolio (0-1)
    absolute_size: float  # Absolute position size in base currency
    risk_amount: float  # Amount at risk in quote currency
    entry_price: float
    stop_loss_price: float
    target_price: float
    win_rate: float
    expected_return: float
    kelly_fraction: float  # Raw Kelly (before cap)
    capped: bool  # Whether position was capped at 71%
    timestamp: datetime
    
    def __str__(self) -> str:
        return (f"KellyPosition({self.symbol}, fraction={self.fraction:.2%}, "
                f"size={self.absolute_size:.4f}, risk={self.risk_amount:.2f})")


class KellyCalculator:
    """
    Kelly criterion position size calculator.
    
    This implements the position sizing strategy from 0xRicker's article:
    - Calculate optimal fraction using win rate and risk/reward
    - Cap at 71% (optimal from empirical results)
    - Use half-Kelly for safety
    - Adjust based on recent performance
    
    Key parameters:
        kelly_cap: Maximum fraction (0.71 from article)
        safety_factor: Reduce Kelly by this factor (0.5 for half-Kelly)
        min_win_rate: Minimum win rate to consider trading (0.55)
        max_position: Maximum position size regardless of Kelly
    """
    
    # Optimal parameters from 0xRicker's results
    KELLY_CAP = 0.71  # Never risk more than 71%
    SAFETY_FACTOR = 0.5  # Use half-Kelly for safety
    MIN_WIN_RATE = 0.55  # Minimum acceptable win rate
    MIN_RISK_REWARD = 1.5  # Minimum risk/reward ratio
    
    def __init__(
        self,
        portfolio_value: float,
        kelly_cap: float = 0.71,
        safety_factor: float = 0.5,
        min_win_rate: float = 0.55,
        max_position_pct: float = 0.90
    ):
        """
        Initialize Kelly calculator.
        
        Args:
            portfolio_value: Total portfolio value in quote currency (USD)
            kelly_cap: Maximum Kelly fraction (default 0.71)
            safety_factor: Safety margin (default 0.5 for half-Kelly)
            min_win_rate: Minimum win rate to trade (default 0.55)
            max_position_pct: Maximum position as % of portfolio
        """
        self.portfolio_value = portfolio_value
        self.kelly_cap = kelly_cap
        self.safety_factor = safety_factor
        self.min_win_rate = min_win_rate
        self.max_position_pct = max_position_pct
        
        # Track recent performance for dynamic adjustment
        self._trade_history: Dict[str, list] = {}
        self._performance_window = 30  # Last 30 trades per symbol
        
        logger.info(
            f"KellyCalculator initialized: portfolio=${portfolio_value:.2f}, "
            f"cap={kelly_cap:.0%}, safety={safety_factor:.0%}"
        )
    
    def calculate_fraction(
        self,
        symbol: str,
        win_rate: float,
        expected_return: float,
        stop_loss_pct: float = 0.03,
        override_kelly: Optional[float] = None
    ) -> float:
        """
        Calculate optimal position fraction using Kelly criterion.
        
        Args:
            symbol: Trading symbol
            win_rate: Probability of winning (0-1)
            expected_return: Expected return as percentage (e.g., 5.0 for 5%)
            stop_loss_pct: Stop loss as percentage (default 3%)
            override_kelly: Override with pre-calculated Kelly fraction
            
        Returns:
            Optimal position fraction (0-1)
        """
        # Validate inputs
        if win_rate < self.min_win_rate:
            logger.warning(
                f"{symbol}: Win rate {win_rate:.0%} below minimum {self.min_win_rate:.0%}"
            )
            return 0.0
        
        if expected_return <= 0:
            logger.warning(f"{symbol}: Non-positive expected return {expected_return}")
            return 0.0
        
        # Calculate risk/reward ratio
        risk_reward = expected_return / stop_loss_pct
        
        if risk_reward < self.MIN_RISK_REWARD:
            logger.warning(
                f"{symbol}: Risk/reward {risk_reward:.2f} below minimum {self.MIN_RISK_REWARD}"
            )
            return 0.0
        
        # Calculate Kelly fraction
        if override_kelly is not None:
            kelly_fraction = override_kelly
        else:
            kelly_fraction = self._calculate_kelly(win_rate, expected_return, stop_loss_pct)
        
        if kelly_fraction <= 0:
            logger.debug(f"{symbol}: Negative Kelly fraction, no position")
            return 0.0
        
        # Apply safety factor (half-Kelly)
        safe_fraction = kelly_fraction * self.safety_factor
        
        # Cap at maximum
        capped_fraction = min(safe_fraction, self.kelly_cap)
        
        # Cap at maximum position percentage
        final_fraction = min(capped_fraction, self.max_position_pct)
        
        logger.info(
            f"{symbol}: Kelly={kelly_fraction:.1%}, Safe={safe_fraction:.1%}, "
            f"Final={final_fraction:.1%} (capped={capped_fraction < safe_fraction})"
        )
        
        return final_fraction
    
    def calculate_position(
        self,
        symbol: str,
        entry_price: float,
        stop_loss_price: float,
        target_price: float,
        win_rate: float,
        expected_return: float,
        model_confidence: Optional[float] = None
    ) -> Optional[KellyPosition]:
        """
        Calculate complete position details.
        
        Args:
            symbol: Trading symbol
            entry_price: Entry price in quote currency
            stop_loss_price: Stop loss price
            target_price: Take profit target
            win_rate: Probability of winning (0-1)
            expected_return: Expected return percentage
            model_confidence: Confidence from opportunity detector (0-1)
            
        Returns:
            KellyPosition with all details, or None if position should not be taken
        """
        # Adjust win rate by model confidence if provided
        if model_confidence is not None:
            # Blend model confidence with historical win rate
            adjusted_win_rate = win_rate * 0.7 + model_confidence * 0.3
        else:
            adjusted_win_rate = win_rate
        
        # Calculate stop loss percentage
        stop_loss_pct = (entry_price - stop_loss_price) / entry_price
        
        # Calculate expected return percentage
        if expected_return <= 0:
            expected_return_pct = (target_price - entry_price) / entry_price * 100
        else:
            expected_return_pct = expected_return
        
        # Calculate Kelly fraction
        fraction = self.calculate_fraction(
            symbol=symbol,
            win_rate=adjusted_win_rate,
            expected_return=expected_return_pct,
            stop_loss_pct=abs(stop_loss_pct)
        )
        
        if fraction <= 0:
            return None
        
        # Calculate absolute position size
        position_value = self.portfolio_value * fraction
        absolute_size = position_value / entry_price  # Base currency amount
        
        # Calculate risk amount
        risk_amount = absolute_size * abs(entry_price - stop_loss_price)
        
        # Calculate raw Kelly for reporting
        raw_kelly = self._calculate_kelly(
            adjusted_win_rate, expected_return_pct, abs(stop_loss_pct)
        )
        
        return KellyPosition(
            symbol=symbol,
            fraction=fraction,
            absolute_size=absolute_size,
            risk_amount=risk_amount,
            entry_price=entry_price,
            stop_loss_price=stop_loss_price,
            target_price=target_price,
            win_rate=adjusted_win_rate,
            expected_return=expected_return_pct,
            kelly_fraction=raw_kelly,
            capped=raw_kelly * self.safety_factor > self.kelly_cap,
            timestamp=datetime.now()
        )
    
    def _calculate_kelly(
        self,
        win_rate: float,
        expected_return_pct: float,
        stop_loss_pct: float
    ) -> float:
        """
        Calculate raw Kelly fraction.
        
        Formula: f* = (p * (b + 1) - 1) / b
        
        Where:
            p = win_rate
            b = expected_return / stop_loss (odds)
        
        Args:
            win_rate: Probability of winning (0-1)
            expected_return_pct: Expected return percentage
            stop_loss_pct: Stop loss percentage
            
        Returns:
            Kelly fraction (can be negative)
        """
        if stop_loss_pct <= 0:
            return 0.0
        
        # Calculate odds (b)
        b = expected_return_pct / stop_loss_pct
        
        if b <= 0:
            return 0.0
        
        # Kelly formula
        kelly = (win_rate * (b + 1) - 1) / b
        
        return max(0, kelly)
    
    def update_with_trade_result(
        self,
        symbol: str,
        win: bool,
        actual_return: float
    ) -> None:
        """
        Update performance tracking with trade result.
        
        Args:
            symbol: Trading symbol
            win: Whether trade was profitable
            actual_return: Actual return percentage
        """
        if symbol not in self._trade_history:
            self._trade_history[symbol] = []
        
        self._trade_history[symbol].append({
            'win': win,
            'return': actual_return,
            'timestamp': datetime.now()
        })
        
        # Keep only recent trades
        self._trade_history[symbol] = self._trade_history[symbol][-self._performance_window:]
    
    def get_symbol_performance(self, symbol: str) -> Dict[str, Any]:
        """
        Get performance statistics for a symbol.
        
        Args:
            symbol: Trading symbol
            
        Returns:
            Dictionary with performance metrics
        """
        trades = self._trade_history.get(symbol, [])
        
        if not trades:
            return {
                'trades': 0,
                'win_rate': 0.0,
                'avg_return': 0.0,
                'total_return': 0.0
            }
        
        wins = sum(1 for t in trades if t['win'])
        returns = [t['return'] for t in trades]
        
        return {
            'trades': len(trades),
            'wins': wins,
            'losses': len(trades) - wins,
            'win_rate': wins / len(trades),
            'avg_return': np.mean(returns),
            'total_return': sum(returns),
            'best_trade': max(returns),
            'worst_trade': min(returns),
        }
    
    def get_dynamic_win_rate(self, symbol: str, base_win_rate: float) -> float:
        """
        Get win rate adjusted by recent performance.
        
        Blends historical base win rate with recent actual performance.
        
        Args:
            symbol: Trading symbol
            base_win_rate: Base win rate from model or historical data
            
        Returns:
            Adjusted win rate
        """
        performance = self.get_symbol_performance(symbol)
        
        if performance['trades'] < 5:
            # Not enough data, use base rate
            return base_win_rate
        
        # Blend: 60% recent performance, 40% base rate
        recent_win_rate = performance['win_rate']
        adjusted = recent_win_rate * 0.6 + base_win_rate * 0.4
        
        logger.debug(
            f"{symbol}: Base win rate={base_win_rate:.1%}, "
            f"Recent={recent_win_rate:.1%}, Adjusted={adjusted:.1%}"
        )
        
        return adjusted
    
    def update_portfolio_value(self, new_value: float) -> None:
        """Update portfolio value after trades."""
        old_value = self.portfolio_value
        self.portfolio_value = new_value
        logger.info(
            f"Portfolio value updated: ${old_value:.2f} → ${new_value:.2f} "
            f"({(new_value/old_value - 1)*100:.2f}%)"
        )
    
    def get_position_summary(self, position: KellyPosition) -> str:
        """
        Get human-readable position summary.
        
        Args:
            position: KellyPosition object
            
        Returns:
            Formatted summary string
        """
        return (
            f"Position: {position.symbol}\n"
            f"  Size: {position.absolute_size:.4f} @ ${position.entry_price:.2f}\n"
            f"  Portfolio fraction: {position.fraction:.1%}\n"
            f"  Risk: ${position.risk_amount:.2f} ({(position.entry_price - position.stop_loss_price)/position.entry_price*100:.1f}%)\n"
            f"  Target: ${position.target_price:.2f} ({(position.target_price - position.entry_price)/position.entry_price*100:.1f}%)\n"
            f"  Stop loss: ${position.stop_loss_price:.2f}\n"
            f"  Win rate: {position.win_rate:.1%}\n"
            f"  Expected return: {position.expected_return:.1f}%\n"
            f"  Kelly fraction: {position.kelly_fraction:.1%} (capped: {position.capped})"
        )


# Convenience function for simple calculations
def calculate_kelly_position(
    portfolio_value: float,
    win_rate: float,
    expected_return: float,
    stop_loss_pct: float = 0.03,
    entry_price: float = 1.0
) -> Dict[str, float]:
    """
    Quick Kelly position calculation.
    
    Args:
        portfolio_value: Total portfolio value
        win_rate: Probability of winning (0-1)
        expected_return: Expected return percentage
        stop_loss_pct: Stop loss percentage
        entry_price: Entry price
        
    Returns:
        Dictionary with position details
    """
    calculator = KellyCalculator(portfolio_value)
    fraction = calculator.calculate_fraction(
        symbol='TEST',
        win_rate=win_rate,
        expected_return=expected_return,
        stop_loss_pct=stop_loss_pct
    )
    
    return {
        'fraction': fraction,
        'position_value': portfolio_value * fraction,
        'size': (portfolio_value * fraction) / entry_price,
        'risk': portfolio_value * fraction * stop_loss_pct
    }


if __name__ == "__main__":
    # Test Kelly calculator
    print("=== Kelly Criterion Calculator Test ===\n")
    
    # Test with typical parameters
    calculator = KellyCalculator(portfolio_value=10000)
    
    # Test case 1: Good opportunity
    print("Test 1: Good opportunity (70% win rate, 10% return)")
    position = calculator.calculate_position(
        symbol='XBT/USD',
        entry_price=50000,
        stop_loss_price=48500,  # 3% stop
        target_price=55000,  # 10% target
        win_rate=0.70,
        expected_return=10.0,
        model_confidence=0.85
    )
    
    if position:
        print(calculator.get_position_summary(position))
    else:
        print("No position recommended")
    
    print()
    
    # Test case 2: Marginal opportunity
    print("Test 2: Marginal opportunity (55% win rate, 5% return)")
    position = calculator.calculate_position(
        symbol='XBT/USD',
        entry_price=50000,
        stop_loss_price=48500,
        target_price=52500,
        win_rate=0.55,
        expected_return=5.0
    )
    
    if position:
        print(calculator.get_position_summary(position))
    else:
        print("No position recommended")
    
    print()
    
    # Test case 3: Bad opportunity (below minimum win rate)
    print("Test 3: Bad opportunity (50% win rate)")
    position = calculator.calculate_position(
        symbol='XBT/USD',
        entry_price=50000,
        stop_loss_price=48500,
        target_price=52500,
        win_rate=0.50,
        expected_return=5.0
    )
    
    if position:
        print(calculator.get_position_summary(position))
    else:
        print("No position recommended (below minimum win rate)")
    
    print()
    
    # Test quick calculation
    print("Test 4: Quick calculation")
    result = calculate_kelly_position(
        portfolio_value=10000,
        win_rate=0.70,
        expected_return=10.0,
        stop_loss_pct=0.03,
        entry_price=50000
    )
    print(f"  Fraction: {result['fraction']:.1%}")
    print(f"  Position value: ${result['position_value']:.2f}")
    print(f"  Size: {result['size']:.4f}")
    print(f"  Risk: ${result['risk']:.2f}")
