# region imports
from AlgorithmImports import *
# endregion

# Source: https://www.investopedia.com/articles/trading/04/091504.asp

class KellyCriterionSMACrossoverAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2010, 1, 1)
        # Add the risky and risk-free assets.
        self._risk_asset = self.add_equity("SPY", Resolution.HOUR)
        self._rf_asset = self.add_equity('SHY', Resolution.HOUR)
        # Add some strategy-specific indicators/variables.
        self._risk_asset.short_sma = self.sma(self._risk_asset.symbol, 1)
        self._risk_asset.long_sma = self.sma(self._risk_asset.symbol, 6)
        # Create the KellyCriterion object.
        self._risk_asset.signal = 0
        self._kelly_criterion = KellyCriterion(self.get_parameter('kelly_period', 10))
        # Add a warm-up period so we some historical performance of the strategy
        # once we start trading.
        self.set_warm_up(timedelta(365))

    def on_data(self, data: Slice):
        # Wait until the market is open.
        if not data.bars or not self.is_market_open(self._risk_asset.symbol):
            return
        # Pass the latest signal to the KellyCriterion object.
        if not self._risk_asset.signal and self._risk_asset.short_sma > self._risk_asset.long_sma:
            self._risk_asset.signal = 1
            self._kelly_criterion.update_signal(1, self._risk_asset.price)
        elif self._risk_asset.signal and self._risk_asset.short_sma < self._risk_asset.long_sma:
            self._risk_asset.signal = 0
            self._kelly_criterion.update_signal(0, self._risk_asset.price)
        # Wait until we can trade.
        if self.is_warming_up or not self._kelly_criterion.is_ready:
            return
        # Update the portfolio holdings based on the signal.
        if self._risk_asset.signal and not self._risk_asset.holdings.is_long:
            # Cap the exposure at 200% to avoid errors.
            weight = min(2, self._kelly_criterion.weight())
            self.set_holdings(
                [
                    PortfolioTarget(self._risk_asset.symbol, weight),
                    # If the target weight for the risky asset is <1, then raise the porfolio
                    # exposure to 100% with the risk-free asset.
                    PortfolioTarget(self._rf_asset.symbol, 0 if weight > 1 else 1-weight)
                ]
            )
        elif not self._risk_asset.signal and self._risk_asset.holdings.is_long:
            # If the signal is 0, put 100% of the portfolio in the risk-free asset.
            self.set_holdings([PortfolioTarget(self._rf_asset.symbol, 1)], True)


class KellyCriterion:

    def __init__(self, period):
        self._period = period
        self._trades = np.array([])

    def update_signal(self, signal, price):
        if signal: # Enter
            self._entry_price = price
        else: # Exit
            self._trades = np.append(self._trades, [price - self._entry_price])[-self._period:]

    def weight(self):
        # Wait until there are enough trade samples.
        if not self.is_ready:
            return None
        # Calculate the Kelly %.
        wins = self._trades[self._trades > 0]
        losses = self._trades[self._trades < 0]
        if not wins.sum() or not losses.sum():
            return 0
        win_loss_ratio = wins.mean() / losses.mean()
        winning_probability = len(wins) / self._period
        return winning_probability - (1-winning_probability)/win_loss_ratio
    
    @property
    def is_ready(self):
        return len(self._trades) == self._period