from AlgorithmImports import *
from datetime import timedelta
import numpy as np

class AdaptiveDayTradingStrategy(QCAlgorithm):
    """
    Day Trading Strategy that implements:
    - Market regime detection using SPY 200 EMA
    - RSI-based fund surfing among different asset groups
    - Volatility hedging with UUP+GLD weighted by inverse volatility
    - Overbought/oversold triggers using multiple timeframes
    """
    
    def Initialize(self):
        """Initialize algorithm settings, data subscriptions, and indicators"""
        # Set start date and cash
        self.SetStartDate(2019, 1, 1)
        self.SetCash(100000)
        
        # Dictionary to store symbols
        self.symbols = {}
        
        # Add all required equities
        tickers = ["SPY", "QQQ", "TQQQ", "SPXL", "SQQQ", "UVXY", "UUP", "GLD", "SHY"]
        for ticker in tickers:
            self.symbols[ticker] = self.AddEquity(ticker, Resolution.Minute).Symbol
        
        # Dictionary to store indicators
        self.indicators = {}
        
        # Initialize indicators for each symbol
        for ticker in tickers:
            symbol = self.symbols[ticker]
            
            # RSI indicators
            self.indicators[f"{ticker}_rsi_10"] = self.RSI(symbol, 10)
            self.indicators[f"{ticker}_rsi_20"] = self.RSI(symbol, 20)
            
            # EMAs and SMAs
            if ticker == "SPY":
                self.indicators[f"{ticker}_ema_200"] = self.EMA(symbol, 200)
            if ticker == "TQQQ":
                self.indicators[f"{ticker}_sma_20"] = self.SMA(symbol, 20)
                
        # Volatility windows for inverse vol weighting
        self.volatility_window = 30
        
        # Schedule execution times
        self.Schedule.On(self.DateRules.EveryDay(), 
                        self.TimeRules.AfterMarketOpen("SPY", 1), 
                        self.CheckSignalsAndTrade)
        
        self.Schedule.On(self.DateRules.EveryDay(), 
                        self.TimeRules.BeforeMarketClose("SPY", 5), 
                        self.ClosePositions)

    def CheckSignalsAndTrade(self):
        """
        Main trading logic implementation following the strategy flowchart
        """
        if not self.AreIndicatorsReady():
            return
        
        # Clear any existing positions first
        self.Liquidate()
        
        # Check market regime using SPY 200 EMA
        spy_price = self.Securities["SPY"].Price
        spy_200_ema = self.indicators["SPY_ema_200"].Current.Value
        
        # Execute the appropriate market regime logic
        if spy_price > spy_200_ema:
            self.ExecuteBullMarketLogic()
        else:
            self.ExecuteBearMarketLogic()

    def ExecuteBullMarketLogic(self):
        """
        Bull Market Logic Branch (SPY > 200 EMA)
        Exactly follows the PlantUML specification
        """
        # Check TQQQ overbought
        if self.GetRSI("TQQQ", 10) > 79:
            self.SetHoldings("UVXY", 1.0)
            return
            
        # Check SPXL overbought
        if self.GetRSI("SPXL", 10) > 79:
            self.SetHoldings("UVXY", 1.0)
            return
            
        # Check QQQ 5-day decline
        qqq_5d_return = self.GetPeriodReturn("QQQ", 5)
        
        if qqq_5d_return < -0.06:  # Less than -6%
            tqqq_1d_return = self.GetPeriodReturn("TQQQ", 1)
            
            if tqqq_1d_return > 0.05:  # Greater than 5%
                self.ExecuteVolatilityHedge()
            elif self.GetRSI("TQQQ", 10) > 31:
                self.ExecuteVolatilityHedge()
            else:
                self.ExecuteFundSurf(["SHY", "TQQQ", "SPXL"])
        else:
            if self.GetRSI("QQQ", 10) < 31:
                self.ExecuteFundSurf(["SHY", "TQQQ", "SPXL"])
            else:
                self.ExecuteFundSurf(["QQQ", "SPY"])

    def ExecuteBearMarketLogic(self):
        """
        Bear Market Logic Branch (SPY <= 200 EMA)
        Exactly follows the PlantUML specification
        """
        # First check TQQQ oversold condition
        if self.GetRSI("TQQQ", 10) < 31:
            self.ExecuteFundSurf(["SHY", "TQQQ", "SPXL"])
            return
        
        # Check UVXY overbought conditions
        uvxy_rsi = self.GetRSI("UVXY", 10)
        if uvxy_rsi > 74:
            if uvxy_rsi > 84:
                self.ExecuteVolatilityHedge()
            else:
                self.SetHoldings("UVXY", 1.0)
            return
        
        # Check TQQQ price vs SMA
        tqqq_price = self.Securities["TQQQ"].Price
        tqqq_sma = self.indicators["TQQQ_sma_20"].Current.Value
        
        if tqqq_price > tqqq_sma:
            if self.GetRSI("SQQQ", 10) < 31:
                self.ExecuteVolatilityHedge()
            else:
                self.ExecuteFundSurf(["SHY", "TQQQ", "SPXL"])
        else:
            self.ExecuteVolatilityHedge()

    def ExecuteVolatilityHedge(self):
        """
        Implements volatility hedging using UUP and GLD
        Weighted by 30-day inverse volatility
        """
        # Calculate 30-day volatilities
        uup_vol = self.GetHistoricalVolatility("UUP", self.volatility_window)
        gld_vol = self.GetHistoricalVolatility("GLD", self.volatility_window)
        
        # Calculate inverse volatility weights
        total_inv_vol = (1/uup_vol + 1/gld_vol)
        uup_weight = (1/uup_vol) / total_inv_vol
        gld_weight = (1/gld_vol) / total_inv_vol
        
        # Set positions
        self.SetHoldings("UUP", uup_weight)
        self.SetHoldings("GLD", gld_weight)

    def ExecuteFundSurf(self, tickers):
        """
        Implements fund surfing by selecting the asset with lowest 20-day RSI
        
        Parameters:
        tickers (list): List of tickers to consider for fund surfing
        """
        lowest_rsi = float('inf')
        selected_ticker = None
        
        # Find the ticker with the lowest 20-day RSI
        for ticker in tickers:
            current_rsi = self.GetRSI(ticker, 20)
            if current_rsi < lowest_rsi:
                lowest_rsi = current_rsi
                selected_ticker = ticker
        
        if selected_ticker:
            self.SetHoldings(selected_ticker, 1.0)

    def ClosePositions(self):
        """Close all positions at end of day"""
        self.Liquidate()

    def GetRSI(self, ticker, period):
        """Helper method to get current RSI value"""
        return self.indicators[f"{ticker}_rsi_{period}"].Current.Value

    def GetHistoricalVolatility(self, ticker, window):
        """Calculate historical volatility for the specified window"""
        history = self.History(self.symbols[ticker], window, Resolution.Daily)
        if len(history) < window:
            return float('inf')
        returns = np.diff(np.log(history['close'].values))
        return np.std(returns) * np.sqrt(252)  # Annualized

    def GetPeriodReturn(self, ticker, days):
        """Calculate return over specified number of days"""
        history = self.History(self.symbols[ticker], days + 1, Resolution.Daily)
        if len(history) < days + 1:
            return 0
        prices = history['close'].values
        return (prices[-1] / prices[0]) - 1

    def AreIndicatorsReady(self):
        """Check if all indicators have enough data"""
        for indicator in self.indicators.values():
            if not indicator.IsReady:
                return False
        return True

    def OnData(self, data):
        """OnData event handler. Required but not used."""
        pass