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

class TBMR_PA_Adaptive(QCAlgorithm):
    def PredictPrice(self, df):
        close = df['close']
        
        # Trend indicators
        recent_mean = close.rolling(10).mean().iloc[-1]
        short_mean = close.rolling(5).mean().iloc[-1]
        long_mean = close.rolling(20).mean().iloc[-1]

        # Momentum / slope
        momentum = (short_mean - long_mean) / long_mean

        # Volatility
        vol = close.pct_change().rolling(10).std().iloc[-1]

        # Dynamic weights based on momentum and volatility
        # Strong upward trend -> rely more on last price
        # High volatility -> rely more on recent mean
        w_last = min(max(0.3 + momentum*2 - vol*5, 0), 1)
        w_mean = 1 - w_last

        # Compute predicted price with small random noise
        noise = np.random.normal(0, 0.002)
        predicted = w_mean * recent_mean + w_last * close.iloc[-1]
        return predicted * (1 + noise)

    def Initialize(self):
        self.SetStartDate(2023,1,1)
        self.SetEndDate(2025,1,1)
        self.SetCash(100000)

        tickers = [
            "AAPL","MSFT","AMZN","GOOG","META","TSLA","NVDA","JPM","V","UNH",
            "HD","PG","MA","DIS","PYPL","ADBE","NFLX","INTC","CRM","KO",
            "PFE","MRK","PEP","CSCO","ABBV","AVGO","T","XOM","CVX","NKE",
            "ORCL","ACN","MCD","ABT","BMY","COST","WMT","QCOM","TXN","LIN",
            "NEE","HON","LOW","UPS","PM","RTX","INTU","SBUX","GS","AMGN","MDT"
        ]
        
        self.symbols = [self.AddEquity(t, Resolution.Daily).Symbol for t in tickers]
        self.lookback = 50
        self.k_multiplier = 1.4  # Dynamic entry/exit threshold multiplier
        self.min_holding_days = 3
        
        # Store last entry date for each symbol
        self.last_entry = {sym: self.Time - timedelta(days=10) for sym in self.symbols}

        # Schedule TBM-PA daily run
        self.Schedule.On(self.DateRules.EveryDay(), 
                         self.TimeRules.AfterMarketOpen("SPY", 5), 
                         self.RunTBMR)

    def RunTBMR(self):
        active_trades = []
        current_time = self.Time

        for sym in self.symbols:
            history = self.History(sym, self.lookback, Resolution.Daily)
            if history.empty or len(history) < self.lookback:
                continue

            df = pd.DataFrame(history['close'])
            df.columns = ['close']

            predicted = self.PredictPrice(df)
            current_price = df['close'].iloc[-1]
            deviation = (current_price - predicted)/predicted

            # Calculate volatility
            sigma = df['close'].pct_change().rolling(10).std().iloc[-1]
            sigma = sigma if not np.isnan(sigma) else 0.01

            # Dynamic thresholds
            entry_thresh = -self.k_multiplier * sigma
            exit_thresh = self.k_multiplier * sigma

            # Minimum holding period check
            days_since_entry = (current_time - self.last_entry[sym]).days

            # Entry
            if deviation < entry_thresh and not self.Portfolio[sym].Invested and days_since_entry >= self.min_holding_days:
                active_trades.append((sym, sigma))
                self.last_entry[sym] = current_time

            # Exit
            elif deviation > exit_thresh and self.Portfolio[sym].Invested:
                self.Liquidate(sym)

            # Plot tracking
            self.Plot(f"TBMR-{sym.Value}", "Predicted", predicted)
            self.Plot(f"TBMR-{sym.Value}", "Price", current_price)

        # ===== Portfolio Allocation =====
        if active_trades:
            # Allocate inversely to volatility
            total_inverse_vol = sum(1/s for _, s in active_trades)
            for sym, sigma in active_trades:
                weight = (1/sigma)/total_inverse_vol * 0.95  # max 95% capital
                self.SetHoldings(sym, weight)

        self.Plot("Portfolio","Equity",self.Portfolio.TotalPortfolioValue)