Blame for quant_agent_engine.py

"""
QuantAgent OS - Core Self-Improving Engine
Implements the QuantAgent two-layer framework:
1. Inner Reasoning Loop: Writer LLM <-> Judge LLM with Shared Context Buffer
2. Outer Feedback Loop: Empirical Backtesting Oracle + Autonomous Knowledge Base
3. XGBoost Multi-Alpha Downstream Ensemble
"""

import os
import re
import sys
import time
import copy
import logging
import traceback
import numpy as np
import pandas as pd
from typing import Dict, List, Optional, Tuple, Any
from dataclasses import dataclass, field
import xgboost as xgb
from sklearn.model_selection import KFold
from sklearn.metrics import mean_squared_error

logger = logging.getLogger(__name__)


# -----------------------------------------------------------------------------
# Base Factor Class (Compatible with idea-factor design from QuantAgent paper)
# -----------------------------------------------------------------------------

class Factor:
    """Base abstract class for all quantitative signals."""
    def __init__(self):
        self.name: str = "BaseFactor"
        self.window_length: int = 14
        self.inputs: List[str] = ["close", "open", "high", "low", "volume"]

    def calc(self, data: Dict[str, pd.DataFrame]) -> pd.DataFrame:
        """
        Calculates signal matrix.
        Parameters:
            data: dict where keys are field names ('close', 'open', etc.)
                  and values are DataFrames (index=dates, columns=tickers)
        Returns:
            pd.DataFrame of factor scores (index=dates, columns=tickers)
        """
        raise NotImplementedError("Factor subclass must implement calc()")


# -----------------------------------------------------------------------------
# Knowledge Base Record Data Structure
# -----------------------------------------------------------------------------

@dataclass
class AlphaRecord:
    id: str
    name: str
    trading_idea: str
    code: str
    inner_iterations: int
    judge_score: float
    judge_critique: str
    # Real-world backtest metrics (Outer Loop)
    ic_mean: float = 0.0
    ic_ir: float = 0.0
    sharpe: float = 0.0
    annualized_return: float = 0.0
    max_drawdown: float = 0.0
    turnover: float = 0.0
    valid_ratio: float = 0.0
    status: str = "PENDING"  # SUCCESS, FAILED, PROBATION
    timestamp: float = field(default_factory=time.time)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "name": self.name,
            "trading_idea": self.trading_idea,
            "judge_score": self.judge_score,
            "ic_mean": round(self.ic_mean, 4),
            "ic_ir": round(self.ic_ir, 3),
            "sharpe": round(self.sharpe, 2),
            "annualized_return_pct": round(self.annualized_return * 100, 2),
            "max_drawdown_pct": round(self.max_drawdown * 100, 2),
            "turnover_pct": round(self.turnover * 100, 1),
            "valid_ratio": round(self.valid_ratio, 2),
            "status": self.status
        }


# -----------------------------------------------------------------------------
# Backtest Oracle & Quantitative Evaluator (Outer Feedback Loop)
# -----------------------------------------------------------------------------

class AlphaBacktester:
    """Executes factor code against cross-sectional data and evaluates performance."""

    def __init__(self, forward_horizon: int = 1):
        self.forward_horizon = forward_horizon

    def execute_and_evaluate(
        self,
        factor_obj: Factor,
        market_data: Dict[str, pd.DataFrame]
    ) -> Tuple[bool, Dict[str, Any], Optional[pd.DataFrame]]:
        """
        Runs factor.calc() and computes Information Coefficient, Sharpe, and Returns.
        """
        try:
            # 1. Prepare 3D-like aligned cross-sectional dictionaries
            tickers = list(market_data.keys())
            if not tickers:
                return False, {"error": "Empty market data"}, None

            fields = factor_obj.inputs
            aligned_data = {}
            for f in fields:
                field_df = pd.DataFrame({
                    ticker: market_data[ticker][f] 
                    for ticker in tickers 
                    if f in market_data[ticker].columns
                }).ffill().bfill()
                aligned_data[f] = field_df

            # 2. Execute calculation
            signal_df = factor_obj.calc(aligned_data)
            
            if not isinstance(signal_df, pd.DataFrame) or signal_df.empty:
                return False, {"error": "Factor returned invalid or empty DataFrame"}, None

            # 3. Clean signal values
            signal_df = signal_df.replace([np.inf, -np.inf], np.nan)
            valid_ratio = float((~signal_df.isna()).mean().mean())

            if valid_ratio < 0.30:
                return False, {
                    "error": f"Low valid data ratio ({valid_ratio:.1%}). Too many NaNs.",
                    "valid_ratio": valid_ratio
                }, None

            # 4. Compute Forward Returns
            close_matrix = pd.DataFrame({
                ticker: market_data[ticker]['close'] for ticker in tickers
            }).ffill().bfill()
            
            fwd_returns = close_matrix.pct_change(self.forward_horizon).shift(-self.forward_horizon)

            # 5. Compute Cross-Sectional Information Coefficient (IC)
            ic_series = []
            for date in signal_df.index:
                if date in fwd_returns.index:
                    s_row = signal_df.loc[date].dropna()
                    r_row = fwd_returns.loc[date].dropna()
                    common_tickers = s_row.index.intersection(r_row.index)
                    if len(common_tickers) >= 3:
                        s_vals = s_row[common_tickers]
                        r_vals = r_row[common_tickers]
                        if s_vals.std() > 1e-7 and r_vals.std() > 1e-7:
                            corr = np.corrcoef(s_vals, r_vals)[0, 1]
                            if not np.isnan(corr):
                                ic_series.append(corr)

            if not ic_series:
                return False, {"error": "Could not compute cross-sectional correlation"}, None

            ic_series = pd.Series(ic_series)
            ic_mean = float(ic_series.mean())
            ic_std = float(ic_series.std()) if ic_series.std() > 0 else 1.0
            ic_ir = float(ic_mean / ic_std)

            # 6. Simulate Long-Short Strategy (Top 30% long, Bottom 30% short)
            strategy_daily_returns = []
            turnover_list = []
            prev_weights = pd.Series(0.0, index=tickers)

            for date in signal_df.index[:-self.forward_horizon]:
                if date in close_matrix.index:
                    sig_row = signal_df.loc[date].dropna()
                    if len(sig_row) >= 4:
                        # Rank-based weights demeaned
                        ranks = sig_row.rank()
                        weights = (ranks - ranks.mean())
                        if weights.abs().sum() > 0:
                            weights = weights / weights.abs().sum()
                            
                            # Realized forward return
                            date_ret = close_matrix.pct_change().shift(-1).loc[date]
                            common = weights.index.intersection(date_ret.dropna().index)
                            pnl = float((weights[common] * date_ret[common]).sum())
                            strategy_daily_returns.append(pnl)

                            # Turnover
                            turnover = float((weights - prev_weights.reindex(weights.index, fill_value=0.0)).abs().sum())
                            turnover_list.append(turnover)
                            prev_weights = weights

            pnl_series = pd.Series(strategy_daily_returns)
            ann_ret = float(pnl_series.mean() * 252) if len(pnl_series) > 0 else 0.0
            ann_vol = float(pnl_series.std() * np.sqrt(252)) if len(pnl_series) > 0 else 0.01
            sharpe = float(ann_ret / ann_vol) if ann_vol > 0 else 0.0

            # Max Drawdown
            cum_ret = (1 + pnl_series).cumprod()
            rolling_max = cum_ret.cummax()
            drawdowns = (cum_ret - rolling_max) / rolling_max
            mdd = float(abs(drawdowns.min())) if len(drawdowns) > 0 else 0.0

            mean_turnover = float(np.mean(turnover_list)) if turnover_list else 0.0

            metrics = {
                "ic_mean": ic_mean,
                "ic_ir": ic_ir,
                "sharpe": sharpe,
                "annualized_return": ann_ret,
                "annualized_volatility": ann_vol,
                "max_drawdown": mdd,
                "turnover": mean_turnover,
                "valid_ratio": valid_ratio,
                "sample_points": len(ic_series)
            }
            return True, metrics, signal_df

        except Exception as e:
            return False, {"error": f"Execution exception: {str(e)}\n{traceback.format_exc()}"}, None


# -----------------------------------------------------------------------------
# Autonomous Knowledge Base (KB)
# -----------------------------------------------------------------------------

class AutonomousKnowledgeBase:
    """Stores all generated alphas, failure records, and provides vector/keyword retrieval."""

    def __init__(self):
        self.records: Dict[str, AlphaRecord] = {}
        self._seed_initial_knowledge()

    def _seed_initial_knowledge(self):
        """Seeds initial verified quantitative principles and factor templates."""
        sample_idea = "Moving Average Convergence Divergence Momentum with Volume confirmation"
        sample_code = '''
class MACDVolumeSignal(Factor):
    def __init__(self):
        self.name = "MACDVolumeSignal"
        self.window_length = 26
        self.inputs = ["close", "volume"]

    def calc(self, data):
        close = data["close"]
        volume = data["volume"]
        
        # Vectorized EMA calculation
        ema12 = close.ewm(span=12, adjust=False).mean()
        ema26 = close.ewm(span=26, adjust=False).mean()
        macd_line = (ema12 - ema26) / close
        
        # Volume trend ratio
        vol_ma = volume.rolling(20).mean()
        vol_surge = (volume / vol_ma.replace(0, np.nan)).clip(0.5, 3.0)
        
        # Bounded continuous signal in [-1, 1]
        raw_signal = macd_line * vol_surge
        signal = np.tanh(raw_signal * 10).fillna(0.0)
        return signal
'''.strip()
        self.add_record(AlphaRecord(
            id="ALPHA_SEED_001",
            name="MACDVolumeSignal",
            trading_idea=sample_idea,
            code=sample_code,
            inner_iterations=3,
            judge_score=9.2,
            judge_critique="Excellent vectorized implementation with tanh continuous bounding and volume confirmation.",
            ic_mean=0.048,
            ic_ir=0.72,
            sharpe=1.65,
            annualized_return=0.185,
            max_drawdown=0.112,
            turnover=0.28,
            valid_ratio=0.98,
            status="SUCCESS"
        ))

        # Seed a failed negative example for pessimistic learning
        fail_code = '''
class BinaryBreakoutFail(Factor):
    def __init__(self):
        self.name = "BinaryBreakoutFail"
        self.window_length = 10
        self.inputs = ["close"]

    def calc(self, data):
        # Flawed: unvectorized Python loop, unranked binary output 0 or 1
        return (data["close"] > data["close"].shift(1)) * 1.0
'''.strip()
        self.add_record(AlphaRecord(
            id="ALPHA_FAIL_001",
            name="BinaryBreakoutFail",
            trading_idea="Simple 1-day upward price momentum",
            code=fail_code,
            inner_iterations=1,
            judge_score=3.5,
            judge_critique="Critique: Binary 0/1 output without cross-sectional ranking; high churn; zero predictive nuance.",
            ic_mean=-0.012,
            ic_ir=-0.15,
            sharpe=-0.42,
            annualized_return=-0.05,
            max_drawdown=0.22,
            turnover=0.95,
            valid_ratio=0.99,
            status="FAILED"
        ))

    def add_record(self, record: AlphaRecord):
        self.records[record.id] = record

    def retrieve_hints(self, query: str, top_k: int = 2) -> List[Dict[str, Any]]:
        """Retrieves most relevant alpha records for in-context prompting."""
        query_words = set(re.findall(r'\w+', query.lower()))
        scored = []
        for r in self.records.values():
            rec_words = set(re.findall(r'\w+', (r.name + " " + r.trading_idea + " " + r.judge_critique).lower()))
            overlap = len(query_words.intersection(rec_words))
            score = overlap + (1.0 if r.status == "SUCCESS" else -0.5)
            scored.append((score, r))
        
        scored.sort(key=lambda x: x[0], reverse=True)
        return [r.to_dict() for _, r in scored[:top_k]]

    def get_all_records(self) -> List[Dict[str, Any]]:
        return [r.to_dict() for r in self.records.values()]

    def get_success_alphas(self) -> List[AlphaRecord]:
        return [r for r in self.records.values() if r.status == "SUCCESS"]


# -----------------------------------------------------------------------------
# QuantAgent Two-Layer Engine
# -----------------------------------------------------------------------------

class QuantAgentEngine:
    """
    Orchestrates the Inner Reasoning Loop (Writer <-> Judge) and 
    the Outer Feedback Loop (Backtest -> KB Update).
    """

    def __init__(self, llm_adapter=None):
        self.llm = llm_adapter
        self.kb = AutonomousKnowledgeBase()
        self.backtester = AlphaBacktester(forward_horizon=1)

    def run_inner_loop(
        self,
        trading_idea: str,
        max_iterations: int = 3,
        score_threshold: float = 7.5
    ) -> Tuple[str, str, float, List[Dict[str, Any]]]:
        """
        Executes Algorithm 1: The Inner Reasoning Loop.
        Returns: (final_alpha_name, final_code, final_score, trajectory_history)
        """
        history = []
        ctx = {"problem": trading_idea, "kb_hints": []}
        
        # 1. Retrieve knowledge hints
        hints = self.kb.retrieve_hints(trading_idea, top_k=2)
        ctx["kb_hints"] = hints

        current_code = ""
        current_name = "CandidateSignal"
        current_score = 0.0
        current_critique = ""

        for t in range(1, max_iterations + 1):
            iteration_log = {"round": t}

            # --- WRITER STEP ---
            current_name, current_code = self._writer_generate(trading_idea, ctx, current_code, current_critique)
            iteration_log["generated_code"] = current_code
            iteration_log["alpha_name"] = current_name

            # --- JUDGE STEP ---
            current_score, current_critique, advice_list = self._judge_evaluate(current_name, current_code, trading_idea)
            iteration_log["judge_score"] = current_score
            iteration_log["judge_critique"] = current_critique
            iteration_log["advice"] = advice_list

            history.append(iteration_log)

            # Update context buffer
            ctx["last_code"] = current_code
            ctx["last_review"] = current_critique

            if current_score >= score_threshold:
                break

        return current_name, current_code, current_score, history

    def run_outer_loop(
        self,
        trading_idea: str,
        market_data: Dict[str, pd.DataFrame],
        max_inner_rounds: int = 3
    ) -> Dict[str, Any]:
        """
        Executes Algorithm 2: The Outer Feedback Loop.
        1. Runs Inner Loop to synthesize and refine factor code.
        2. Submits code to the Real Backtesting Oracle.
        3. Updates Autonomous KB with metrics and failure/success status.
        """
        # Step 1: Inner Loop
        alpha_name, alpha_code, judge_score, history = self.run_inner_loop(
            trading_idea=trading_idea,
            max_iterations=max_inner_rounds
        )

        # Step 2: Instantiate dynamic code safely
        factor_instance, inst_err = self._instantiate_factor_code(alpha_code, alpha_name)
        if factor_instance is None:
            # Record compilation error in KB for pessimistic learning
            rec_id = f"ALPHA_{int(time.time()*1000)%100000}"
            rec = AlphaRecord(
                id=rec_id,
                name=alpha_name,
                trading_idea=trading_idea,
                code=alpha_code,
                inner_iterations=len(history),
                judge_score=judge_score,
                judge_critique=f"Compilation Failed: {inst_err}",
                status="FAILED"
            )
            self.kb.add_record(rec)
            return {
                "success": False,
                "error": f"Failed to instantiate code: {inst_err}",
                "alpha_name": alpha_name,
                "code": alpha_code,
                "trajectory": history
            }

        # Step 3: Real Environment Backtest (The Outer Feedback Oracle)
        success, metrics, signal_matrix = self.backtester.execute_and_evaluate(factor_instance, market_data)

        # Step 4: Knowledge Base Update
        rec_id = f"ALPHA_{int(time.time()*1000)%100000}"
        if success:
            is_viable = metrics["sharpe"] > 0.5 and abs(metrics["ic_mean"]) > 0.015
            status = "SUCCESS" if is_viable else "PROBATION"
            rec = AlphaRecord(
                id=rec_id,
                name=alpha_name,
                trading_idea=trading_idea,
                code=alpha_code,
                inner_iterations=len(history),
                judge_score=judge_score,
                judge_critique=history[-1]["judge_critique"],
                ic_mean=metrics["ic_mean"],
                ic_ir=metrics["ic_ir"],
                sharpe=metrics["sharpe"],
                annualized_return=metrics["annualized_return"],
                max_drawdown=metrics["max_drawdown"],
                turnover=metrics["turnover"],
                valid_ratio=metrics["valid_ratio"],
                status=status
            )
        else:
            rec = AlphaRecord(
                id=rec_id,
                name=alpha_name,
                trading_idea=trading_idea,
                code=alpha_code,
                inner_iterations=len(history),
                judge_score=judge_score,
                judge_critique=f"Backtest Oracle Rejection: {metrics.get('error', 'Execution error')}",
                status="FAILED"
            )

        self.kb.add_record(rec)

        return {
            "success": success,
            "record": rec.to_dict(),
            "metrics": metrics if success else {},
            "code": alpha_code,
            "trajectory": history,
            "signal_matrix": signal_matrix
        }

    # -------------------------------------------------------------------------
    # Writer and Judge Helpers
    # -------------------------------------------------------------------------

    def _writer_generate(
        self,
        idea: str,
        ctx: Dict[str, Any],
        prev_code: str,
        feedback: str
    ) -> Tuple[str, str]:
        """Generates or refines candidate alpha code via LLM or specialized generator."""
        if self.llm is not None:
            prompt = f"""
You are the Quant Alpha Writer Agent.
Task: Write a Python quantitative trading signal subclassing `Factor`.
Idea: {idea}
Relevant KB Hints: {ctx.get('kb_hints', [])}
Previous Code:
{prev_code if prev_code else 'None (First round)'}
Mentor Feedback on Previous Code:
{feedback if feedback else 'None'}

Rules:
1. Subclass `Factor`.
2. Set self.name, self.window_length, self.inputs.
3. In `calc(self, data)`, use vectorized pandas/numpy operations across all assets.
4. Normalize signal between -1.0 and 1.0 or use np.tanh / z-score. Avoid crude binary 0/1 outputs.
5. Handle NaNs with .fillna(0.0) or -0.1 for absence of pattern.
Return ONLY valid Python code block.
"""
            response = self.llm.generate(prompt)
            code = self._extract_code_block(response)
            name_match = re.search(r'class\s+([A-Za-z0-9_]+)\s*\(\s*Factor\s*\)', code)
            name = name_match.group(1) if name_match else "EvolvedSignal"
            return name, code

        # Intelligent Heuristic Writer if LLM is not hooked up
        return self._heuristic_writer(idea, prev_code, feedback)

    def _judge_evaluate(
        self,
        name: str,
        code: str,
        idea: str
    ) -> Tuple[float, str, List[str]]:
        """Evaluates factor code against rigorous quant criteria."""
        if self.llm is not None:
            prompt = f"""
You are the Senior Quantitative Mentor Judge.
Evaluate this proposed Alpha signal:
Idea: {idea}
Alpha Name: {name}
Code:
{code}

Grade the code on 6 criteria (0-10 total score):
1. Class structure (inherits Factor, defines inputs, window_length, calc)
2. Vectorization efficiency (no slow stock-by-stock loops)
3. Lookahead bias prevention (no shift(-1) or future leakage)
4. Nuance & continuity (continuous bounded signal in [-1, 1], not flat binary 0/1)
5. Numerical stability (.clip, .fillna, no zero division)
6. Alignment with trading intuition.

Output in format:
SCORE: <float 0-10>
CRITIQUE: <brief summary>
ADVICE:
- <point 1>
- <point 2>
"""
            resp = self.llm.generate(prompt)
            score_match = re.search(r'SCORE:\s*([0-9\.]+)', resp)
            score = float(score_match.group(1)) if score_match else 7.0
            critique = resp
            advice = ["Review vectorization and scaling."]
            return score, critique, advice

        # Deterministic Quant Quality Heuristic Judge
        return self._heuristic_judge(code, idea)

    def _extract_code_block(self, text: str) -> str:
        matches = re.findall(r'```python(.*?)```', text, re.DOTALL)
        if matches:
            return matches[0].strip()
        matches_generic = re.findall(r'```(.*?)```', text, re.DOTALL)
        if matches_generic:
            return matches_generic[0].strip()
        return text.strip()

    def _instantiate_factor_code(self, code_str: str, class_name: str) -> Tuple[Optional[Factor], Optional[str]]:
        """Safely compiles and instantiates factor code in an execution namespace."""
        import builtins
        safe_globals = {
            "Factor": Factor,
            "np": np,
            "pd": pd,
            "scipy": __import__("scipy"),
            "__builtins__": builtins.__dict__
        }
        try:
            compiled = compile(code_str, "<quant_agent_sandbox>", "exec")
            local_vars = {}
            exec(compiled, safe_globals, local_vars)
            
            # Find class
            factor_cls = local_vars.get(class_name)
            if not factor_cls:
                for v in local_vars.values():
                    if isinstance(v, type) and issubclass(v, Factor) and v is not Factor:
                        factor_cls = v
                        break
            
            if factor_cls is None:
                return None, f"No Factor subclass named '{class_name}' found."

            instance = factor_cls()
            return instance, None
        except Exception as e:
            return None, str(e)

    def _heuristic_writer(self, idea: str, prev_code: str, feedback: str) -> Tuple[str, str]:
        """Provides expert procedural implementations that evolve across iterations."""
        idea_lower = idea.lower()

        # If previous code exists and was criticized, refine it
        if "binary" in feedback.lower() or "continuous" in feedback.lower() or "nuance" in feedback.lower():
            name = "ContinuousNormalizedAlpha"
            code = '''
class ContinuousNormalizedAlpha(Factor):
    def __init__(self):
        self.name = "ContinuousNormalizedAlpha"
        self.window_length = 20
        self.inputs = ["close", "high", "low", "volume"]

    def calc(self, data):
        close = data["close"]
        high = data["high"]
        low = data["low"]
        volume = data["volume"]
        
        # Volatility normalization via True Range
        tr = np.maximum(high - low, np.maximum((high - close.shift(1)).abs(), (low - close.shift(1)).abs()))
        atr = tr.rolling(self.window_length).mean().replace(0, np.nan)
        
        # Continuous momentum normalized by volatility
        mom = (close - close.shift(5)) / atr
        
        # Volume acceleration multiplier
        vol_ratio = volume / volume.rolling(self.window_length).mean().replace(0, np.nan)
        vol_scaled = np.clip(vol_ratio, 0.5, 2.5)
        
        raw_signal = mom * vol_scaled
        # Bounded between -1.0 and 1.0 with smooth activation
        signal = np.tanh(raw_signal / 2.0).fillna(0.0)
        return signal
'''.strip()
            return name, code

        if "breakout" in idea_lower or "volatility" in idea_lower:
            name = "AdaptiveVolatilityBreakout"
            code = '''
class AdaptiveVolatilityBreakout(Factor):
    def __init__(self):
        self.name = "AdaptiveVolatilityBreakout"
        self.window_length = 20
        self.inputs = ["close", "high", "low", "volume"]

    def calc(self, data):
        close = data["close"]
        high = data["high"]
        low = data["low"]
        
        # Donchian channel
        upper = high.shift(1).rolling(self.window_length).max()
        lower = low.shift(1).rolling(self.window_length).min()
        channel_width = (upper - lower).replace(0, np.nan)
        
        # Distance from center normalized to [-1, 1]
        mid = (upper + lower) / 2.0
        signal = ((close - mid) / (channel_width / 2.0)).clip(-1.0, 1.0).fillna(0.0)
        return signal
'''.strip()
            return name, code

        if "mean reversion" in idea_lower or "reversal" in idea_lower or "rsi" in idea_lower:
            name = "NormalizedMeanReversionAlpha"
            code = '''
class NormalizedMeanReversionAlpha(Factor):
    def __init__(self):
        self.name = "NormalizedMeanReversionAlpha"
        self.window_length = 15
        self.inputs = ["close"]

    def calc(self, data):
        close = data["close"]
        sma = close.rolling(self.window_length).mean()
        std = close.rolling(self.window_length).std().replace(0, np.nan)
        
        # Cross-sectional Z-score inversion
        z_score = (close - sma) / std
        signal = -np.tanh(z_score).fillna(0.0)  # Negative sign for mean reversion
        return signal
'''.strip()
            return name, code

        # Default Multi-Factor Trend
        name = "MultiPeriodTrendAlpha"
        code = '''
class MultiPeriodTrendAlpha(Factor):
    def __init__(self):
        self.name = "MultiPeriodTrendAlpha"
        self.window_length = 25
        self.inputs = ["close", "volume"]

    def calc(self, data):
        close = data["close"]
        ret5 = close.pct_change(5)
        ret20 = close.pct_change(20)
        composite = 0.6 * ret5 + 0.4 * ret20
        vol20 = close.pct_change().rolling(20).std().replace(0, np.nan)
        sharpe_signal = composite / vol20
        return np.tanh(sharpe_signal).fillna(0.0)
'''.strip()
        return name, code

    def _heuristic_judge(self, code: str, idea: str) -> Tuple[float, str, List[str]]:
        """Rigorous automated quantitative critique."""
        score = 8.5
        critiques = []
        advice = []

        # Check 1: Vectorization
        if "for stock in" in code or ".iterrows()" in code:
            score -= 2.5
            critiques.append("Unvectorized stock-by-stock loop detected. Slow execution.")
            advice.append("Vectorize calculations across entire DataFrame matrices.")

        # Check 2: Lookahead bias
        if "shift(-" in code:
            score -= 4.0
            critiques.append("CRITICAL: Negative shift shift(-1) detected! Lookahead bias.")
            advice.append("Remove future shifts; only use historical bars shift(1) or rolling.")

        # Check 3: Flat binary signal
        if "== 1" in code or "* 1.0" in code and "tanh" not in code and "clip" not in code:
            score -= 2.0
            critiques.append("Signal appears binary or non-continuous. Models require ranked nuance.")
            advice.append("Use continuous scaling like np.tanh, z-score, or distance normalization.")

        # Check 4: NaN protection
        if "fillna" not in code:
            score -= 1.0
            critiques.append("Missing .fillna() handling; may cause propagation of NaNs.")
            advice.append("Append .fillna(0.0) to output.")

        # Check 5: Inheritance
        if "class " not in code or "(Factor)" not in code:
            score -= 4.0
            critiques.append("Class does not properly inherit from `Factor`.")
            advice.append("Subclass Factor and implement calc(self, data).")

        score = max(1.0, min(10.0, score))
        summary = " | ".join(critiques) if critiques else "Code complies with vectorized, continuous, leakage-free quantitative standards."
        return score, summary, advice

    # -------------------------------------------------------------------------
    # Downstream XGBoost Multi-Alpha Ensemble (Replicating Paper Figure 5)
    # -------------------------------------------------------------------------

    def train_downstream_xgboost(
        self,
        market_data: Dict[str, pd.DataFrame]
    ) -> Dict[str, Any]:
        """
        Gathers all successful alphas from KB, evaluates features X across time/stocks,
        and trains an XGBoost model predicting forward returns.
        Demonstrates reduction in prediction MSE as KB expands (Figure 5).
        """
        success_alphas = self.kb.get_success_alphas()
        if len(success_alphas) < 1:
            return {"error": "Need at least 1 successful alpha in KB to train ensemble"}

        tickers = list(market_data.keys())
        close_matrix = pd.DataFrame({
            t: market_data[t]['close'] for t in tickers
        }).ffill().bfill()
        target_returns = close_matrix.pct_change(1).shift(-1)

        # Build feature dataset
        feature_matrices = {}
        for a in success_alphas:
            factor_inst, _ = self._instantiate_factor_code(a.code, a.name)
            if factor_inst:
                ok, _, sig_df = self.backtester.execute_and_evaluate(factor_inst, market_data)
                if ok and sig_df is not None:
                    feature_matrices[a.name] = sig_df

        if not feature_matrices:
            return {"error": "Failed to compute signal matrices for ensemble features"}

        # Stack into tabular (date x ticker) format
        rows = []
        for date in target_returns.index[:-1]:
            for ticker in tickers:
                y_val = target_returns.loc[date, ticker]
                if not np.isnan(y_val):
                    row_data = {"date": date, "ticker": ticker, "target": y_val}
                    for f_name, f_mat in feature_matrices.items():
                        if date in f_mat.index and ticker in f_mat.columns:
                            row_data[f_name] = f_mat.loc[date, ticker]
                        else:
                            row_data[f_name] = 0.0
                    rows.append(row_data)

        df_all = pd.DataFrame(rows).dropna()
        if len(df_all) < 50:
            return {"error": "Insufficient tabular observations for training"}

        feature_cols = list(feature_matrices.keys())
        X = df_all[feature_cols].values
        y = df_all["target"].values

        # K-Fold Cross Validation
        kf = KFold(n_splits=3, shuffle=False)
        mse_list = []
        model = xgb.XGBRegressor(n_estimators=50, max_depth=3, learning_rate=0.05, random_state=42)

        for train_idx, val_idx in kf.split(X):
            X_tr, y_tr = X[train_idx], y[train_idx]
            X_va, y_va = X[val_idx], y[val_idx]
            model.fit(X_tr, y_tr)
            preds = model.predict(X_va)
            mse_list.append(mean_squared_error(y_va, preds))

        mean_mse = float(np.mean(mse_list))
        importances = {col: round(float(imp), 4) for col, imp in zip(feature_cols, model.feature_importances_)}

        return {
            "status": "SUCCESS",
            "alphas_combined_count": len(feature_cols),
            "alpha_names": feature_cols,
            "mean_squared_error": round(mean_mse, 6),
            "feature_importances": importances,
            "observations_count": len(df_all)
        }