from QuantConnect.Algorithm import QCAlgorithm
from QuantConnect.Data.Market import TradeBar
from QuantConnect import Resolution
from QuantConnect.Indicators import BollingerBands, IndicatorDataPoint
from datetime import timedelta

class PercentB:
    def __init__(self, name, period, standardDeviations):
        self.Name = name
        self.BB = BollingerBands(period, standardDeviations)
        self.WarmUpPeriod = period
        self.current_value = 0
        self.previous_value = None

    def Update(self, input):
        if isinstance(input, TradeBar):
            input = IndicatorDataPoint(input.EndTime, input.Close)

        if not self.BB.Update(input):
            return False

        if self.BB.IsReady:
            self.previous_value = self.current_value
            lower_band = self.BB.LowerBand.Current.Value
            upper_band = self.BB.UpperBand.Current.Value

            if upper_band != lower_band:
                self.current_value = (input.Value - lower_band) / (upper_band - lower_band)
            else:
                self.current_value = 0

        return True

    @property
    def Value(self):
        return self.current_value

    @property
    def IsReady(self):
        return self.BB.IsReady

    @property
    def Current(self):
        return IndicatorDataPoint(self.Name, self.BB.LowerBand.Current.EndTime, self.current_value)

    @property
    def Previous(self):
        if self.previous_value is not None:
            return IndicatorDataPoint(self.Name, self.BB.LowerBand.Current.EndTime, self.previous_value)
        return None


class SPYStockStrategy(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 1, 1)
        self.SetEndDate(2024, 9, 1)
        self.SetCash(50000)  # Starting with $50,000
        
        self.symbol = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.percent_b = PercentB("PercentB", 20, 2)
        self.RegisterIndicator(self.symbol, self.percent_b, Resolution.Daily)
        self.SetWarmUp(20)
        
        self.position = 0
        self.entry_price = 0
        self.entry_spy_price = 0
        self.sold_50_percent = False
        self.allow_new_trade = True
        self.original_position_size = 0
        self.last_trade_date = None
        self.first_entry_made = False
        self.first_entry_price = 0

    def OnData(self, data):
        if self.IsWarmingUp or not self.Securities[self.symbol].HasData:
            return

        current_date = self.Time.date()
        price = round(self.Securities[self.symbol].Close, 2)
        daily_return = self.CalculateDailyReturn()

        if not self.percent_b.IsReady:
            return

        percent_b = self.percent_b.Current.Value
        percent_b_prev = self.percent_b.Previous.Value

        self.Log(f"Date: {current_date}, Price: {price:.2f}, Daily Return: {daily_return:.2%}, %B: {percent_b:.4f}, Prev %B: {percent_b_prev:.4f}")

        if self.position > 0:
            self.ManageExistingPosition(price, current_date, percent_b, percent_b_prev)
        elif self.CanTakeTrade(current_date):
            self.TryEnterTrade(percent_b, percent_b_prev, price, current_date, daily_return)

    def CalculateDailyReturn(self):
        today_price = self.Securities[self.symbol].Close
        yesterday_price = self.History(self.symbol, 2, Resolution.Daily)["close"].values[0]
        return (today_price - yesterday_price) / yesterday_price

    def ManageExistingPosition(self, price, current_date, percent_b, percent_b_prev):
        price_change = (price - self.entry_spy_price) / self.entry_spy_price

        # Check for complete exit conditions
        if price_change <= -0.15 or (percent_b > 1 and percent_b_prev <= 1):
            self.SellStock(1, price, current_date)
            self.Log(f"Exited entire position due to {price_change:.2%} loss or %B crossing above 1 on {current_date}. SPY price: {price:.2f}")
            return

        if price_change >= 0.04 and not self.sold_50_percent:
            self.SellStock(0.5, price, current_date)
            self.sold_50_percent = True
            self.Log(f"Sold 50% of position at {price:.2f} on {current_date} after 4% gain.")

        elif self.sold_50_percent and price_change >= 0.05:
            additional_percent = price_change - 0.04
            sell_fraction = 1 / (2 ** (int(additional_percent * 100)))
            sell_amount = min(self.position, int(self.position * sell_fraction))

            if sell_amount > 0:
                self.SellStock(sell_amount / self.position, price, current_date)
                self.Log(f"Sold {sell_amount} shares (remaining portion) after {additional_percent * 100:.1f}% move on {current_date}. SPY price: {price:.2f}")

        if self.position <= self.original_position_size * 0.5 and self.sold_50_percent:
            self.allow_new_trade = True

        if price_change >= 0.10:
            self.SellStock(1, price, current_date)
            self.Log(f"Exited entire position due to 10% gain on {current_date}. SPY price: {price:.2f}")

    def TryEnterTrade(self, percent_b, percent_b_prev, price, current_date, daily_return):
        entry_condition = percent_b < 0.0 and percent_b_prev >= 0.0 and daily_return < -0.0125

        if entry_condition:
            # Invest 100% of available cash
            cash_to_invest = self.Portfolio.Cash
            quantity = int(cash_to_invest // price)

            if quantity > 0:
                self.MarketOrder(self.symbol, quantity)
                self.position += quantity
                self.original_position_size = quantity
                self.entry_spy_price = price
                self.last_trade_date = current_date
                self.first_entry_made = True
                self.first_entry_price = price
                self.allow_new_trade = False
                self.Log(f"Full entry: Bought {quantity} shares at {price:.2f} on {current_date}. %B: {percent_b:.4f}, Daily Return: {daily_return:.2%}")
            else:
                self.Log(f"Entry condition met but insufficient cash to buy on {current_date}.")
        else:
            self.Log(f"No entry on {current_date}. %B: {percent_b:.4f}, Daily Return: {daily_return:.2%}")

    def SellStock(self, portion, current_price, current_date):
        shares_to_sell = max(1, int(self.position * portion))
        self.MarketOrder(self.symbol, -shares_to_sell)
        self.position -= shares_to_sell

        self.Log(f"Sold {shares_to_sell} shares on {current_date}. SPY price: {current_price:.2f}")

        if self.position == 0:
            self.ResetPositionVariables()

    def CanTakeTrade(self, current_date):
        if self.last_trade_date is None:
            return True
        days_since_last_trade = (current_date - self.last_trade_date).days
        self.Log(f"Days since last trade: {days_since_last_trade}, Allow new trade: {self.allow_new_trade}")
        return days_since_last_trade > 5 and self.allow_new_trade

    def ResetPositionVariables(self):
        self.position = 0
        self.entry_price = 0
        self.entry_spy_price = 0
        self.sold_50_percent = False
        self.allow_new_trade = True
        self.first_entry_made = False
        self.first_entry_price = 0
        self.Log("Position variables reset. Ready for new trade.")