from AlgorithmImports import *
import numpy as np

class MomentumFactorAlgorithm(QCAlgorithm):
    
    def initialize(self):
        self.set_start_date(2018, 1, 1)
        self.set_end_date(2024, 12, 31)
        self.set_cash(100000)
        
        self.universe_settings.resolution = Resolution.DAILY
        self.add_universe(self.coarse_selection)
        
        self.mom_window    = 252
        self.skip_window   = 21
        self.num_long      = 10
        self.rebalance_day = -1
        
        self.momentum_data = {}

        # Performance tracking
        self.daily_returns        = []
        self.prev_portfolio_value = self.portfolio.total_portfolio_value
        self.peak_value           = self.portfolio.total_portfolio_value
        self.max_drawdown         = 0

    # ------------------------------------------------------------------ #
    #  UNIVERSE                                                            #
    # ------------------------------------------------------------------ #

    def coarse_selection(self, coarse):
        filtered = [x for x in coarse
                    if x.has_fundamental_data
                    and x.price > 10
                    and x.dollar_volume > 5e7]
        sorted_by_volume = sorted(filtered,
                                  key=lambda x: x.dollar_volume,
                                  reverse=True)
        return [x.symbol for x in sorted_by_volume[:150]]

    def on_securities_changed(self, changes):
        for security in changes.added_securities:
            self.momentum_data[security.symbol] = True

        for security in changes.removed_securities:
            self.momentum_data.pop(security.symbol, None)
            if self.portfolio[security.symbol].invested:
                self.liquidate(security.symbol)

    # ------------------------------------------------------------------ #
    #  REBALANCE                                                           #
    # ------------------------------------------------------------------ #

    def on_data(self, data: Slice):
        if self.time.month == self.rebalance_day:
            return
        self.rebalance_day = self.time.month
        self._rebalance(data)

    def _rebalance(self, data):
        scores = {}
        vols   = {}

        for symbol in list(self.momentum_data.keys()):
            if symbol not in data.bars:
                continue

            history = self.history(symbol,
                                   self.mom_window + self.skip_window,
                                   Resolution.DAILY)
            if history.empty or len(history) < self.mom_window + self.skip_window:
                continue

            closes       = history['close']
            past_price   = closes.iloc[0]
            recent_price = closes.iloc[-(self.skip_window + 1)]

            if past_price <= 0:
                continue

            scores[symbol] = (recent_price - past_price) / past_price

            # 21-day realized vol for position sizing
            daily_rets     = closes.pct_change().dropna().tail(21)
            vol            = daily_rets.std()
            vols[symbol]   = vol if vol > 0 else 1e-6

        if len(scores) < self.num_long:
            return

        ranked       = sorted(scores.items(), key=lambda x: x[1], reverse=True)
        long_symbols = [s for s, _ in ranked[:self.num_long]]

        # Exit stale positions
        for symbol in self.portfolio.keys():
            if self.portfolio[symbol].invested:
                if symbol not in long_symbols:
                    self.liquidate(symbol)

        # Vol-adjusted sizing — inverse vol weighting, 5% cash buffer
        inv_vols = {s: 1.0 / vols[s] for s in long_symbols if s in vols}
        total    = sum(inv_vols.values())

        for symbol in long_symbols:
            weight = inv_vols.get(symbol, 1.0 / self.num_long)
            self.set_holdings(symbol, 0.95 * weight / total)

    # ------------------------------------------------------------------ #
    #  PERFORMANCE TRACKING                                                #
    # ------------------------------------------------------------------ #

    def on_end_of_day(self, symbol):
        if symbol != next(iter(self.active_securities.keys()), None):
            return

        current_value = self.portfolio.total_portfolio_value

        if self.prev_portfolio_value > 0:
            daily_ret = (current_value - self.prev_portfolio_value) / self.prev_portfolio_value
            self.daily_returns.append(daily_ret)

        if current_value > self.peak_value:
            self.peak_value = current_value
        drawdown = (self.peak_value - current_value) / self.peak_value
        if drawdown > self.max_drawdown:
            self.max_drawdown = drawdown

        self.prev_portfolio_value = current_value

    def on_end_of_algorithm(self):
        returns = np.array(self.daily_returns)

        if len(returns) < 2:
            print("Not enough data for stats")
            return

        total_return = (self.portfolio.total_portfolio_value - 100000) / 100000
        ann_return   = (1 + total_return) ** (252 / len(returns)) - 1
        ann_vol      = returns.std() * np.sqrt(252)
        sharpe       = (ann_return - 0.05) / ann_vol if ann_vol > 0 else 0
        wins         = np.sum(returns > 0)
        win_rate     = wins / len(returns)
        calmar       = ann_return / self.max_drawdown if self.max_drawdown > 0 else 0

        print("=" * 50)
        print(f"TOTAL RETURN:    {total_return:.2%}")
        print(f"ANN. RETURN:     {ann_return:.2%}")
        print(f"ANN. VOLATILITY: {ann_vol:.2%}")
        print(f"SHARPE RATIO:    {sharpe:.3f}")
        print(f"MAX DRAWDOWN:    {self.max_drawdown:.2%}")
        print(f"CALMAR RATIO:    {calmar:.3f}")
        print(f"WIN RATE:        {win_rate:.2%}")
        print(f"TRADING DAYS:    {len(returns)}")
        print("=" * 50)