from AlgorithmImports import *
from AlgorithmImports import *
from datetime import timedelta

class MyAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2022, 1, 1)
        self.SetEndDate(2022, 12, 1)
        self.SetCash(100000)

        self.signals = ["VIXM", "BND", "BIL", "TLT", "SPY"]
        self.risk_on_no_inflation = ["TECL", "TQQQ", "UPRO", "TMF"]
        self.risk_on_falling_rates = ["UPRO", "TMF"]
        self.risk_off_market_crash = ["SHY"]
        self.risk_off_rising_rates = ["USDU", "QID", "TBF"]
        self.refined_risk_off = ["IEI", "GLD", "TIP", "BSV"]

        for symbol in self.signals + self.risk_on_no_inflation + self.risk_on_falling_rates + self.risk_off_market_crash + self.risk_off_rising_rates + self.refined_risk_off:
            self.AddEquity(symbol)
        
        self.current_condition = None
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.AfterMarketOpen("SPY", 10), self.CheckConditions)

    def OnData(self, data):
        pass

    def CheckConditions(self):
        history = self.History(self.signals + self.risk_on_no_inflation + self.risk_on_falling_rates + self.risk_off_market_crash + self.risk_off_rising_rates + self.refined_risk_off, 60, Resolution.Daily)
        
        if history.empty:
            return
        
        last_prices = history['close'].unstack(level=0).iloc[-1]
        prev_prices = history['close'].unstack(level=0).iloc[-60]
        
        vixm_momentum = self.CalculateMomentum(history['close']['VIXM'], 40)
        
        if vixm_momentum > 0.69:
            new_condition = "RiskOffMarketCrash"
        else:
            bnd_return_60 = self.CalculateReturn(last_prices['BND'], prev_prices['BND'], 60)
            bil_return_60 = self.CalculateReturn(last_prices['BIL'], prev_prices['BIL'], 60)
            
            if bnd_return_60 > bil_return_60:
                new_condition = "RiskOnNormal"
            elif bnd_return_60 <= bil_return_60 and self.CheckTLTCondition(last_prices, prev_prices):
                new_condition = "RiskOffRisingRates"
            else:
                spy_drawdown = self.CalculateMaxDrawdown(history['close']['SPY'], 5)
                
                if spy_drawdown < 0.05:
                    new_condition = "RiskOnFallingRates"
                else:
                    new_condition = "RefinedRiskOff"
        
        if new_condition != self.current_condition:
            self.current_condition = new_condition
            self.Rebalance(self.current_condition, history, last_prices)
    
    def CheckTLTCondition(self, last_prices, prev_prices):
        tlt_return_20 = self.CalculateCumulativeReturn(last_prices['TLT'], prev_prices['TLT'], 20)
        bil_return_20 = self.CalculateCumulativeReturn(last_prices['BIL'], prev_prices['BIL'], 20)
        return tlt_return_20 < bil_return_20
    
    def Rebalance(self, condition, history, last_prices):
        self.Log(f"Rebalancing for condition: {condition}")
        self.Liquidate()
        if condition == "RiskOffMarketCrash":
            self.SetHoldings("SHY", 1)
        elif condition == "RiskOnNormal":
            self.RiskOnNormal(history, last_prices)
        elif condition == "RiskOffRisingRates":
            self.RiskOffRisingRates(history)
        elif condition == "RiskOnFallingRates":
            self.SetHoldings("UPRO", 0.55)
            self.SetHoldings("TMF", 0.45)
        elif condition == "RefinedRiskOff":
            self.RefinedRiskOff()
    
    def RiskOnNormal(self, history, last_prices):
        relative_strength = {asset: self.CalculateMomentum(history['close'][asset], 10) for asset in self.risk_on_no_inflation}
        top_assets = sorted(relative_strength, key=lambda x: relative_strength[x], reverse=True)[:3]
        weight = 0.33
        for asset in top_assets:
            self.SetHoldings(asset, weight)
    
    def RiskOffRisingRates(self, history):
        self.SetHoldings("USDU", 0.5)
        qid_tbf_momentum = {asset: self.CalculateMomentum(history['close'][asset], 20) for asset in self.risk_off_rising_rates}
        top_asset = max(qid_tbf_momentum, key=qid_tbf_momentum.get)
        self.SetHoldings(top_asset, 0.5)
    
    def RefinedRiskOff(self):
        weight = 0.25
        for asset in self.refined_risk_off:
            self.SetHoldings(asset, weight)
    
    def CalculateReturn(self, current_price, previous_price, days):
        return (current_price - previous_price) / previous_price
    
    def CalculateCumulativeReturn(self, current_price, previous_price, days):
        return (current_price / previous_price) - 1
    
    def CalculateMaxDrawdown(self, prices, days):
        max_drawdown = 0
        for i in range(1, len(prices) - days + 1):
            drawdown = (prices[i:i+days].min() - prices[i-1]) / prices[i-1]
            max_drawdown = min(max_drawdown, drawdown)
        return max_drawdown
    
    def CalculateMomentum(self, prices, days):
        return prices.pct_change(periods=days).mean()