afres_v2.py

"""
AFRES v2: Agentic Factor Revision and Evaluation System

A faithful adaptation of APRES's rubric-discovery pipeline to quantitative
factor generation.  Key design choices (per user request):

  1. LLM agent generates factor expressions from a discovered rubric.
  2. A regression model (factor-value → future-return) provides the fitness
     signal.  Fitness = –MAE (or IC) on a held-out test set.
  3. NO LLM-as-judge.  Factor quality is determined entirely by the data.
  4. NO QD / MAP-Elites.  The search budget is spent on rubric discovery.
  5. The rubric-discovery loop mirrors APRES Section 3.1:
        Propose → Generate Factors → Evaluate via Regression → Select & Refine
     with MultiAIDE-style tree search (branching, debug-and-retry).

Architecture
────────────
┌─────────────┐     ┌─────────────────┐     ┌──────────────────┐
│  Rubric     │────→│  LLM Factor     │────→│  Regression      │
│  Proposer   │     │  Generator      │     │  Evaluator       │
│  (LLM)      │     │  (prompt → expr)│     │  (sklearn)       │
└─────────────┘     └─────────────────┘     └──────────────────┘
       ↑                                              │
       │           Select & Refine (MultiAIDE tree)    │
       └──────────────────────────────────────────────┘

References
  • APRES   – arXiv:2603.03142 (Sec. 3.1 rubric search)
  • MultiAIDE – Zhao et al. 2025 (tree-search scaffold)
"""

import json
import re
import time
import copy
import warnings
from dataclasses import dataclass, field
from typing import List, Dict, Optional, Tuple, Callable
from enum import Enum
from collections import defaultdict

import numpy as np
import pandas as pd
from sklearn.linear_model import Ridge
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_absolute_error, r2_score

warnings.filterwarnings("ignore")

# ═════════════════════════════════════════════════════════════════════════════
# 1.  DATA STRUCTURES
# ═════════════════════════════════════════════════════════════════════════════

class SignalType(Enum):
    PRICE_BASED = "price_based"
    VOLUME_BASED = "volume_based"
    FUNDAMENTAL = "fundamental"
    TECHNICAL = "technical"
    CROSS_SECTIONAL = "cross_sectional"
    TIME_SERIES = "time_series"

@dataclass
class RubricItem:
    """One actionable design principle for factor generation."""
    id: str
    description: str
    # Actionable constraints passed to the generator
    feature_focus: List[str] = field(default_factory=list)
    preferred_ops: List[str] = field(default_factory=list)
    time_horizon_hint: str = "any"          # short / medium / long / any
    complexity_hint: str = "any"            # low / medium / high / any
    weight: float = 1.0

    def to_dict(self) -> Dict:
        return {
            "id": self.id, "description": self.description,
            "feature_focus": self.feature_focus,
            "preferred_ops": self.preferred_ops,
            "time_horizon_hint": self.time_horizon_hint,
            "complexity_hint": self.complexity_hint,
            "weight": self.weight,
        }

@dataclass
class FactorRubric:
    items: List[RubricItem]
    def to_dict(self):
        return {"items": [i.to_dict() for i in self.items]}

@dataclass
class Factor:
    id: str
    expression: str
    rubric_id: str          # which rubric variant produced this factor
    ic: float = 0.0
    mae: float = 0.0        # regression MAE (lower = better)
    r2: float = 0.0         # regression R²
    sharpe: float = 0.0
    returns: float = 0.0
    generation: int = 0
    valid: bool = True      # False if expression failed to evaluate

    def to_dict(self):
        return {
            "id": self.id, "expression": self.expression,
            "rubric_id": self.rubric_id, "ic": self.ic,
            "mae": self.mae, "r2": self.r2, "sharpe": self.sharpe,
            "returns": self.returns, "generation": self.generation,
            "valid": self.valid,
        }

# ═════════════════════════════════════════════════════════════════════════════
# 2.  MARKET DATA  (panel: days × stocks, pre-computed features)
# ═════════════════════════════════════════════════════════════════════════════

class MarketData:
    """
    Synthetic market panel with baked-in predictive structure so that
    well-designed factors can genuinely outperform random ones.
    """

    def __init__(self, n_stocks: int = 80, n_days: int = 600, seed: int = 42):
        rng = np.random.RandomState(seed)
        self.n_stocks = n_stocks
        self.n_days = n_days
        self.dates = pd.date_range("2020-01-01", periods=n_days, freq="B")
        self.symbols = [f"S{i:03d}" for i in range(n_stocks)]

        # ── raw price / volume panels ──
        self.close = np.zeros((n_days, n_stocks))
        self.open_ = np.zeros((n_days, n_stocks))
        self.high = np.zeros((n_days, n_stocks))
        self.low = np.zeros((n_days, n_stocks))
        self.volume = np.zeros((n_days, n_stocks))

        for i in range(n_stocks):
            # base returns
            ret = rng.normal(0.0002, 0.018, n_days)

            # ── predictable structure (weak but real) ──
            # momentum: past 5-day ret → next day (+)
            mom = np.zeros(n_days)
            mom[5:] = 0.25 * ret[:-5]
            # mean-reversion: deviation from 10-day mean → next day (–)
            rev = np.zeros(n_days)
            for t in range(10, n_days):
                rev[t] = -0.15 * (ret[t - 1] - ret[t - 10:t].mean())
            # volume-return interaction
            vol_signal = np.zeros(n_days)
            vol_signal[5:] = 0.10 * np.abs(ret[:-5]) * rng.lognormal(0, 0.3, n_days - 5)

            ret[10:] += mom[10:] + rev[10:] + vol_signal[10:]

            prices = 100 * np.exp(np.cumsum(ret))
            self.close[:, i] = prices
            self.open_[:, i] = prices * (1 + rng.normal(0, 0.001, n_days))
            self.high[:, i] = prices * (1 + np.abs(rng.normal(0, 0.008, n_days)))
            self.low[:, i] = prices * (1 - np.abs(rng.normal(0, 0.008, n_days)))
            self.volume[:, i] = rng.lognormal(15, 0.3, n_days)

        # ── pre-compute feature panels ──
        self.ret_1d = np.diff(self.close, axis=0, prepend=self.close[:1]) / (self.close + 1e-10)
        self.ret_5d = np.zeros_like(self.close)
        self.ret_5d[5:] = (self.close[5:] - self.close[:-5]) / (self.close[:-5] + 1e-10)
        self.ret_20d = np.zeros_like(self.close)
        self.ret_20d[20:] = (self.close[20:] - self.close[:-20]) / (self.close[:-20] + 1e-10)

        df_close = pd.DataFrame(self.close)
        self.sma_5 = df_close.rolling(5, min_periods=1).mean().values
        self.sma_10 = df_close.rolling(10, min_periods=1).mean().values
        self.sma_20 = df_close.rolling(20, min_periods=1).mean().values
        self.vol_20d = pd.DataFrame(self.ret_1d).rolling(20, min_periods=1).std().values

        df_volume = pd.DataFrame(self.volume)
        self.vol_sma_20 = df_volume.rolling(20, min_periods=1).mean().values
        self.high_20d = pd.DataFrame(self.high).rolling(20, min_periods=1).max().values
        self.low_20d = pd.DataFrame(self.low).rolling(20, min_periods=1).min().values
        self.vwap = self.close * (1 + rng.normal(0, 0.0003, (n_days, n_stocks)))

        # ── target: next-day return ──
        self.future_ret = np.zeros_like(self.close)
        self.future_ret[:-1] = np.diff(self.close, axis=0) / (self.close[:-1] + 1e-10)

        # ── train / test split ──
        self.train_idx = np.arange(30, 450)   # skip first 30 for feature warmup
        self.test_idx = np.arange(450, min(580, n_days - 1))

    def eval_expr(self, expr: str) -> Optional[np.ndarray]:
        """Evaluate a factor expression → (days × stocks) array."""
        ns = {
            'close': self.close, 'open': self.open_, 'high': self.high,
            'low': self.low, 'volume': self.volume, 'vwap': self.vwap,
            'returns_1d': self.ret_1d, 'returns_5d': self.ret_5d,
            'returns_20d': self.ret_20d, 'sma_5': self.sma_5,
            'sma_10': self.sma_10, 'sma_20': self.sma_20,
            'volatility_20d': self.vol_20d, 'volume_sma_20': self.vol_sma_20,
            'high_20d': self.high_20d, 'low_20d': self.low_20d,
            'np': np, 'abs': np.abs, 'log': np.log, 'sqrt': np.sqrt,
            'sign': np.sign,
            'rank': lambda x: self._rank(x),
            'ts_mean': lambda x, w: self._ts(x, w, 'mean'),
            'ts_std': lambda x, w: self._ts(x, w, 'std'),
            'ts_max': lambda x, w: self._ts(x, w, 'max'),
            'ts_min': lambda x, w: self._ts(x, w, 'min'),
            'ts_zscore': lambda x, w: (x - self._ts(x, w, 'mean')) /
                                        (self._ts(x, w, 'std') + 1e-10),
            'ts_delta': lambda x, w: self._delta(x, w),
            'ts_corr': lambda x, y, w: self._corr(x, y, w),
            'ts_cov': lambda x, y, w: self._cov(x, y, w),
            'ts_rank': lambda x, w: self._tsrank(x, w),
        }
        try:
            result = eval(expr, {"__builtins__": {}}, ns)
            if isinstance(result, np.ndarray) and result.shape == (self.n_days, self.n_stocks):
                return result
            return None
        except Exception:
            return None

    # ---- helpers ----
    def _rank(self, x):
        r = np.zeros_like(x)
        for t in range(x.shape[0]):
            valid = np.isfinite(x[t])
            if valid.sum() > 0:
                r[t, valid] = pd.Series(x[t, valid]).rank(pct=True).values
        return r

    def _ts(self, x, w, method):
        df = pd.DataFrame(x)
        if method == 'mean':   return df.rolling(w, min_periods=1).mean().values
        if method == 'std':    return df.rolling(w, min_periods=1).std().values
        if method == 'max':    return df.rolling(w, min_periods=1).max().values
        if method == 'min':    return df.rolling(w, min_periods=1).min().values
        return x

    def _delta(self, x, w):
        out = np.zeros_like(x)
        out[w:] = x[w:] - x[:-w]
        return out

    def _corr(self, x, y, w):
        r = np.zeros_like(x)
        for t in range(x.shape[0]):
            a = x[max(0, t - w + 1):t + 1].flatten()
            b = y[max(0, t - w + 1):t + 1].flatten()
            if len(a) > 1 and np.std(a) > 0 and np.std(b) > 0:
                r[t] = np.corrcoef(a, b)[0, 1]
        return r

    def _cov(self, x, y, w):
        c = np.zeros_like(x)
        for t in range(x.shape[0]):
            a = x[max(0, t - w + 1):t + 1].flatten()
            b = y[max(0, t - w + 1):t + 1].flatten()
            c[t] = np.cov(a, b)[0, 1] if len(a) > 1 else 0.0
        return c

    def _tsrank(self, x, w):
        r = np.zeros_like(x)
        for t in range(x.shape[0]):
            vals = x[max(0, t - w + 1):t + 1].flatten()
            if len(vals) >= w:
                r[t] = pd.Series(vals).rank(pct=True).values[-1]
            else:
                r[t] = np.nan
        return r

# ═════════════════════════════════════════════════════════════════════════════
# 3.  REGRESSION EVALUATOR  (factor value → future return)
# ═════════════════════════════════════════════════════════════════════════════

class RegressionEvaluator:
    """
    Trains a regression model: factor_value → next-day return.
    Fitness is reported as:
        • MAE  (Mean Absolute Error)  – primary metric, lower = better
        • IC   (Spearman rank correlation) – cross-sectional predictive power
        • R²   (coefficient of determination)
        • Sharpe & returns from a simple long-top-quintile strategy
    """

    def __init__(self, data: MarketData, model_type: str = "ridge"):
        self.data = data
        self.model_type = model_type

    def evaluate(self, factor: Factor) -> Dict[str, float]:
        vals = self.data.eval_expr(factor.expression)
        if vals is None:
            factor.valid = False
            return {"mae": 1e6, "ic": 0.0, "r2": -1.0, "sharpe": 0.0, "returns": 0.0}

        # ── flatten panel to (sample, feature) for regression ──
        X_train, y_train = self._flatten(vals, self.data.train_idx)
        X_test,  y_test  = self._flatten(vals, self.data.test_idx)

        if len(X_train) < 100 or len(X_test) < 50:
            factor.valid = False
            return {"mae": 1e6, "ic": 0.0, "r2": -1.0, "sharpe": 0.0, "returns": 0.0}

        # ── train regression ──
        model = Ridge(alpha=1.0) if self.model_type == "ridge" else \
                RandomForestRegressor(n_estimators=50, max_depth=6, random_state=42, n_jobs=-1)
        model.fit(X_train, y_train)
        pred_test = model.predict(X_test)

        mae = mean_absolute_error(y_test, pred_test)
        r2 = r2_score(y_test, pred_test)

        # ── IC (cross-sectional rank correlation per day) ──
        ics = []
        for t in self.data.test_idx:
            f_t = vals[t]
            r_t = self.data.future_ret[t]
            valid = np.isfinite(f_t) & np.isfinite(r_t)
            if valid.sum() >= 10:
                ic = np.corrcoef(pd.Series(f_t[valid]).rank().values,
                                  pd.Series(r_t[valid]).rank().values)[0, 1]
                if np.isfinite(ic):
                    ics.append(ic)
        ic = float(np.mean(ics)) if ics else 0.0

        # ── simple portfolio Sharpe ──
        port_rets = []
        for t in self.data.test_idx:
            f_t = vals[t]
            r_t = self.data.future_ret[t]
            valid = np.isfinite(f_t) & np.isfinite(r_t)
            if valid.sum() >= 10:
                q80 = np.percentile(f_t[valid], 80)
                mask = (f_t >= q80) & valid
                if mask.sum() > 0:
                    port_rets.append(float(np.mean(r_t[mask])))

        if len(port_rets) > 2:
            rets_arr = np.array(port_rets)
            ann_ret = float(np.mean(rets_arr) * 252)
            sharpe = float((np.mean(rets_arr) / (np.std(rets_arr) + 1e-10)) * np.sqrt(252))
        else:
            ann_ret = 0.0
            sharpe = 0.0

        factor.mae = mae
        factor.ic = ic
        factor.r2 = r2
        factor.sharpe = sharpe
        factor.returns = ann_ret
        factor.valid = True
        return {"mae": mae, "ic": ic, "r2": r2, "sharpe": sharpe, "returns": ann_ret}

    def _flatten(self, vals: np.ndarray, idx: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """Flatten panel slices to (n_samples, 1) feature matrix + target vector."""
        X = vals[idx].flatten().reshape(-1, 1)
        y = self.data.future_ret[idx].flatten()
        mask = np.isfinite(X[:, 0]) & np.isfinite(y)
        return X[mask], y[mask]

# ═════════════════════════════════════════════════════════════════════════════
# 4.  LLM FACTOR GENERATOR  (interface + simulated implementation)
# ═════════════════════════════════════════════════════════════════════════════

class LLMFactorGenerator:
    """
    Generates factor expressions from a rubric.
    In production this would call an LLM API (GPT-4, Claude, etc.).
    The prototype uses a template-based simulator that respects the
    actionable constraints in each RubricItem.
    """

    # operator library
    OPS = ["ts_mean", "ts_std", "ts_max", "ts_min", "ts_zscore",
           "ts_delta", "ts_corr", "ts_cov", "ts_rank", "rank",
           "abs", "sign", "log", "sqrt"]

    FEATURES = ["close", "open", "high", "low", "volume", "vwap",
                "returns_1d", "returns_5d", "returns_20d",
                "sma_5", "sma_10", "sma_20",
                "volatility_20d", "volume_sma_20",
                "high_20d", "low_20d"]

    WINDOWS = [3, 5, 10, 20]

    def __init__(self, seed: int = 99):
        self.rng = np.random.RandomState(seed)
        self._counter = 0

    def generate(self, rubric: FactorRubric, n: int = 5) -> List[Factor]:
        """Generate n factor expressions constrained by the rubric."""
        # Aggregate constraints from rubric items
        features = self._extract_features(rubric)
        ops = self._extract_ops(rubric)
        complexity_target = self._extract_complexity(rubric)

        factors = []
        for _ in range(n):
            self._counter += 1
            expr = self._build_expression(features, ops, complexity_target)
            factors.append(Factor(
                id=f"f_{self._counter}",
                expression=expr,
                rubric_id=id(rubric),  # temporary ID, will be overwritten
                generation=0,
            ))
        return factors

    # ── constraint extraction ──
    def _extract_features(self, rubric: FactorRubric) -> List[str]:
        feats = []
        for item in rubric.items:
            feats.extend(item.feature_focus)
        return list(set(feats)) if feats else self.FEATURES

    def _extract_ops(self, rubric: FactorRubric) -> List[str]:
        ops = []
        for item in rubric.items:
            ops.extend(item.preferred_ops)
        return list(set(ops)) if ops else self.OPS

    def _extract_complexity(self, rubric: FactorRubric) -> int:
        hints = [item.complexity_hint for item in rubric.items]
        if "low" in hints:
            return 2
        if "high" in hints:
            return 5
        return 3  # default medium

    # ── expression builder ──
    def _build_expression(self, features, ops, complexity_target: int) -> str:
        """Build a random expression of roughly target complexity."""
        n_terms = max(1, complexity_target - 1 + self.rng.randint(-1, 2))
        terms = []
        for _ in range(n_terms):
            terms.append(self._random_term(features, ops))
        if len(terms) == 1:
            return terms[0]
        # Combine with binary ops
        expr = terms[0]
        for t in terms[1:]:
            op = self.rng.choice([" + ", " - ", " * "])
            expr = f"({expr}){op}({t})"
        return expr

    def _random_term(self, features, ops) -> str:
        feat = self.rng.choice(features)
        # Decide: raw feature, windowed operator, or cross-feature interaction
        choice = self.rng.choice(["raw", "op", "interaction"], p=[0.25, 0.50, 0.25])
        if choice == "raw":
            return feat
        if choice == "op":
            op = self.rng.choice([o for o in ops if o.startswith("ts_") or o in ("rank", "abs", "sign", "log", "sqrt")])
            if op.startswith("ts_") and op not in ("ts_corr", "ts_cov"):
                w = self.rng.choice(self.WINDOWS)
                return f"{op}({feat}, {w})"
            if op in ("rank", "abs", "sign", "log", "sqrt"):
                return f"{op}({feat})"
            # fallback
            return feat
        # interaction
        feat2 = self.rng.choice([f for f in features if f != feat])
        op = self.rng.choice([" * ", " + ", " - "])
        return f"{feat}{op}{feat2}"

# ═════════════════════════════════════════════════════════════════════════════
# 5.  RUBRIC PROPOSER  (MultiAIDE-style tree search)
# ═════════════════════════════════════════════════════════════════════════════

class RubricProposer:
    """
    Implements the APRES rubric-discovery loop:
        Propose → Generate Factors → Evaluate via Regression → Select & Refine
    With MultiAIDE-style tree search:
        • N0 initial rubric branches
        • Each step: branch N new variants from best or buggy (prob p_debug)
        • Max debug depth D_max per node
    """

    # Seed rubrics to start the search
    SEED_RUBRICS = [
        FactorRubric([
            RubricItem("price_momentum", "Prefer price-based momentum signals",
                       feature_focus=["close", "returns_1d", "returns_5d"],
                       preferred_ops=["ts_mean", "ts_delta", "rank"],
                       time_horizon_hint="short", complexity_hint="low"),
        ]),
        FactorRubric([
            RubricItem("volume_confirm", "Use volume to confirm price signals",
                       feature_focus=["volume", "volume_sma_20", "returns_1d"],
                       preferred_ops=["ts_corr", "rank", "ts_mean"],
                       time_horizon_hint="medium", complexity_hint="medium"),
        ]),
        FactorRubric([
            RubricItem("mean_reversion", "Capture mean-reversion patterns",
                       feature_focus=["close", "sma_20", "sma_10"],
                       preferred_ops=["ts_zscore", "ts_delta", "abs"],
                       time_horizon_hint="short", complexity_hint="medium"),
        ]),
    ]

    def __init__(self,
                 N0: int = 3,
                 N: int = 3,
                 p_debug: float = 0.3,
                 D_max: int = 5,
                 seed: int = 77):
        self.N0 = N0
        self.N = N
        self.p_debug = p_debug
        self.D_max = D_max
        self.rng = np.random.RandomState(seed)
        self._debug_counts = defaultdict(int)   # node_id → debug attempts

    def propose_initial(self) -> List[FactorRubric]:
        """N0 diverse seed rubrics."""
        rubrics = [copy.deepcopy(r) for r in self.SEED_RUBRICS[:self.N0]]
        # Pad with randomised variants if fewer seeds than N0
        while len(rubrics) < self.N0:
            base = copy.deepcopy(self.rng.choice(self.SEED_RUBRICS))
            rubrics.append(self._mutate_rubric(base))
        return rubrics

    def propose_from_parent(self, parent: FactorRubric, is_buggy: bool = False) -> List[FactorRubric]:
        """Branch N new rubric variants from a parent."""
        children = []
        for _ in range(self.N):
            child = self._mutate_rubric(copy.deepcopy(parent))
            if is_buggy:
                # More aggressive mutation for debug branches
                child = self._mutate_rubric(child)
            children.append(child)
        return children

    # ── mutation operators ──
    def _mutate_rubric(self, rubric: FactorRubric) -> FactorRubric:
        """Apply one structural mutation to a rubric."""
        items = rubric.items
        mutation = self.rng.choice(["add", "remove", "replace", "tweak"])

        if mutation == "add" or len(items) == 0:
            items.append(self._random_item())
        elif mutation == "remove" and len(items) > 1:
            items.pop(self.rng.randint(0, len(items)))
        elif mutation == "replace" and len(items) > 0:
            idx = self.rng.randint(0, len(items))
            items[idx] = self._random_item()
        elif mutation == "tweak" and len(items) > 0:
            idx = self.rng.randint(0, len(items))
            items[idx] = self._tweak_item(items[idx])

        return FactorRubric(items)

    def _random_item(self) -> RubricItem:
        templates = [
            ("momentum", "Focus on momentum signals",
             ["close", "returns_5d", "returns_20d"], ["ts_mean", "ts_delta", "rank"]),
            ("volatility", "Exploit volatility patterns",
             ["volatility_20d", "returns_1d", "close"], ["ts_zscore", "ts_std", "abs"]),
            ("volume_price", "Volume-price interaction",
             ["volume", "close", "returns_1d"], ["ts_corr", "rank", "ts_mean"]),
            ("cross_sectional", "Cross-sectional ranking",
             ["close", "volume", "returns_1d"], ["rank", "ts_zscore", "ts_rank"]),
            ("mean_reversion", "Mean-reversion",
             ["close", "sma_10", "sma_20"], ["ts_zscore", "ts_delta", "sign"]),
            ("breakout", "Breakout patterns",
             ["high_20d", "low_20d", "close"], ["ts_max", "ts_min", "ts_delta"]),
            ("vwap", "VWAP deviation",
             ["vwap", "close", "volume"], ["ts_zscore", "ts_mean", "ts_corr"]),
        ]
        t = templates[self.rng.randint(0, len(templates))]
        return RubricItem(
            id=t[0], description=t[1],
            feature_focus=t[2], preferred_ops=t[3],
            time_horizon_hint=self.rng.choice(["short", "medium", "long"]),
            complexity_hint=self.rng.choice(["low", "medium", "high"]),
            weight=1.0,
        )

    def _tweak_item(self, item: RubricItem) -> RubricItem:
        """Small perturbation of one rubric item."""
        tweak = self.rng.choice(["feature", "op", "horizon", "complexity"])
        if tweak == "feature" and item.feature_focus:
            all_feats = ["close", "volume", "returns_1d", "returns_5d", "vwap",
                         "sma_10", "sma_20", "volatility_20d", "high_20d", "low_20d"]
            item.feature_focus = list(set(
                item.feature_focus + [self.rng.choice(all_feats)]
            ))[:3]
        elif tweak == "op" and item.preferred_ops:
            all_ops = ["ts_mean", "ts_std", "ts_zscore", "ts_delta", "rank",
                       "ts_corr", "ts_rank", "abs", "sign", "log"]
            item.preferred_ops = list(set(
                item.preferred_ops + [self.rng.choice(all_ops)]
            ))[:3]
        elif tweak == "horizon":
            item.time_horizon_hint = self.rng.choice(["short", "medium", "long"])
        elif tweak == "complexity":
            item.complexity_hint = self.rng.choice(["low", "medium", "high"])
        return item

# ═════════════════════════════════════════════════════════════════════════════
# 6.  MAIN AFRES SYSTEM  (Rubric Discovery Loop)
# ═════════════════════════════════════════════════════════════════════════════

class AFRES:
    """
    Agentic Factor Revision and Evaluation System.

    Phase 1 – Rubric Discovery (faithful to APRES §3.1):
        1. Propose:      RubricProposer creates rubric variants
        2. Generate:     LLMFactorGenerator produces factors per rubric
        3. Evaluate:     RegressionEvaluator trains model, reports MAE
        4. Select&Refine: MultiAIDE tree search keeps / branches best rubrics

    Phase 2 – Best-Rubric Factor Generation:
        Use the discovered rubric to generate a final pool of factors.
    """

    def __init__(self,
                 data: MarketData,
                 generator: LLMFactorGenerator,
                 evaluator: RegressionEvaluator,
                 proposer: RubricProposer,
                 factors_per_rubric: int = 4):
        self.data = data
        self.generator = generator
        self.evaluator = evaluator
        self.proposer = proposer
        self.k = factors_per_rubric

        # Search state
        self.best_rubric: Optional[FactorRubric] = None
        self.best_fitness: float = -1e9
        self.all_factors: List[Factor] = []
        self.history: List[Dict] = []

    def discover_rubric(self, max_iterations: int = 20) -> Tuple[FactorRubric, List[Factor]]:
        """
        Run the APRES-style rubric-discovery loop.

        Returns the best discovered rubric and the best factors found under it.
        """
        print("=" * 70)
        print("  PHASE 1: RUBRIC DISCOVERY (APRES-style MultiAIDE search)")
        print("=" * 70)

        # ── initial population ──
        population = self.proposer.propose_initial()
        scores = []
        for rubric in population:
            fitness, factors = self._evaluate_rubric(rubric)
            scores.append((fitness, rubric, factors))
            self.history.append({
                "iteration": 0, "rubric_id": id(rubric),
                "fitness": fitness, "n_factors": len(factors),
                "mean_ic": np.mean([f.ic for f in factors]) if factors else 0.0,
                "mean_mae": np.mean([f.mae for f in factors]) if factors else 1e6,
            })

        # Track best
        best = max(scores, key=lambda x: x[0])
        self.best_fitness, self.best_rubric, best_factors = best
        print(f"\nInitial best fitness = {self.best_fitness:.4f}  "
              f"(IC={np.mean([f.ic for f in best_factors]):.4f}, "
              f"MAE={np.mean([f.mae for f in best_factors]):.4f})")

        # ── iterative tree search ──
        for it in range(1, max_iterations + 1):
            t0 = time.time()

            # Pick parent: best with prob (1-p_debug), else random buggy
            if self.proposer.rng.random() > self.proposer.p_debug:
                parent = self.best_rubric
                is_buggy = False
            else:
                # Pick a random previously-evaluated rubric (could be bad)
                parent = self.proposer.rng.choice([s[1] for s in scores])
                is_buggy = True

            children = self.proposer.propose_from_parent(parent, is_buggy)

            child_scores = []
            for child in children:
                fitness, factors = self._evaluate_rubric(child)
                child_scores.append((fitness, child, factors))
                self.history.append({
                    "iteration": it, "rubric_id": id(child),
                    "fitness": fitness, "n_factors": len(factors),
                    "mean_ic": np.mean([f.ic for f in factors]) if factors else 0.0,
                    "mean_mae": np.mean([f.mae for f in factors]) if factors else 1e6,
                    "buggy_branch": is_buggy,
                })

            # Update global best
            local_best = max(child_scores, key=lambda x: x[0])
            if local_best[0] > self.best_fitness:
                self.best_fitness = local_best[0]
                self.best_rubric = local_best[1]
                best_factors = local_best[2]
                print(f"  Iter {it:02d}: NEW BEST fitness={self.best_fitness:.4f}  "
                      f"IC={np.mean([f.ic for f in best_factors]):.4f}  "
                      f"MAE={np.mean([f.mae for f in best_factors]):.4f}  "
                      f"({time.time()-t0:.1f}s)")
            else:
                print(f"  Iter {it:02d}: no improvement  "
                      f"best={self.best_fitness:.4f}  "
                      f"({time.time()-t0:.1f}s)")

            scores.extend(child_scores)

        print(f"\n{'='*70}")
        print(f"  RUBRIC DISCOVERY COMPLETE  –  best fitness = {self.best_fitness:.4f}")
        print(f"{'='*70}")
        self._print_rubric(self.best_rubric)
        return self.best_rubric, best_factors

    def _evaluate_rubric(self, rubric: FactorRubric) -> Tuple[float, List[Factor]]:
        """
        Evaluate a rubric:
            1. Generate k factor expressions
            2. Evaluate each on market data via regression
            3. Aggregate fitness (mean of valid factors)
        """
        rubric_id = f"rubric_{id(rubric)}"
        factors = self.generator.generate(rubric, n=self.k)
        for f in factors:
            f.rubric_id = rubric_id

        valid_factors = []
        for f in factors:
            self.evaluator.evaluate(f)
            self.all_factors.append(f)
            if f.valid:
                valid_factors.append(f)

        if not valid_factors:
            return -1e6, factors  # severe penalty for rubric that generates no valid factors

        # Fitness = mean IC (higher = better)
        # Could also use -MAE or a blend; IC is cleaner for ranking
        fitness = float(np.mean([f.ic for f in valid_factors]))
        return fitness, factors

    def generate_final_pool(self, n: int = 20) -> List[Factor]:
        """Generate a large factor pool from the best discovered rubric."""
        if self.best_rubric is None:
            raise RuntimeError("Run discover_rubric() first")
        print(f"\n{'='*70}")
        print(f"  PHASE 2: FINAL FACTOR POOL  (best rubric, n={n})")
        print(f"{'='*70}")
        factors = self.generator.generate(self.best_rubric, n=n)
        for f in factors:
            f.rubric_id = "best"
            self.evaluator.evaluate(f)
            self.all_factors.append(f)

        valid = [f for f in factors if f.valid]
        valid.sort(key=lambda f: f.ic, reverse=True)
        return valid

    @staticmethod
    def _print_rubric(rubric: FactorRubric):
        print("\nBest discovered rubric:")
        for item in rubric.items:
            print(f"  • {item.id}: {item.description}")
            print(f"    features={item.feature_focus}  ops={item.preferred_ops}  "
                  f"horizon={item.time_horizon_hint}  complexity={item.complexity_hint}")

    def summary(self) -> Dict:
        """Produce a JSON-serialisable summary."""
        valid_factors = [f for f in self.all_factors if f.valid]
        if not valid_factors:
            return {}

        top = max(valid_factors, key=lambda f: f.ic)
        return {
            "best_rubric": self.best_rubric.to_dict() if self.best_rubric else None,
            "best_fitness": float(self.best_fitness),
            "total_factors_generated": len(self.all_factors),
            "valid_factors": len(valid_factors),
            "top_factor": top.to_dict(),
            "mean_ic": float(np.mean([f.ic for f in valid_factors])),
            "mean_mae": float(np.mean([f.mae for f in valid_factors])),
            "mean_sharpe": float(np.mean([f.sharpe for f in valid_factors])),
            "search_history": self.history,
        }

# ═════════════════════════════════════════════════════════════════════════════
# 7.  DRIVER
# ═════════════════════════════════════════════════════════════════════════════

def main():
    print("\n" + "=" * 70)
    print("  AFRES v2: Agentic Factor Revision and Evaluation System")
    print("  (Faithful APRES rubric-discovery adapted to factor generation)")
    print("=" * 70 + "\n")

    t0 = time.time()

    # 1. Market data
    print("[1/4]  Loading market data …")
    data = MarketData(n_stocks=80, n_days=600, seed=42)

    # 2. Components
    print("[2/4]  Initialising components …")
    generator = LLMFactorGenerator(seed=99)
    evaluator = RegressionEvaluator(data, model_type="ridge")
    proposer = RubricProposer(N0=3, N=3, p_debug=0.3, D_max=5, seed=77)

    # 3. AFRES system
    afres = AFRES(data, generator, evaluator, proposer, factors_per_rubric=4)

    # 4. Run rubric discovery
    print("[3/4]  Running rubric-discovery loop …")
    best_rubric, best_factors = afres.discover_rubric(max_iterations=15)

    # 5. Final pool
    print("\n[4/4]  Generating final factor pool …")
    final_pool = afres.generate_final_pool(n=20)

    # 6. Report
    print("\n" + "=" * 70)
    print("  FINAL REPORT")
    print("=" * 70)
    print(f"\nTotal factors evaluated: {len(afres.all_factors)}")
    print(f"Valid factors: {len([f for f in afres.all_factors if f.valid])}")
    print(f"\nTop 5 factors by IC:")
    for f in final_pool[:5]:
        print(f"  {f.id:8s}  IC={f.ic:+.4f}  MAE={f.mae:.4f}  "
              f"Sharpe={f.sharpe:6.2f}  {f.expression}")

    # 7. Save
    results = afres.summary()
    with open("/app/afres_v2_results.json", "w") as fp:
        json.dump(results, fp, indent=2, default=str)
    print(f"\nResults saved to /app/afres_v2_results.json")
    print(f"Total wall time: {time.time()-t0:.1f}s")


if __name__ == "__main__":
    main()