#region imports
from AlgorithmImports import *
#endregion


def CalculateTrendIndicators(self):
    
    top_pct= 0.1

    # Moving average calculation
    moving_averages = {}
    for symbol, prices in self.historical_data.items():
        if len(prices) >= self.lookback:
            moving_averages[symbol] = prices.mean()

    top_symbols = sorted(moving_averages, key=moving_averages.get, reverse=True)[:int(len(moving_averages) * top_pct)]
    
    # # Compounded Return Calculations
    # moving_averages = {}
    # compounded_returns = {}
    # for symbol, prices in self.historical_data.items():
    #     if len(prices) >= self.lookback:
    #         daily_returns = prices.pct_change().dropna()
    #         compounded_return = (1 + daily_returns).prod() - 1
    #         compounded_returns[symbol] = compounded_return

    # top_symbols = sorted(compounded_returns, key=compounded_returns.get, reverse=True)[:int(len(compounded_returns) * top_pct)]

    return top_symbols

#region imports
from AlgorithmImports import *
from pypfopt import BlackLittermanModel
import pandas as pd
import numpy as np
#endregion

def OptimizePortfolio(self, mu, S):

    # Black-Litterman views (neutral in this case)
    Q = pd.Series(index=mu.index, data=mu.values)
    P = np.identity(len(mu.index))

    # Optimize via Black-Litterman
    bl = BlackLittermanModel(S,Q=Q, P=P, pi=mu)
    bl_weights = bl.bl_weights()
    
    # Normalize weights
    total_weight = sum(bl_weights.values())
    normalized_weights = {symbol: weight / total_weight for symbol, weight in bl_weights.items()}
    
    return normalized_weights

#region imports
from pypfopt import risk_models, expected_returns
from AlgorithmImports import *
#endregion

def CalculateRiskParameters(self, top_symbols):

    # Get historical prices for selected symbols
    selected_history = self.History(top_symbols, self.lookback, Resolution.Daily)['close'].unstack(level=0)
        
    mu = expected_returns.mean_historical_return(selected_history)
    S = risk_models.sample_cov(selected_history)

    return mu, S

#region imports
from AlgorithmImports import *
#endregion

def Execute_Trades(self, position_list):

    # Place market orders
    for symbol, weight in position_list.items():
        self.SetHoldings(symbol, weight)


def Exit_Positions(self, position_list):
    
    # Liquidate positions not in the target weights
    for holding in self.Portfolio.Values:
        if holding.Symbol not in position_list and holding.Invested:
            self.Liquidate(holding.Symbol)

# region imports
from AlgorithmImports import *
from Alpha_Models import CalculateTrendIndicators
from Risk_Models import CalculateRiskParameters
from Portfolio_Construction import OptimizePortfolio
from Trade_Execution import Execute_Trades, Exit_Positions
# endregion

class NCSU_Strategy_2024_Q3(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2023, 12, 1)  # Set Start Date
        self.SetEndDate(2024, 12, 31)  # Set End Date
        self.SetCash(1000000)  # Set Strategy Cash

        # ETF to get constituents from
        self.etf = "SPY"

        self.universe_settings.leverage = 2.0

        self.AddEquity(self.etf, Resolution.Daily)
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverseSelection(ETFConstituentsUniverseSelectionModel(self.etf))

        self.historical_data = {}
        self.lookback = 252  # MAX lookback period for moving average calculation

        self.equity_high_water_mark = self.Portfolio.TotalPortfolioValue
        self.drawdown_threshold = 0.04  # 6% drawdown threshold
        self.rebalanced = False

    def OnSecuritiesChanged(self, changes):
        # Evaluate if performance is better by trading out of holdings dropped from ETF
        for security in changes.AddedSecurities:
            self.historical_data[security.Symbol] = self.History(security.Symbol, self.lookback, Resolution.Daily)['close']
        
        for security in changes.RemovedSecurities:
            if security.Symbol in self.historical_data:
                del self.historical_data[security.Symbol]
            
            if self.Portfolio[security.Symbol].Invested:
                self.Liquidate(security.Symbol)
                self.Debug(f"Liquidating {security.Symbol} as it is removed from the ETF")

    def OnData(self, data):
        if not self.rebalanced:
            # Check for drawdown condition
            current_equity = self.Portfolio.TotalPortfolioValue
            if (self.equity_high_water_mark - current_equity) / self.equity_high_water_mark >= self.drawdown_threshold:
                self.Debug(f"Drawdown exceeded {self.drawdown_threshold}. Rebalancing...")
                self.Rebalance()
                self.equity_high_water_mark = current_equity  # Update high water mark

    def Rebalance(self):
        self.Debug(f"Rebalancing on {self.Time}")

        # Alpha Model Output
        sorted_symbols = CalculateTrendIndicators(self)

        # Risk Model Output
        mu, S = CalculateRiskParameters(self, top_symbols=sorted_symbols)
        
        # Reduce mu by transaction costs
        transaction_cost = 0.001  # Assumed transaction cost per trade
        for symbol in sorted_symbols:
            if symbol in mu:
                mu[symbol] -= transaction_cost
            else:
                self.Debug(f"Symbol {symbol} not found in mu")

        # Portfolio Construction Output
        target_positions = OptimizePortfolio(self, mu=mu, S=S)

        Exit_Positions(self, position_list=target_positions)

        Execute_Trades(self, position_list=target_positions)

        self.rebalanced = True

    def OnOrderEvent(self, orderEvent):
        self.rebalanced = False  # Reset rebalanced flag after trades have been executed

    def OnEndOfDay(self):
        # Check for end of day to reset rebalanced flag if necessary
        if self.rebalanced == False:
            self.Rebalance()
            self.rebalanced = True