afres_system.py

"""
AFRES: Agentic Factor Revision and Evaluation System

An adaptation of APRES for financial factor generation with QD algorithms.
Architecture:
  1. Factor Evaluation Rubric Discovery (APRES-style)
  2. QD-Enhanced Factor Archive (MAP-Elites)
  3. Iterative Factor Revision
"""

import json
import random
import numpy as np
import pandas as pd
from dataclasses import dataclass, field
from typing import List, Dict, Tuple, Optional, Callable, Any
from enum import Enum
import copy
import re
import time

# ============================================================================
# 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"

class TimeHorizon(Enum):
    SHORT = "short"
    MEDIUM = "medium"
    LONG = "long"

@dataclass
class Factor:
    id: str
    name: str
    description: str
    expression: str
    signal_type: SignalType
    time_horizon: TimeHorizon
    complexity: int
    ic: float = 0.0
    sharpe: float = 0.0
    returns: float = 0.0
    max_drawdown: float = 0.0
    rubric_scores: Dict[str, float] = field(default_factory=dict)
    overall_score: float = 0.0
    generation: int = 0
    parent_ids: List[str] = field(default_factory=list)

    def to_dict(self):
        return {
            "id": self.id, "name": self.name, "description": self.description,
            "expression": self.expression,
            "signal_type": self.signal_type.value, "time_horizon": self.time_horizon.value,
            "complexity": self.complexity,
            "ic": self.ic, "sharpe": self.sharpe, "returns": self.returns,
            "max_drawdown": self.max_drawdown,
            "rubric_scores": self.rubric_scores, "overall_score": self.overall_score,
            "generation": self.generation, "parent_ids": self.parent_ids,
        }

@dataclass
class RubricItem:
    id: str
    name: str
    description: str
    weight: float = 1.0

    def to_dict(self):
        return {"id": self.id, "name": self.name, "description": self.description, "weight": self.weight}

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

# ============================================================================
# EFFICIENT MARKET DATA (pre-computed)
# ============================================================================

class PrecomputedMarketData:
    """Pre-computed features for fast factor evaluation."""

    def __init__(self, n_stocks: int = 50, n_days: int = 500, seed: int = 42):
        np.random.seed(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)]

        # Generate panel data: each column is a stock, each row is a day
        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):
            ret = np.random.normal(0.0003, 0.015, n_days)
            # Add momentum signal: past 5-day ret predicts next day
            mom = np.zeros(n_days)
            mom[5:] = 0.2 * ret[:-5]
            ret[5:] += mom[5:]

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

        # Pre-compute common features as 2D arrays (days x stocks)
        self.ret_1d = np.diff(self.close, axis=0, prepend=self.close[0: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)

        self.vol_20d = pd.DataFrame(self.ret_1d).rolling(20, min_periods=1).std().values
        self.sma_5 = pd.DataFrame(self.close).rolling(5, min_periods=1).mean().values
        self.sma_20 = pd.DataFrame(self.close).rolling(20, min_periods=1).mean().values
        self.vol_sma_20 = pd.DataFrame(self.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 + np.random.normal(0, 0.0003, (n_days, n_stocks)))

        # Future returns (next day)
        self.future_ret_1d = np.zeros_like(self.close)
        self.future_ret_1d[:-1] = np.diff(self.close, axis=0) / (self.close[:-1] + 1e-10)

    def eval_expression(self, expr: str) -> Optional[np.ndarray]:
        """Evaluate a factor expression across all (days, stocks)."""
        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,
            'volatility_20d': self.vol_20d, 'sma_5': self.sma_5, 'sma_20': self.sma_20,
            '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_rolling(x, w, 'mean'),
            'ts_std': lambda x, w: self._ts_rolling(x, w, 'std'),
            'ts_max': lambda x, w: self._ts_rolling(x, w, 'max'),
            'ts_min': lambda x, w: self._ts_rolling(x, w, 'min'),
            'ts_zscore': lambda x, w: (x - self._ts_rolling(x, w, 'mean')) / (self._ts_rolling(x, w, 'std') + 1e-10),
            'ts_delta': lambda x, w: self._ts_delta(x, w),
            'ts_corr': lambda x, y, w: self._ts_corr(x, y, w),
            'ts_cov': lambda x, y, w: self._ts_cov(x, y, w),
            'ts_rank': lambda x, w: self._ts_rank(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

    def _rank(self, x):
        """Cross-sectional rank per day."""
        r = np.zeros_like(x)
        for t in range(x.shape[0]):
            valid = ~np.isnan(x[t])
            if valid.sum() > 0:
                r[t, valid] = pd.Series(x[t, valid]).rank(pct=True).values
        return r

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

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

    def _ts_corr(self, x, y, w):
        r = np.zeros_like(x)
        for t in range(x.shape[0]):
            start = max(0, t - w + 1)
            if start < t:
                a = x[start:t+1].flatten()
                b = y[start: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 _ts_cov(self, x, y, w):
        c = np.zeros_like(x)
        for t in range(x.shape[0]):
            start = max(0, t - w + 1)
            if start < t:
                a = x[start:t+1].flatten()
                b = y[start:t+1].flatten()
                c[t] = np.cov(a, b)[0, 1] if len(a) > 1 else 0
        return c

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

# ============================================================================
# FACTOR EVALUATOR
# ============================================================================

class FactorEvaluator:
    def __init__(self, data: PrecomputedMarketData):
        self.data = data
        # Use days 100-350 for train, 350-450 for test (skip first 20 for feature warmup)
        self.test_idx = np.arange(350, min(450, data.n_days - 1))

    def evaluate(self, factor: Factor) -> Dict[str, float]:
        vals = self.data.eval_expression(factor.expression)
        if vals is None:
            return {"ic": 0.0, "sharpe": 0.0, "returns": 0.0, "max_drawdown": 0.0}

        # IC calculation: rank correlation per day
        ics = []
        for t in self.test_idx:
            f_t = vals[t]
            r_t = self.data.future_ret_1d[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)

        mean_ic = np.mean(ics) if ics else 0.0

        # Portfolio simulation: long top 20% each day
        port_rets = []
        for t in self.test_idx:
            f_t = vals[t]
            r_t = self.data.future_ret_1d[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(np.mean(r_t[mask]))

        if len(port_rets) > 2:
            rets = np.array(port_rets)
            ann_ret = np.mean(rets) * 252
            ann_vol = np.std(rets) * np.sqrt(252)
            sharpe = ann_ret / (ann_vol + 1e-10)
            cum = np.cumsum(rets)
            running_max = np.maximum.accumulate(cum)
            dd = cum - running_max
            max_dd = np.min(dd) if len(dd) > 0 else 0.0
        else:
            ann_ret = 0.0
            sharpe = 0.0
            max_dd = 0.0

        return {"ic": mean_ic, "sharpe": sharpe, "returns": ann_ret, "max_drawdown": max_dd}

# ============================================================================
# AGENTS
# ============================================================================

class RubricProposer:
    INITIAL_RUBRIC = [
        RubricItem("predictive_power", "Predictive Power", "How well does the factor predict future returns?"),
        RubricItem("robustness", "Robustness", "How stable is the factor across market conditions?"),
        RubricItem("interpretability", "Interpretability", "How easy to understand the economic rationale?"),
        RubricItem("complexity", "Appropriate Complexity", "Complex enough but not overfitting."),
        RubricItem("diversity", "Diversity", "How different from existing factors?"),
        RubricItem("turnover", "Turnover Friendliness", "Reasonable transaction costs."),
    ]

    def __init__(self):
        self.rng = np.random.RandomState(42)

    def propose(self, existing=None, feedback=None):
        if existing is None:
            return FactorRubric([copy.deepcopy(i) for i in self.INITIAL_RUBRIC])
        items = [copy.deepcopy(i) for i in existing.items]
        if feedback and "best_item" in feedback:
            for item in items:
                if item.id == feedback["best_item"]:
                    item.weight = min(2.0, item.weight * 1.2)
        if self.rng.random() < 0.3:
            extras = [
                RubricItem("nonlinearity", "Non-linearity Capture", "Captures non-linear dynamics"),
                RubricItem("regime_adaptivity", "Regime Adaptivity", "Performs across regimes"),
                RubricItem("cs_consistency", "Cross-sectional Consistency", "Works across stocks"),
                RubricItem("lag_structure", "Appropriate Lag", "Uses available information only"),
            ]
            new_item = copy.deepcopy(self.rng.choice(extras))
            if not any(i.id == new_item.id for i in items):
                items.append(new_item)
        return FactorRubric(items)

class FactorReviewer:
    def __init__(self):
        self.rng = np.random.RandomState(43)

    def review(self, factor: Factor, rubric: FactorRubric, evaluator: FactorEvaluator):
        metrics = evaluator.evaluate(factor)
        factor.ic = metrics["ic"]
        factor.sharpe = metrics["sharpe"]
        factor.returns = metrics["returns"]
        factor.max_drawdown = metrics["max_drawdown"]

        scores = {}
        for item in rubric.items:
            if item.id == "predictive_power":
                s = min(1.0, max(0.0, (factor.ic + 0.05) / 0.15))
            elif item.id == "robustness":
                s = min(1.0, max(0.0, (factor.sharpe + 0.5) / 2.0))
            elif item.id == "interpretability":
                s = min(1.0, max(0.0, 1.0 - factor.complexity / 10.0))
            elif item.id == "complexity":
                s = min(1.0, max(0.0, 1.0 - abs(factor.complexity - 4) / 4.0))
            elif item.id == "diversity":
                s = self.rng.uniform(0.3, 0.8)
            elif item.id == "turnover":
                s = 0.8 if factor.time_horizon == TimeHorizon.LONG else 0.6 if factor.time_horizon == TimeHorizon.MEDIUM else 0.4
            elif item.id == "nonlinearity":
                s = 0.5 + 0.3 * min(1.0, factor.complexity / 8.0)
            elif item.id == "regime_adaptivity":
                s = min(1.0, max(0.0, (factor.sharpe + 0.5) / 2.0))
            elif item.id == "cs_consistency":
                s = min(1.0, max(0.0, (abs(factor.ic) + 0.02) / 0.12))
            elif item.id == "lag_structure":
                s = 0.7 if factor.time_horizon != TimeHorizon.SHORT else 0.5
            else:
                s = self.rng.uniform(0.4, 0.7)
            scores[item.id] = min(1.0, max(0.0, s + self.rng.normal(0, 0.05)))

        total_w = sum(i.weight for i in rubric.items)
        overall = sum(scores.get(i.id, 0.0) * i.weight for i in rubric.items) / max(total_w, 1e-10)
        factor.rubric_scores = scores
        factor.overall_score = overall
        return scores

class FactorGenerator:
    SEEDS = [
        ("returns_5d", SignalType.PRICE_BASED, TimeHorizon.SHORT, 1),
        ("returns_20d", SignalType.PRICE_BASED, TimeHorizon.MEDIUM, 1),
        ("ts_mean(returns_1d, 5)", SignalType.PRICE_BASED, TimeHorizon.SHORT, 2),
        ("ts_corr(close, volume, 20)", SignalType.VOLUME_BASED, TimeHorizon.MEDIUM, 2),
        ("(close - sma_20) / ts_std(close, 20)", SignalType.TECHNICAL, TimeHorizon.MEDIUM, 3),
        ("rank(volume) * rank(returns_5d)", SignalType.VOLUME_BASED, TimeHorizon.SHORT, 3),
        ("ts_delta(close, 5) / ts_mean(close, 20)", SignalType.TECHNICAL, TimeHorizon.SHORT, 3),
        ("ts_zscore(volume, 20) * ts_zscore(returns_1d, 5)", SignalType.CROSS_SECTIONAL, TimeHorizon.SHORT, 4),
        ("(high_20d - low) / (high_20d - low_20d + 1e-10)", SignalType.TECHNICAL, TimeHorizon.MEDIUM, 3),
        ("log(volume / volume_sma_20) * returns_5d", SignalType.VOLUME_BASED, TimeHorizon.SHORT, 3),
    ]

    def __init__(self):
        self.rng = np.random.RandomState(44)
        self._c = 0

    def generate_seeds(self, n=10):
        factors = []
        for i, (expr, st, th, comp) in enumerate(self.SEEDS[:n]):
            factors.append(Factor(
                id=f"seed_{i}", name=f"Seed {i}",
                description=f"Seed using {st.value}", expression=expr,
                signal_type=st, time_horizon=th, complexity=comp, generation=0))
        return factors

    def mutate(self, parent: Factor, mutation="random"):
        self._c += 1
        if mutation == "operator_replacement":
            expr = self._replace_op(parent.expression)
        elif mutation == "parameter_change":
            expr = self._change_param(parent.expression)
        elif mutation == "feature_swap":
            expr = self._swap_feat(parent.expression)
        else:
            expr = self._random_mod(parent.expression)

        st = parent.signal_type
        th = parent.time_horizon
        if self.rng.random() < 0.2:
            st = self.rng.choice(list(SignalType))
        if self.rng.random() < 0.2:
            th = self.rng.choice(list(TimeHorizon))

        return Factor(
            id=f"mut_{self._c}", name=f"Mutated {parent.name}",
            description=f"Mutation of {parent.id}", expression=expr,
            signal_type=st, time_horizon=th,
            complexity=max(1, parent.complexity + self.rng.randint(-1, 2)),
            generation=parent.generation + 1, parent_ids=[parent.id])

    def _replace_op(self, expr):
        repl = {
            'ts_mean': ['ts_median', 'ts_sum'],
            'ts_std': ['ts_var', 'ts_mad'],
            'ts_corr': ['ts_cov', 'ts_beta'],
            'ts_zscore': ['ts_rank', 'ts_delta'],
            'rank': ['sign', 'abs'],
            'log': ['sqrt', 'sign'],
        }
        for old, new_list in repl.items():
            if old in expr:
                expr = expr.replace(old, self.rng.choice(new_list), 1)
                break
        return expr

    def _change_param(self, expr):
        nums = re.findall(r'\d+', expr)
        if nums:
            n = self.rng.choice(nums)
            new_n = str(max(1, int(n) + self.rng.randint(-3, 4)))
            if new_n != n:
                expr = expr.replace(n, new_n, 1)
        return expr

    def _swap_feat(self, expr):
        feats = ['close', 'open', 'high', 'low', 'volume', 'vwap',
                 'returns_1d', 'returns_5d', 'returns_20d', 'sma_5', 'sma_20']
        for _ in range(3):
            old = self.rng.choice(feats)
            if old in expr:
                new = self.rng.choice([f for f in feats if f != old])
                expr = expr.replace(old, new, 1)
                break
        return expr

    def _random_mod(self, expr):
        mods = [self._replace_op, self._change_param, self._swap_feat]
        return self.rng.choice(mods)(expr)

class FactorRewriter:
    def __init__(self):
        self.rng = np.random.RandomState(45)

    def revise(self, factor: Factor, scores: Dict, rubric: FactorRubric):
        sorted_scores = sorted(scores.items(), key=lambda x: x[1])
        expr = factor.expression
        for item_id, score in sorted_scores[:2]:
            if score > 0.7:
                continue
            if item_id == "predictive_power":
                expr = f"ts_zscore(({expr}), 20)"
            elif item_id == "robustness":
                expr = f"ts_mean(({expr}), 5)"
            elif item_id == "interpretability":
                expr = self._simplify(expr)
            elif item_id == "complexity":
                if factor.complexity > 6:
                    expr = self._simplify(expr)
                else:
                    expr = f"({expr}) * rank(volume)"
            elif item_id == "diversity":
                expr = f"sign({expr}) * abs(returns_5d)"
            elif item_id == "turnover":
                expr = f"ts_mean(({expr}), 10)"
            elif item_id == "nonlinearity":
                expr = f"abs({expr}) * sign(returns_1d)"
            elif item_id == "regime_adaptivity":
                expr = f"({expr}) / (volatility_20d + 1e-10)"
        return Factor(
            id=f"rev_{factor.id}", name=f"Revised {factor.name}",
            description=f"Revision of {factor.id}", expression=expr,
            signal_type=factor.signal_type, time_horizon=factor.time_horizon,
            complexity=factor.complexity + 1, generation=factor.generation + 1,
            parent_ids=[factor.id])

    def _simplify(self, expr):
        expr = re.sub(r'ts_mean\(([^,]+),\s*\d+\)', r'\1', expr)
        expr = re.sub(r'ts_zscore\(([^,]+),\s*\d+\)', r'\1', expr)
        return expr

# ============================================================================
# QD ARCHIVE
# ============================================================================

class QDArchive:
    def __init__(self, behavior_dims, fitness_func):
        self.behavior_dims = behavior_dims
        self.fitness_func = fitness_func
        self.rng = np.random.RandomState(123)
        self.grid_shape = tuple(len(d[1]) for d in behavior_dims)
        self.archive = {}
        self.fitness_grid = np.full(self.grid_shape, -np.inf)
        self.all_factors = []

    def _idx(self, factor):
        idx = []
        for dim_name, dim_vals in self.behavior_dims:
            if dim_name == "signal_type":
                val = factor.signal_type
            elif dim_name == "time_horizon":
                val = factor.time_horizon
            elif dim_name == "complexity_bin":
                val = "low" if factor.complexity <= 2 else "medium" if factor.complexity <= 5 else "high"
            else:
                val = "unknown"
            try:
                idx.append(dim_vals.index(val))
            except ValueError:
                idx.append(0)
        return tuple(idx)

    def add(self, factor):
        idx = self._idx(factor)
        fitness = self.fitness_func(factor)
        self.all_factors.append(factor)
        if fitness > self.fitness_grid[idx]:
            self.archive[idx] = factor
            self.fitness_grid[idx] = fitness
            return True
        return False

    def get_random_elite(self):
        if not self.archive:
            return None
        return self.rng.choice(list(self.archive.values()))

    def get_best(self):
        if not self.archive:
            return None
        return max(self.archive.values(), key=self.fitness_func)

    def get_coverage(self):
        total = np.prod(self.grid_shape)
        return len(self.archive) / total if total > 0 else 0

    def get_all_elites(self):
        return list(self.archive.values())

    def summary(self):
        elites = self.get_all_elites()
        if not elites:
            return {"coverage": 0.0, "num_elites": 0, "mean_fitness": 0.0, "max_fitness": 0.0}
        fits = [self.fitness_func(e) for e in elites]
        return {"coverage": self.get_coverage(), "num_elites": len(elites),
                "mean_fitness": float(np.mean(fits)), "max_fitness": float(np.max(fits)),
                "min_fitness": float(np.min(fits))}

# ============================================================================
# AFRES SYSTEM
# ============================================================================

class AFRES:
    def __init__(self, data: PrecomputedMarketData):
        self.data = data
        self.evaluator = FactorEvaluator(data)
        self.rubric_proposer = RubricProposer()
        self.factor_reviewer = FactorReviewer()
        self.factor_generator = FactorGenerator()
        self.factor_rewriter = FactorRewriter()
        self.rubric = None
        self.qd_archive = None
        self.factor_library = []
        self.iteration = 0
        self.rubric_history = []

    def discover_rubric(self, n_iter=5, n_per_iter=10):
        print("=" * 60)
        print("PHASE 1: FACTOR EVALUATION RUBRIC DISCOVERY")
        print("=" * 60)
        seeds = self.factor_generator.generate_seeds(n_per_iter)
        best_rubric = None
        best_mae = float('inf')

        for it in range(n_iter):
            t0 = time.time()
            if it == 0:
                rubric = self.rubric_proposer.propose()
            else:
                feedback = {"best_item": self._find_best_item(best_rubric, seeds)}
                rubric = self.rubric_proposer.propose(best_rubric, feedback)

            # Evaluate seeds with this rubric
            for f in seeds:
                self.factor_reviewer.review(f, rubric, self.evaluator)

            mae = self._compute_mae(seeds, rubric)
            print(f"Iter {it+1}/{n_iter}: {len(rubric.items)} items, MAE={mae:.4f} ({time.time()-t0:.1f}s)")

            if mae < best_mae:
                best_mae = mae
                best_rubric = rubric
                print(f"  -> New best rubric!")
            self.rubric_history.append({"iteration": it, "n_items": len(rubric.items), "mae": mae})

        self.rubric = best_rubric
        print(f"\nBest rubric ({len(best_rubric.items)} items, MAE={best_mae:.4f}):")
        for item in best_rubric.items:
            print(f"  - {item.name} (w={item.weight:.2f})")
        return best_rubric

    def _find_best_item(self, rubric, factors):
        best_corr = -1
        best_id = rubric.items[0].id
        for item in rubric.items:
            s = [f.rubric_scores.get(item.id, 0.0) for f in factors]
            ics = [f.ic for f in factors]
            if len(s) > 1 and np.std(s) > 0:
                corr = abs(np.corrcoef(s, ics)[0, 1])
                if not np.isnan(corr) and corr > best_corr:
                    best_corr = corr
                    best_id = item.id
        return best_id

    def _compute_mae(self, factors, rubric):
        pred = np.array([f.overall_score for f in factors])
        actual = np.array([f.ic for f in factors])
        if np.std(pred) > 0:
            p_scaled = (pred - np.mean(pred)) / (np.std(pred) + 1e-10)
            p_scaled = p_scaled * np.std(actual) + np.mean(actual)
            return float(np.mean(np.abs(p_scaled - actual)))
        return float(np.mean(np.abs(actual)))

    def run_qd_search(self, n_iter=15, n_mutate=5):
        print("\n" + "=" * 60)
        print("PHASE 2: QD-ENHANCED FACTOR SEARCH (MAP-Elites)")
        print("=" * 60)

        if self.rubric is None:
            raise ValueError("Discover rubric first!")

        behavior_dims = [
            ("signal_type", list(SignalType)),
            ("time_horizon", list(TimeHorizon)),
            ("complexity_bin", ["low", "medium", "high"]),
        ]
        self.qd_archive = QDArchive(behavior_dims,
                                     lambda f: f.ic + f.sharpe / 3.0 + f.overall_score / 2.0)

        seeds = self.factor_generator.generate_seeds(10)
        for f in seeds:
            self.factor_reviewer.review(f, self.rubric, self.evaluator)
            self.qd_archive.add(f)
            self.factor_library.append(f)

        s = self.qd_archive.summary()
        print(f"Initial: coverage={s['coverage']:.1%}, elites={s['num_elites']}, mean_fit={s['mean_fitness']:.3f}")

        for it in range(n_iter):
            t0 = time.time()
            added_count = 0
            for _ in range(n_mutate):
                parent = self.qd_archive.get_random_elite()
                if parent is None:
                    continue
                mutation = np.random.choice(["operator", "param", "feature", "random"])
                child = self.factor_generator.mutate(parent, mutation)
                self.factor_reviewer.review(child, self.rubric, self.evaluator)
                if self.qd_archive.add(child):
                    added_count += 1
                self.factor_library.append(child)
            s = self.qd_archive.summary()
            print(f"Iter {it+1}/{n_iter}: +{added_count} new, coverage={s['coverage']:.1%}, "
                  f"elites={s['num_elites']}, max_fit={s['max_fitness']:.3f} ({time.time()-t0:.1f}s)")

        return self.qd_archive

    def iterative_revision(self, n_iter=3):
        print("\n" + "=" * 60)
        print("PHASE 3: ITERATIVE FACTOR REVISION")
        print("=" * 60)

        if self.rubric is None:
            raise ValueError("Discover rubric first!")

        if self.qd_archive:
            candidates = sorted(self.qd_archive.get_all_elites(),
                                  key=lambda f: self.qd_archive.fitness_func(f), reverse=True)[:5]
        else:
            candidates = self.factor_generator.generate_seeds(5)
            for f in candidates:
                self.factor_reviewer.review(f, self.rubric, self.evaluator)

        improved = []
        for factor in candidates:
            current = factor
            print(f"\nRevising {current.id} (score={current.overall_score:.3f}, IC={current.ic:.4f})")
            for it in range(n_iter):
                scores = self.factor_reviewer.review(current, self.rubric, self.evaluator)
                revised = self.factor_rewriter.revise(current, scores, self.rubric)
                self.factor_reviewer.review(revised, self.rubric, self.evaluator)
                print(f"  Iter {it+1}: prev={current.overall_score:.3f} -> rev={revised.overall_score:.3f} "
                      f"(IC: {current.ic:.4f} -> {revised.ic:.4f})")
                if revised.overall_score > current.overall_score:
                    current = revised
            improved.append(current)
            print(f"  Final: score={current.overall_score:.3f}, IC={current.ic:.4f}, Sharpe={current.sharpe:.3f}")
        return improved

    def run(self, rubric_iter=5, qd_iter=15, revision_iter=3):
        print("\n" + "=" * 70)
        print(" AFRES: AGENTIC FACTOR REVISION AND EVALUATION SYSTEM")
        print("=" * 70)
        t0_total = time.time()

        self.discover_rubric(rubric_iter)
        self.run_qd_search(qd_iter)
        best = self.iterative_revision(revision_iter)

        print("\n" + "=" * 70)
        print(" FINAL RESULTS")
        print("=" * 70)

        if self.qd_archive:
            elites = self.qd_archive.get_all_elites()
            if elites:
                top = max(elites, key=lambda f: f.overall_score)
                print(f"\nBest factor overall: {top.id}")
                print(f"  Expression: {top.expression}")
                print(f"  IC: {top.ic:.4f} | Sharpe: {top.sharpe:.3f} | Score: {top.overall_score:.3f}")
                print(f"  Type: {top.signal_type.value} | Horizon: {top.time_horizon.value} | Gen: {top.generation}")

        print(f"\nArchive coverage: {self.qd_archive.get_coverage():.1%}")
        print(f"Total factors: {len(self.factor_library)}")
        print(f"Total time: {time.time()-t0_total:.1f}s")

        return {
            "rubric": self.rubric.to_dict() if self.rubric else None,
            "archive_summary": self.qd_archive.summary() if self.qd_archive else None,
            "best_factors": [f.to_dict() for f in best],
            "total_factors": len(self.factor_library),
            "rubric_history": self.rubric_history,
        }

# ============================================================================
# MAIN
# ============================================================================

def main():
    print("Creating synthetic market data...")
    data = PrecomputedMarketData(n_stocks=50, n_days=500, seed=42)
    afres = AFRES(data)
    results = afres.run(rubric_iter=5, qd_iter=15, revision_iter=3)

    with open("/app/afres_results.json", "w") as f:
        json.dump(results, f, indent=2, default=str)
    print("\nResults saved to /app/afres_results.json")

if __name__ == "__main__":
    main()