#region imports
from AlgorithmImports import *
#endregion


# Your New Python File
class MomentumAlpha(AlphaModel):
    
    def __init__(
        self, 
        long_period=252, #12-month
        short_period=21, #1-month
        long_percent=0.2, #quintile
        short_percent=0.2
    ):
        self.long_period = long_period
        self.short_period = short_period
        self.long_percent = long_percent
        self.short_percent = short_percent
        self.mom_by_symbol = {}
        self.last_month = -1
    

    def update(self, algorithm, data):
        # emit monthly insight for now
        if self.last_month == algorithm.Time.month:
            return []
        self.last_month = algorithm.Time.month

        mom_scores = []
        for symbol in self.mom_by_symbol:
            if self.mom_by_symbol[symbol].is_ready():
                #algorithm.log(f"{algorithm.Time}: {symbol} is ready")
                long_momentum = self.mom_by_symbol[symbol].get_momentum_percent(self.short_period, self.long_period-1)
                short_momentum = self.mom_by_symbol[symbol].get_momentum_percent(0, self.short_period-1)
                mom_scores.append([symbol, long_momentum, short_momentum])
        # size_in_bytes = sys.getsizeof(self.mom_by_symbol[symbol])
        # algorithm.log(f"mom rolling window size: {size_in_bytes}")
                
        if not mom_scores:
            return []
        
        def assign_quintiles(data):
            return pd.qcut(data, 5, labels=[1, 2, 3, 4, 5])

        # algorithm.log(f"{algorithm.Time}: Number of assets available is {len(mom_scores)}")
        mom_scores_df = pd.DataFrame(mom_scores, columns=['symbol', 'long_momentum', 'short_momentum'])
        mom_scores_df['long_momentum_quintile'] = assign_quintiles(mom_scores_df['long_momentum'])
        mom_scores_df['short_momentum_quintile'] = assign_quintiles(mom_scores_df['short_momentum'])

        long_stocks = mom_scores_df[(mom_scores_df['long_momentum_quintile'] == 5) & (mom_scores_df['short_momentum_quintile'] == 1)]
        short_stocks = mom_scores_df[(mom_scores_df['long_momentum_quintile'] == 1) & (mom_scores_df['short_momentum_quintile'] == 5)]

        algorithm.log(f"Long Stocks: {long_stocks['symbol']}")
        algorithm.log(f"Short Stocks: {short_stocks['symbol']}")
        
        insights = []
        for symbol in long_stocks['symbol'].unique():
            insights.append(Insight.Price(symbol, timedelta(days=20), InsightDirection.Up, None, None))
        for symbol in short_stocks['symbol'].unique():
            insights.append(Insight.Price(symbol, timedelta(days=20), InsightDirection.Down, None, None))

        return insights
    


    def on_securities_changed(self, algorithm, changes):
        for security in changes.RemovedSecurities:
            # algorithm.log(f"{algorithm.time}: Removed {security.Symbol}")
            if security.Symbol in self.mom_by_symbol:
                del self.mom_by_symbol[security.Symbol]
                
        for security in changes.AddedSecurities:
            # algorithm.log(f"{algorithm.time}: Added {security.Symbol}")
            if security not in self.mom_by_symbol:
                history_by_symbol = algorithm.History(security.Symbol, self.long_period, Resolution.DAILY)
                self.mom_by_symbol[security.Symbol] = LongMomentumShortReversalIndicator(security.Symbol, self.long_period, algorithm, history_by_symbol)
                #register indicator for daily update
                algorithm.RegisterIndicator(security.Symbol, self.mom_by_symbol[security.Symbol], Resolution.DAILY)
                del history_by_symbol
                
class LongMomentumShortReversalIndicator(PythonIndicator):
    
    def __init__(self, symbol, period, algorithm, history):
        self.symbol = symbol
        self.period = period
        self.value = 0
        self.rollingWindow = RollingWindow[float](self.period)
        #warm up indicator with history
        if not history.empty:
            for _, tradebar in history.loc[symbol].iterrows():
                self.update(tradebar)
            # algorithm.log(f"{algorithm.Time}: {symbol}")
        
    
    def update(self, bar):
        self.rollingWindow.add(bar.close)
        return self.rollingWindow.is_ready
    
    def is_ready(self):
        return self.rollingWindow.is_ready
    
    # percentage momentum between time_0 and time_n calculated as:
    # (price[time_0] - price[time_n]) / price[time_n]
    def get_momentum_percent(self, time_0, time_n):
        return 100 * (self.rollingWindow[time_0] - self.rollingWindow[time_n]) / self.rollingWindow[time_n]

# region imports
from AlgorithmImports import *
from AlphaModels import MomentumAlpha
# endregion

class CompetitionTemplate(QCAlgorithm):

    def Initialize(self):
        self.set_start_date(2024, 1, 1)
        self.set_cash(1000000)

        self.last_month = -1
        self.num_coarse = 1000
        self.universe_settings.Resolution = Resolution.DAILY
        self.add_universe(self.CoarseSelectionFunction, self.FineSelectionFunction)
        # self.add_alpha(NullAlphaModel())
        self.add_alpha(MomentumAlpha()) #customized alpha model
        self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel(self.IsRebalanceDue))
        self.set_risk_management(NullRiskManagementModel())
        self.set_execution(ImmediateExecutionModel())


    #universe rebalance: monthly
    def IsRebalanceDue(self, time):
        if time.month == self.last_month:
            return None
        self.last_month = time.month
        return time

    #coarse selection: 1000 most liquid equity
    def CoarseSelectionFunction(self, coarse):
        if not self.IsRebalanceDue(self.Time):
            return Universe.UNCHANGED

        sorted_by_dollar_volume = sorted([x for x in coarse if x.HasFundamentalData],
                            key=lambda x: x.DollarVolume, reverse=True)
        return [x.Symbol for x in sorted_by_dollar_volume if x.Price > 5][:self.num_coarse]
    
    #fine selection: sorted by market cap
    def FineSelectionFunction(self, fine):
        #sorted_by_market_cap = sorted(fine, key=lambda f: f.market_cap, reverse=True)
        return [f.Symbol for f in fine]