"""Kelly Criterion position sizing for prediction market betting.

Implements fractional Kelly to control risk while maximizing
log-wealth growth based on model edge vs market implied probability.
"""

import numpy as np
from dataclasses import dataclass

from config import KELLY_FRACTION, MAX_POSITION_USDC


@dataclass
class KellyResult:
    """Result of Kelly sizing calculation."""
    full_kelly_fraction: float    # Fraction of bankroll to bet (full Kelly)
    fractional_kelly: float       # Fraction after applying Kelly fraction
    size_usdc: float              # Absolute bet size in USDC
    kelly_active: bool            # Whether full Kelly recommends a bet
    edge: float                   # Edge in decimal (not bps)
    log_utility: float            # Expected log-utility gain


class KellyCriterion:
    """
    Kelly Criterion for binary prediction markets.

    For binary markets:
        f* = p - q / (b)
    where:
        p = our estimated probability of winning
        q = 1 - p
        b = net odds received (payout / bet - 1)

    In prediction markets:
        If we buy YES at price P, and it resolves YES, we get 1.
        So b = (1-P)/P if buying YES, or P/(1-P) if buying NO.
    """

    def __init__(self, bankroll_usdc: float = 1000.0, fraction: float = KELLY_FRACTION):
        self.bankroll = bankroll_usdc
        self.fraction = fraction

    def size_bet(
        self,
        our_probability: float,     # Our probability (0-100 or 0-1)
        market_probability: float,  # Market probability (0-100 or 0-1)
        side: str = "buy_yes",      # "buy_yes" or "buy_no"
        max_size: float = MAX_POSITION_USDC,
    ) -> KellyResult:
        """
        Calculate Kelly-optimal bet size.

        our_probability: Our model's probability of YES outcome (0-1 or 0-100)
        market_probability: Market-implied probability of YES outcome (0-1 or 0-100)
        """
        # Normalize to 0-1 range
        if our_probability > 1:
            our_probability /= 100.0
        if market_probability > 1:
            market_probability /= 100.0

        # Clamp to avoid division by zero or log(0)
        our_probability = np.clip(our_probability, 0.001, 0.999)
        market_probability = np.clip(market_probability, 0.001, 0.999)

        if side == "buy_yes":
            # Buy YES: we win 1-P per share at cost P
            b = (1.0 - market_probability) / market_probability  # Net odds
            p = our_probability
            q = 1.0 - our_probability
        else:
            # Buy NO: symmetric
            b = market_probability / (1.0 - market_probability)
            p = 1.0 - our_probability   # We win if NO
            q = our_probability

        # Kelly formula: f* = (p * b - q) / b = p - q/b
        if b > 0:
            full_kelly = p - q / b
        else:
            full_kelly = 0.0

        # Edge in decimal
        if side == "buy_yes":
            edge = our_probability - market_probability
        else:
            edge = (1.0 - our_probability) - (1.0 - market_probability)
            edge = market_probability - our_probability  # Same thing

        # Only bet when we have positive edge
        kelly_active = full_kelly > 0.001

        if not kelly_active:
            return KellyResult(
                full_kelly_fraction=0.0,
                fractional_kelly=0.0,
                size_usdc=0.0,
                kelly_active=False,
                edge=edge,
                log_utility=0.0,
            )

        # Apply fraction for safety
        fractional_kelly = full_kelly * self.fraction
        size_usdc = min(fractional_kelly * self.bankroll, max_size)

        # Log utility uses fractions of bankroll
        log_utility = self._expected_log_utility(
            p_win=our_probability,
            market_price=market_probability,
            side=side,
            bet_fraction=min(fractional_kelly, 0.99) if kelly_active else 0.0,
        )

        return KellyResult(
            full_kelly_fraction=full_kelly,
            fractional_kelly=fractional_kelly,
            size_usdc=size_usdc,
            kelly_active=kelly_active,
            edge=edge,
            log_utility=log_utility,
        )

    def compare_sides(
        self,
        our_probability: float,
        market_probability: float,
    ) -> dict:
        """Compare betting YES vs NO and return the better side."""
        yes_result = self.size_bet(our_probability, market_probability, "buy_yes")
        no_result = self.size_bet(our_probability, market_probability, "buy_no")

        if yes_result.size_usdc > no_result.size_usdc:
            return {
                "recommended_side": "buy_yes",
                "size_usdc": yes_result.size_usdc,
                "edge": yes_result.edge,
                "log_utility": yes_result.log_utility,
            }
        else:
            return {
                "recommended_side": "buy_no",
                "size_usdc": no_result.size_usdc,
                "edge": no_result.edge,
                "log_utility": no_result.log_utility,
            }

    def update_bankroll(self, new_bankroll: float):
        """Update bankroll after wins/losses."""
        self.bankroll = new_bankroll

    @staticmethod
    def _expected_log_utility(
        p_win: float,
        market_price: float,
        side: str,
        bet_fraction: float,
    ) -> float:
        """Calculate expected log-utility (Kelly criterion) of a fractional bet."""
        if bet_fraction <= 0:
            return 0.0

        if side == "buy_yes":
            win_mult = (1.0 - market_price) / market_price
        else:
            win_mult = market_price / (1.0 - market_price)

        # bet_fraction is fraction of bankroll
        # Win: bankroll becomes bankroll * (1 + bet_fraction * win_mult)
        # Lose: bankroll becomes bankroll * (1 - bet_fraction)
        total_after_win = 1.0 + bet_fraction * win_mult
        total_after_loss = 1.0 - bet_fraction

        if total_after_loss <= 0:
            return -999.0

        if side == "buy_yes":
            return p_win * np.log(max(1e-10, total_after_win)) + (1.0 - p_win) * np.log(max(1e-10, total_after_loss))
        else:
            return (1.0 - p_win) * np.log(max(1e-10, total_after_win)) + p_win * np.log(max(1e-10, total_after_loss))