database/models.py

from datetime import datetime, timezone
from sqlalchemy import Column, Integer, String, Float, DateTime, Text
from sqlalchemy.orm import declarative_base

Base = declarative_base()

class Alpha(Base):
    """
    SQLAlchemy model representing a WorldQuant BRAIN Alpha candidate in QuantForge v2.0.
    Stores alpha metadata, simulation configuration parameters, and backtest results.
    """
    __tablename__ = "alphas"

    id = Column(Integer, primary_key=True, autoincrement=True)
    expression = Column(String, unique=True, nullable=False, index=True)
    language = Column(String, default="FASTEXPR", nullable=False)
    region = Column(String, nullable=False)
    universe = Column(String, nullable=False)
    delay = Column(Integer, nullable=False)
    decay = Column(Integer, nullable=False)
    truncation = Column(Float, nullable=False)
    neutralization = Column(String, nullable=False)
    pasteurization = Column(String, default="On", nullable=False)
    
    # Advanced settings (nullable for versatility or auto-detection)
    nan_handling = Column(String, nullable=True)
    unit_handling = Column(String, nullable=True)
    lookback = Column(Integer, nullable=True)
    test_period = Column(String, nullable=True)
    max_trade = Column(Float, nullable=True)
    max_position = Column(Float, nullable=True)
    
    # Backtest Metrics
    sharpe = Column(Float, nullable=True)
    fitness = Column(Float, nullable=True)
    turnover = Column(Float, nullable=True)
    returns = Column(Float, nullable=True)
    
    # Mining Orchestration State
    # Expected States: PENDING_GEN, SIMULATING, REFINING, PENDING_CORR, TRASHED, SUBMITTED
    status = Column(String, default="PENDING_GEN", nullable=False)
    retry_count = Column(Integer, default=0, nullable=False)
    error_log = Column(String, nullable=True)
    trash_reason = Column(String, nullable=True)  # Category: Syntax/Schema, Low Performance, High Correlation, No Trades
    alpha_id = Column(String, nullable=True)
    raw_result_json = Column(Text, nullable=True)
    
    date_created = Column(DateTime, default=lambda: datetime.now(timezone.utc), nullable=False)
    date_submitted = Column(DateTime, nullable=True)
    created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), nullable=False)

    def __repr__(self) -> str:
        return (
            f"<Alpha(id={self.id}, expression={self.expression[:40]}..., "
            f"region={self.region}, universe={self.universe}, status={self.status})>"
        )