#region imports
from AlgorithmImports import *
#endregion


# Your New Python File
class MomentumAlpha(AlphaModel):
    
    def __init__(
        self, 
        long_period=252, 
        short_period=21,
        long_percent=0.2, 
        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])
                
        
        if not mom_scores:
            return []
        
        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_rank'] = mom_scores_df['long_momentum'].rank(ascending=False) #high long momentum receives higher(smaller) rank
        mom_scores_df['long_bucket'] = pd.qcut(mom_scores_df['long_rank'], 10, labels=False) + 1
        highest_rank_bucket = mom_scores_df[mom_scores_df['long_bucket'] == 1]
        sorted_by_short_momentum = highest_rank_bucket.sort_values(by='short_momentum', ascending=True)

        insights = []
        num_long = int(self.long_percent * len(sorted_by_short_momentum))
        for index, row in sorted_by_short_momentum.iloc[:num_long].iterrows():
            symbol = row['symbol']
            insights.append(Insight.Price(symbol, timedelta(days=10), InsightDirection.Up))

        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]
                algorithm.liquidate(security.Symbol)
            

        for security in changes.AddedSecurities:
            # algorithm.log(f"{algorithm.time}: Added {security.Symbol}")
            if security not in self.mom_by_symbol:
                self.mom_by_symbol[security.Symbol] = LongMomentumShortReversalIndicator(security.Symbol, period=self.long_period)
                #warm up indicator with history
                history_by_symbol = algorithm.History(security.Symbol, self.long_period, Resolution.DAILY)
                for _, row in history_by_symbol.loc[security.Symbol].iterrows():
                    bar = TradeBar(row.name, security.Symbol, row['open'], row['high'], row['low'], row['close'], row['volume'])
                    self.mom_by_symbol[security.Symbol].update(bar)
                #register indicator for daily update
                algorithm.RegisterIndicator(security.Symbol, self.mom_by_symbol[security.Symbol], Resolution.DAILY)


class LongMomentumShortReversalIndicator(PythonIndicator):
    
    def __init__(self, symbol, period):
        self.symbol = symbol
        self.period = period
        self.value = 0
        self.rollingWindow = RollingWindow[float](self.period)
    
    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 = 500
        self.universe_settings.Resolution = Resolution.DAILY
        self.add_universe(self.CoarseSelectionFunction, self.FineSelectionFunction)
        self.add_alpha(MomentumAlpha()) #customized alpha model
        self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel(self.IsRebalanceDue))
        self.set_risk_management(NullRiskManagementModel())
        self.set_execution(ImmediateExecutionModel())

        self.set_warm_up(7)


    #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

        selected = sorted([x for x in coarse if x.HasFundamentalData],
                            key=lambda x: x.DollarVolume, reverse=True)
        return [x.Symbol for x in selected[: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 sorted_by_market_cap]