import sys, re

with open('main.py', 'r', encoding='utf-8') as f:
    text = f.read()

old_bonus = '''            # Price Preference Bias: User prefers stocks < , then < 
            price_bonus = 0.0
            current_price = float(df_5m["close"].iloc[-1])
            if current_price < 50:
                price_bonus = 0.50  # Huge boost for first preference
            elif current_price < 100:
                price_bonus = 0.25  # Moderate boost for second preference

            # Alpha-weighted composite ranking score
            rank_score = (
                0.30 * signal_strength +           # pure strategy score
                0.25 * relative_strength +         # RS vs SPY
                0.20 * min(vol_adj_momentum, 1.0) + # vol-adjusted momentum
                0.15 * volume_score +              # volume confirmation
                0.10 * trend_consistency +          # higher-highs persistence
                price_bonus                        # user price preference
            )'''

new_bonus = '''            # Strict Price Filtering
            current_price = float(df_5m["close"].iloc[-1])
            if current_price >= 200:
                logger.debug("%s skipped: price $%.2f >=  limit", symbol, current_price)
                continue
            
            price_tier = 1 if current_price < 100 else 2

            # Alpha-weighted composite ranking score
            rank_score = (
                0.30 * signal_strength +           # pure strategy score
                0.25 * relative_strength +         # RS vs SPY
                0.20 * min(vol_adj_momentum, 1.0) + # vol-adjusted momentum
                0.15 * volume_score +              # volume confirmation
                0.10 * trend_consistency           # higher-highs persistence
            )'''

text = text.replace(old_bonus, new_bonus)

# Update candidates.append
old_append = '''            candidates.append({
                "symbol": symbol,
                "final": final,
                "price_target": price_target,
                "df_5m": df_5m,
                "rank_score": rank_score,
                "signal_strength": signal_strength,
                "relative_strength": relative_strength,
                "vol_adj_momentum": vol_adj_momentum,
                "volume_score": volume_score,
                "trend_consistency": trend_consistency,
                "sector": _get_sector(symbol),
                "atr_pct": atr_pct,
            })'''

new_append = '''            candidates.append({
                "symbol": symbol,
                "final": final,
                "price_target": price_target,
                "df_5m": df_5m,
                "rank_score": rank_score,
                "signal_strength": signal_strength,
                "relative_strength": relative_strength,
                "vol_adj_momentum": vol_adj_momentum,
                "volume_score": volume_score,
                "trend_consistency": trend_consistency,
                "sector": _get_sector(symbol),
                "atr_pct": atr_pct,
                "price_tier": price_tier,
                "current_price": current_price,
            })'''

text = text.replace(old_append, new_append)

# Fix the sorting!
old_sort = '''    candidates.sort(key=lambda c: c["rank_score"], reverse=True)'''
new_sort = '''    candidates.sort(key=lambda c: (c["price_tier"], -c["rank_score"]))'''
text = text.replace(old_sort, new_sort)

old_ml_sort = '''            candidates.sort(key=lambda c: c["final_rank"], reverse=True)'''
new_ml_sort = '''            candidates.sort(key=lambda c: (c["price_tier"], -c["final_rank"]))'''
text = text.replace(old_ml_sort, new_ml_sort)

with open('main.py', 'w', encoding='utf-8') as f:
    f.write(text)