from AlgorithmImports import *
import math
from calendar import monthrange
from pytz import timezone






class SharpeBasedAthenaMultiSymbol(QCAlgorithm):
    def Initialize(self):

       

        self.SetStartDate(2000, 1, 1)
        self.SetEndDate(2020, 1, 1)
        self.SetCash(100000)




        tickers = ["SPY", "SPLG", "SOXX", "IWB", "VV", "XLI", "IVE", "IWD","SMH",'QLD','EWT','VGT','XLK','VUG']
        tickerslev = [ "TQQQ", "SOXL", "UPRO"]
        stocks = ['NFLX','MA','NVDA']
        tickers = tickers + tickerslev+stocks

        #tickers = ['SQQQ','SPXS']

        self.SetWarmUp(timedelta(hours=750)) 

        self.SetBrokerageModel(BrokerageName.Alpaca, AccountType.Margin)

        self.max_leverage_for_account = 1  # max account Leverage BTW this needs to be higher then the leverage
        self.leverage = 1   # Max Leverage for stratagy
        self.MAX_WEIGHT = 0.3*self.leverage 

      

        self.symbols = []
        for ticker in tickers:
            security = self.AddEquity(ticker, Resolution.Hour)
            security.SetLeverage(self.max_leverage_for_account)
            self.symbols.append(security.Symbol)

        self.yearly_returns = {s: [] for s in self.symbols}
        self.yearly_prices = {s: None for s in self.symbols}
        self.excluded_symbols = set()

        self.per = 518
        self.mult = 1.0
        
        self.lookback = 30

        self.price_windows = {s: RollingWindow[float](self.per + 2) for s in self.symbols}
        self.last_filts = {s: None for s in self.symbols}
        self.upward = {s: 0 for s in self.symbols}
        self.downward = {s: 0 for s in self.symbols}
        self.return_windows = {s: RollingWindow[float](self.lookback) for s in self.symbols}
        self.last_prices = {s: None for s in self.symbols}

        self.annual_margin_rate = 0.07
        self.monthly_interest_accrued = 0
        self.last_interest_month = self.Time.month  # track when to bill

        self.pending_orders = []



    def is_after_market_close(self):
        eastern = timezone('US/Eastern')
        current_et = self.Time.astimezone(eastern)
        return current_et.hour >= 16  # 4:00 PM or later

    def is_market_open_morning(self):
        eastern = timezone('US/Eastern')
        current_et = self.Time.astimezone(eastern)
        return current_et.hour == 9 and current_et.minute == 31  # Just after open

    def OnWarmupFinished(self):
        self.Debug("Warm-up complete.")


    def OnData(self, data: Slice):
        sharpe_scores = {}
        profitable = set()




        ###############################################################
        # Prosess Pending Orders

        if self.is_market_open_morning() and self.pending_orders:
            self.Debug(f"Executing {len(self.pending_orders)} queued orders...")
            
            for (symbol, quantity, weight, price, side) in self.pending_orders:
                if side == 'long':
                    self.MarketOrder(symbol, quantity)
                elif side == 'flat':
                    self.Liquidate(symbol)
    
            
            self.pending_orders.clear()

        ###############################################################

        for symbol in self.symbols:
            if not data.Bars.ContainsKey(symbol):
                continue

            if symbol in self.excluded_symbols:
                continue

            bar = data[symbol]
            src = (bar.High + bar.Low) / 2
            self.price_windows[symbol].Add(src)

            last_close = self.last_prices[symbol]
            if last_close is not None and last_close > 0:
                ret = math.log(bar.Close / last_close)
                self.return_windows[symbol].Add(ret)
            self.last_prices[symbol] = bar.Close

            if not self.price_windows[symbol].IsReady or not self.return_windows[symbol].IsReady:
                continue

            window = self.price_windows[symbol]
            price_diffs = [abs(window[i] - window[i - 1]) for i in range(1, self.per + 1)]
            avg_range = sum(price_diffs) / len(price_diffs)
            filt = self.rngfilt(src, avg_range * self.mult, self.last_filts[symbol])
            self.last_filts[symbol] = filt

            if filt is None:
                continue

            if src > filt:
                self.upward[symbol] += 1
                self.downward[symbol] = 0
            elif src < filt:
                self.downward[symbol] += 1
                self.upward[symbol] = 0

            # === Sharpe Ratio + Profitability Filter ===
            returns = list(self.return_windows[symbol])
            cumulative_ret = sum(returns)
            if cumulative_ret <= 0:
                continue  # Still tracked, but not allocated
            profitable.add(symbol)

            mean_ret = cumulative_ret / len(returns)
            std_dev = math.sqrt(sum((r - mean_ret) ** 2 for r in returns) / len(returns))
            if std_dev > 0:
                sharpe_scores[symbol] = mean_ret / std_dev

        if not sharpe_scores:
            return

        total_score = sum(max(0, s) for s in sharpe_scores.values())
        if total_score == 0:
            return

        raw_allocations = {
            s: max(0, sharpe_scores[s]) / total_score
            for s in sharpe_scores if s in profitable
        }

        MIN_ALLOC = 1e-4
        allocations = {s: w for s, w in raw_allocations.items() if w >= MIN_ALLOC}

        MAX_WEIGHT = self.MAX_WEIGHT
        allocations = {s: min(w, MAX_WEIGHT) for s, w in allocations.items()}

        # === Margin-Aware Scaling ===
        available_margin = self.Portfolio.MarginRemaining
        portfolio_value = self.Portfolio.TotalPortfolioValue
        if portfolio_value == 0:
            return

        max_total_allocation = min(available_margin / portfolio_value, self.leverage)
        total_alloc = sum(allocations.values())
        if total_alloc == 0:
            return

        scaling_factor = min(1.0, max_total_allocation / total_alloc)
        allocations = {s: w * scaling_factor for s, w in allocations.items()}

        # === Execute Trading Logic ===
        for symbol in self.symbols:
            if not data.Bars.ContainsKey(symbol):
                continue
            if symbol not in profitable:
                continue  # skip unprofitable symbols

            if not self.price_windows[symbol].IsReady or self.last_filts[symbol] is None:
                continue

            src = self.price_windows[symbol][0]
            filt = self.last_filts[symbol]
            long_condition = src > filt and self.upward[symbol] > 0
            short_condition = src < filt and self.downward[symbol] > 0
            invested = self.Portfolio[symbol].Invested

            if not invested and long_condition and symbol in allocations:
                weight = self.leverage * allocations[symbol]
                target_value = weight * self.Portfolio.TotalPortfolioValue
                current_value = self.Portfolio[symbol].HoldingsValue
                price = self.Securities[symbol].Price

                if price <= 0:
                    continue

                value_diff = target_value - current_value
                quantity = int(value_diff / price)

                if quantity != 0:
                    required_margin = price * abs(quantity) / self.Securities[symbol].Leverage
                    if self.Portfolio.MarginRemaining >= required_margin:
                        if self.is_after_market_close():
                            # Queue it for morning execution
                            self.pending_orders.append((symbol, quantity, weight, price, 'long'))
                            self.Debug(f"Queued {symbol} for next morning: qty={quantity}")
                        elif not self.IsWarmingUp:
                            self.MarketOrder(symbol, quantity)


            elif invested and short_condition:

                if not self.IsWarmingUp:
                    if self.is_after_market_close():
                        # Queue it for morning execution
                        quantity = -self.Portfolio[symbol].Quantity  # sell entire position
                        price = self.Securities[symbol].Price
                        weight = 0  # optional for selling
                        self.pending_orders.append((symbol, quantity, weight, price, 'flat'))
                        self.Debug(f"Queued SELL {symbol} for next morning: qty={quantity}")
                   
                    elif not self.IsWarmingUp:
                        self.Liquidate(symbol)

                
    def rngfilt(self, x, r, prev_filt):
        if prev_filt is None:
            return x
        if x > prev_filt:
            return prev_filt if x - r < prev_filt else x - r
        else:
            return prev_filt if x + r > prev_filt else x + r

    def OnEndOfDay(self):
        holdings_value = self.Portfolio.TotalHoldingsValue
        portfolio_value = self.Portfolio.TotalPortfolioValue
        borrowed = max(0, holdings_value - portfolio_value)

        daily_interest = borrowed * (self.annual_margin_rate / 252)
        self.monthly_interest_accrued += daily_interest

        # Bill at the end of the current month
        last_day = monthrange(self.Time.year, self.Time.month)[1]
        if self.Time.day == last_day:
            self.Portfolio.SetCash(self.Portfolio.Cash - self.monthly_interest_accrued)
            # Optional debug:
            # self.Debug(f"Billing ${self.monthly_interest_accrued:.2f} interest for month {self.Time.month}")
            self.monthly_interest_accrued = 0




        if self.Time.month == 12 and self.Time.day == 31:
            for symbol in self.symbols:
                if symbol not in self.Securities:
                    continue

                price = self.Securities[symbol].Price
                last_price = self.yearly_prices[symbol]

                if last_price is not None and last_price > 0:
                    yearly_return = (price - last_price) / last_price
                    self.yearly_returns[symbol].append(yearly_return)

                    # Check for 2 consecutive negative years
                    if len(self.yearly_returns[symbol]) >= 2:
                        if self.yearly_returns[symbol][-1] < 0 and self.yearly_returns[symbol][-2] < 0:
                            self.excluded_symbols.add(symbol)
                            if not self.IsWarmingUp:

                                self.Liquidate(symbol)  # remove any current position


                # Update for next year
                self.yearly_prices[symbol] = price