# region imports
from AlgorithmImports import *
# endregion
# Andreas Clenow Momentum (Static Assets), Framework

from datetime import timedelta
from collections import deque
from scipy import stats
import numpy as np

# ==============================
# STRATEGY INPUTS
# ==============================
# All relevant strategy variables are declared here for easy customization

# List of tradable tickers (symbols) to be used in the strategy
TRADABLE_SYMBOLS = ['GS', 'JPM', 'HD', 'COST', 'DIS']  

# Period length for the custom momentum indicator (e.g., 50-day lookback)
MOMENTUM_PERIOD = 50  

# Number of top securities to be selected based on momentum (e.g., top 3)
TOP_N_ASSETS = 3  

# Starting capital for the strategy
STARTING_CAPITAL = 100000  

# Start and (optional) end date for the backtest
START_DATE = (2014, 1, 1)
# END_DATE = (2024, 10, 1)  # Uncomment to specify an end date

# ==============================
# END OF STRATEGY INPUTS
# ==============================

class ClenowMomentum(AlphaModel): 
    def __init__(self):
        self.PERIOD = MOMENTUM_PERIOD  # Period for calculating momentum (from strategy inputs)
        self.N = TOP_N_ASSETS  # Number of top assets to generate insights for (from strategy inputs)
        
        self.indi = {}  # Dictionary to store indicators
        self.indi_Update = {}  # Updated indicators
        self.securities = []  # List of securities
        
    def OnSecuritiesChanged(self, algorithm, changes):
        # Handle when the universe of securities changes
        for security in changes.AddedSecurities:
            if security.Symbol.Value == 'SPY':
                continue  # Skip SPY as it's used for benchmarking, not trading
            self.securities.append(security)
            symbol = security.Symbol
            
            # Create and register custom momentum indicator for each security
            self.indi[symbol] = My_Custom('My_Custom', symbol, self.PERIOD)
            algorithm.RegisterIndicator(symbol, self.indi[symbol], Resolution.Daily)
        
            # Warm up the indicator with historical data
            history = algorithm.History(symbol, self.PERIOD, Resolution.Daily)
            self.indi[symbol].Warmup(history)

          
    def Update(self, algorithm, data):
        # Generate trading insights based on updated momentum indicator values
        insights = []

        # Check which indicators are ready for generating signals
        ready = [indicator for symbol, indicator in self.indi.items() if indicator.IsReady]
        
        # Sort the indicators by momentum value and select top N assets
        ordered = sorted(ready, key=lambda x: x.Value, reverse=False)[:self.N]
        
        # Generate an insight for each of the top N securities
        for x in ordered:
            insights.append(Insight.Price(x.symbol, timedelta(1), InsightDirection.Up))
        
        # Plot momentum indicator values for all tradable symbols
        for idx, symbol in enumerate(self.indi.keys()):
            algorithm.Plot('Custom_Slope', f'Value {symbol.Value}', list(self.indi.values())[idx].Value)

        return insights

class FrameworkAlgorithm(QCAlgorithm):

    def Initialize(self):
        # Set the start date and initial capital (from strategy inputs)
        self.SetStartDate(*START_DATE)
        self.SetCash(STARTING_CAPITAL)

        # Create the tradable symbols based on the defined variable
        symbols = [Symbol.Create(t, SecurityType.Equity, Market.USA) for t in TRADABLE_SYMBOLS]
        
        # Use manual universe selection based on the symbols
        self.SetUniverseSelection(ManualUniverseSelectionModel(symbols))    
        self.UniverseSettings.Resolution = Resolution.Daily        
        
        # Add the custom Alpha model
        self.AddAlpha(ClenowMomentum())        
        
        # Set the portfolio construction and execution models
        self.Settings.RebalancePortfolioOnInsightChanges = False          
        self.Settings.RebalancePortfolioOnSecurityChanges = True
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(self.DateRules.Every(DayOfWeek.Monday)))
        self.SetExecution(ImmediateExecutionModel()) 
        
        # Add SPY as the market benchmark for tracking performance
        self.MKT = self.AddEquity('SPY', Resolution.Daily).Symbol
        self.mkt = []

        # Set up a daily trade bar consolidator for SPY
        self.consolidator = TradeBarConsolidator(timedelta(days=1))
        self.consolidator.DataConsolidated += self.consolidation_handler
        self.SubscriptionManager.AddConsolidator(self.MKT, self.consolidator)
        
        # Fetch historical data for SPY
        self.history = self.History(self.MKT, 2, Resolution.Daily)
        self.history = self.history['close'].unstack(level=0).dropna()
        
        
    def consolidation_handler(self, sender, consolidated):
        # Update historical SPY data
        self.history.loc[consolidated.EndTime, consolidated.Symbol] = consolidated.Close
        self.history = self.history.iloc[-2:] 
        
            
    def OnEndOfDay(self):
        # Track and plot SPY performance
        mkt_price = self.history[[self.MKT]].iloc[-1]
        self.mkt.append(mkt_price)
        mkt_perf = self.mkt[-1] / self.mkt[0] * STARTING_CAPITAL  # Use starting capital from inputs
        self.Plot('Strategy Equity', 'SPY', mkt_perf)       
        
        # Plot portfolio leverage
        account_leverage = self.Portfolio.TotalHoldingsValue / self.Portfolio.TotalPortfolioValue
        self.Plot('Holdings', 'leverage', round(account_leverage, 2)) 
        
class My_Custom:
    def __init__(self, name, symbol, period):
        # Custom momentum indicator initialization
        self.symbol = symbol
        self.Name = name
        self.Time = datetime.min
        self.Value = 0
        self.Slope = 0
        self.Corr = 0

        self.queue = deque(maxlen=period)  # Queue to store rolling prices
        self.IsReady = False  # Flag to check readiness of the indicator
        

    def Update(self, input):
        # Update the indicator with the latest price
        return self.Update2(input.Time, input.Close)
        
    
    def Update2(self, time, value):
        # Append the new price and calculate the indicator if enough data is available
        self.queue.appendleft(value)
        count = len(self.queue)
        self.Time = time
        
        self.IsReady = count == self.queue.maxlen
        
        # Perform linear regression to calculate momentum if the indicator is ready
        if self.IsReady:    
            y = np.log(self.queue)  # Log-transformed price data
            x = [range(len(y))]
            reg = stats.linregress(x, y)
            slope, corr = reg[0], reg[2]
            self.Slope = slope 
            self.Corr = corr  
            self.annualized_slope = float(np.power(np.exp(self.Slope), 252) - 1) * 2.00 
            self.Value = (self.annualized_slope) * float(corr**2)

        return self.IsReady 
        
    def Warmup(self,history):
        # Warm up the indicator with historical price data
        for index, row in history.loc[self.symbol].iterrows():
            self.Update2(index, row['close'])