submission_engine.py

import asyncio
import logging
import time
from datetime import datetime, timezone, timedelta
from database import SessionLocal, AlphaRecord
from wq_client import BrainClient
import config

logger = logging.getLogger("SubmissionEngine")

class SubmissionEngine:
    def __init__(self, client: BrainClient):
        self.client = client
        self.daily_submission_limit = 20  # Updated limit from 8 to 20
        
        # Throttle polling to at most once every 5 minutes (300 seconds)
        self.last_status_poll_time = 0.0

    def get_daily_submission_count(self, db) -> int:
        """Returns the number of alphas submitted in the last 24 hours."""
        cutoff = datetime.utcnow() - timedelta(hours=24)
        return db.query(AlphaRecord).filter(
            AlphaRecord.status == "SUBMITTED",
            AlphaRecord.date_submitted >= cutoff
        ).count()

    async def clear_legacy_queue(self):
        """
        Process and submit all alphas that are currently in the legacy submission queue
        (PASSED_CORRELATION, QUEUED_FOR_NEXT_BATCH) using existing queue contents.
        """
        logger.info("Clearing legacy submission queue...")
        db = SessionLocal()
        try:
            # Query all legacy queued alphas
            legacy_alphas = db.query(AlphaRecord).filter(
                AlphaRecord.status.in_(["PASSED_CORRELATION", "QUEUED_FOR_NEXT_BATCH"])
            ).all()
            
            if not legacy_alphas:
                logger.info("No legacy queued alphas found.")
                return
                
            logger.info(f"Found {len(legacy_alphas)} legacy queued alphas. Processing submission...")
            
            # Update all status to PASSED_CORRELATION for consistency, then submit
            for record in legacy_alphas:
                if record.status == "QUEUED_FOR_NEXT_BATCH":
                    record.status = "PASSED_CORRELATION"
            db.commit()
            
            # Submit ready alphas up to daily limit
            await self.submit_ready_alphas_with_db(db)
            
        except Exception as e:
            logger.exception("Error during legacy queue clearing.")
            db.rollback()
        finally:
            db.close()

    async def submit_ready_alphas_with_db(self, db):
        """Helper to submit any PASSED_CORRELATION alphas up to daily cap."""
        ready_alphas = db.query(AlphaRecord).filter_by(status="PASSED_CORRELATION").all()
        if not ready_alphas:
            return
            
        logger.info(f"Attempting to submit {len(ready_alphas)} ready alphas...")
        
        # Rank by Fitness * Sharpe to submit the best first if close to limit
        def score_alpha(rec):
            fit = rec.fitness or 0.0
            shp = rec.sharpe or 0.0
            return fit * shp
            
        ready_alphas.sort(key=score_alpha, reverse=True)
        
        for record in ready_alphas:
            daily_count = self.get_daily_submission_count(db)
            if daily_count >= self.daily_submission_limit:
                logger.warning(f"Daily submission limit ({self.daily_submission_limit}) reached. Cannot submit more alphas today.")
                break
                
            alpha_id = record.alpha_id
            # Fallback search if alpha ID is missing
            if not alpha_id or alpha_id == "mock_id_for_pipeline":
                logger.info(f"Alpha ID missing for {record.expression[:45]}... Searching on platform...")
                status_data = await self.client.check_alpha_status_async(record.expression)
                if status_data and status_data.get("id") and status_data.get("id") != "mock_id_for_pipeline":
                    alpha_id = status_data.get("id")
                    record.alpha_id = alpha_id
                    db.commit()
                else:
                    logger.warning(f"Could not retrieve a valid alpha ID for {record.expression[:45]}... Skipping.")
                    continue
                    
            logger.info(f"Submitting alpha {record.expression[:45]}... (ID: {alpha_id})")
            submit_success = await self.client.submit_alpha_async(alpha_id)
            if submit_success:
                record.status = "SUBMITTED"
                record.date_submitted = datetime.utcnow()
                logger.info(f"SUCCESS! Alpha {record.expression[:45]}... officially submitted immediately.")
            else:
                logger.error(f"Failed to submit alpha {record.expression[:45]}...")
            db.commit()
            await asyncio.sleep(5)

    async def poll_pending_correlation_statuses(self):
        """
        Polls the correlation status on the WQ platform for any alphas with status 'PENDING_CORRELATION'.
        Runs status checks every 5 minutes to avoid rate limits.
        """
        now_ts = time.time()
        if now_ts - self.last_status_poll_time < 300:
            return
            
        db = SessionLocal()
        try:
            pending_alphas = db.query(AlphaRecord).filter_by(status="PENDING_CORRELATION").all()
            if not pending_alphas:
                return
                
            logger.info(f"Polling correlation status for {len(pending_alphas)} alphas on WQ platform...")
            self.last_status_poll_time = now_ts
            
            for record in pending_alphas:
                logger.info(f"Checking correlation status for alpha ID: {record.alpha_id} (Expr: {record.expression[:45]}...)")
                status_data = await self.client.check_alpha_status_async(record.expression, record.alpha_id)
                
                if not status_data:
                    logger.warning(f"Could not retrieve status for alpha: {record.expression}. Skipping.")
                    continue
                    
                status = status_data.get("status", "")
                
                # Update alpha ID in database if it was missing or mock
                if status_data.get("id") and status_data.get("id") != "mock_id_for_pipeline":
                    record.alpha_id = status_data.get("id")
                    db.commit()
                
                if status in ("ACTIVE", "SUBMITTED"):
                    record.status = "SUBMITTED"
                    if not record.date_submitted:
                        record.date_submitted = datetime.utcnow()
                    logger.info(f"Alpha {record.expression[:45]}... is already ACTIVE/SUBMITTED on WorldQuant!")
                    db.commit()
                    continue
                
                # Extract checks from 'is' and 'os' blocks
                checks_info = []
                if isinstance(status_data.get("is"), dict):
                    checks_info.extend(status_data["is"].get("checks", []))
                if isinstance(status_data.get("os"), dict):
                    checks_info.extend(status_data["os"].get("checks", []))
                
                # Check for pending status or pending checks
                has_pending_checks = False
                for check in checks_info:
                    if isinstance(check, dict) and check.get("result") == "PENDING":
                        if check.get("name") == "SELF_CORRELATION":
                            continue
                        has_pending_checks = True
                        break
                
                if status in ("PENDING", "UNTESTED", "CHECKING") or has_pending_checks:
                    logger.info(f"Alpha {record.expression[:45]}... is still pending checks.")
                    continue
                
                # Perform strict check validations (Sub-universe Sharpe >= 0.58, Weight Concentration)
                check_failed = False
                failure_reasons = []
                for check in checks_info:
                    if not isinstance(check, dict):
                        continue
                    name = check.get("name")
                    result = check.get("result")
                    value = check.get("value")
                    
                    if result == "FAIL" or check.get("passed") == False:
                        check_failed = True
                        failure_reasons.append(f"{name} check failed on platform")
                    elif name == "LOW_SUB_UNIVERSE_SHARPE" and value is not None:
                        try:
                            if float(value) < 0.58:
                                check_failed = True
                                failure_reasons.append(f"Sub-universe Sharpe {value} < cutoff 0.58")
                        except (ValueError, TypeError):
                            pass
                
                if check_failed:
                    record.status = "FAILED_CORRELATION"
                    record.error_log = "FAILED_CHECKS: " + "; ".join(failure_reasons)
                    logger.warning(f"Alpha {record.expression[:45]}... failed strict checks: {record.error_log}")
                    db.commit()
                    continue
                
                if status in ("FAIL", "FAILED", "REJECTED", "CORRELATED", "FAILED_CORRELATION"):
                    record.status = "QUEUED_FOR_ORTHOGONALIZATION"
                    correlation_score = status_data.get("correlation", "Unknown (>0.70)")
                    record.error_log = f"FAILED_CORRELATION: Overlaps existing alphas. Score = {correlation_score}"
                    logger.warning(f"Alpha {record.expression[:45]}... failed correlation check (Score: {correlation_score}). Queueing for orthogonalization.")
                    db.commit()
                    
                elif status in ("PASS", "PASSED", "UNCORRELATED", "SUCCESS", "UNSUBMITTED"):
                    record.status = "PASSED_CORRELATION"
                    logger.info(f"Alpha {record.expression[:45]}... passed correlation check! Setting status to PASSED_CORRELATION.")
                    db.commit()
                    
                    # Submit immediately if limit allows
                    daily_count = self.get_daily_submission_count(db)
                    if daily_count < self.daily_submission_limit:
                        logger.info(f"Submitting alpha {record.expression[:45]}... immediately.")
                        submit_success = await self.client.submit_alpha_async(record.alpha_id)
                        if submit_success:
                            record.status = "SUBMITTED"
                            record.date_submitted = datetime.utcnow()
                            logger.info(f"SUCCESS! Alpha {record.expression[:45]}... officially submitted immediately.")
                        else:
                            logger.error(f"Failed to submit alpha {record.expression[:45]}...")
                        db.commit()
                    else:
                        logger.warning(f"Daily submission limit reached ({daily_count}/{self.daily_submission_limit}). Keeping status as PASSED_CORRELATION.")
                
                await asyncio.sleep(2)
                
        except Exception as e:
            logger.exception("Error during correlation status polling.")
            db.rollback()
        finally:
            db.close()

    async def sync_submitted_alphas(self):
        """
        Synchronizes the local database status of 'SUBMITTED' alphas with the WorldQuant platform.
        Runs once at startup to realign local DB status with reality.
        """
        logger.info("Starting synchronization of SUBMITTED alphas with WorldQuant platform...")
        db = SessionLocal()
        try:
            submitted_records = db.query(AlphaRecord).filter_by(status="SUBMITTED").all()
            if not submitted_records:
                logger.info("No SUBMITTED alphas found in local database to sync.")
                return
                
            for record in submitted_records:
                logger.info(f"Syncing status for alpha: {record.expression[:45]}... (ID: {record.alpha_id})")
                status_data = await self.client.check_alpha_status_async(record.expression, record.alpha_id)
                
                if not status_data:
                    logger.warning(f"Could not retrieve platform status for: {record.expression[:45]}. Keeping current status.")
                    continue
                
                # Update alpha ID if it was missing or mock
                if status_data.get("id") and status_data.get("id") != "mock_id_for_pipeline":
                    record.alpha_id = status_data.get("id")
                    
                status = status_data.get("status", "")
                logger.info(f"Platform status for {record.expression[:45]}... is '{status}'")
                
                if status in ("ACTIVE", "SUBMITTED"):
                    # Keep as SUBMITTED
                    record.status = "SUBMITTED"
                    if not record.date_submitted:
                        record.date_submitted = datetime.utcnow()
                    db.commit()
                    continue
                
                # Extract checks from 'is' and 'os' blocks
                checks_info = []
                if isinstance(status_data.get("is"), dict):
                    checks_info.extend(status_data["is"].get("checks", []))
                if isinstance(status_data.get("os"), dict):
                    checks_info.extend(status_data["os"].get("checks", []))
                    
                # Check for pending status or pending checks
                has_pending_checks = False
                for check in checks_info:
                    if isinstance(check, dict) and check.get("result") == "PENDING":
                        if check.get("name") == "SELF_CORRELATION":
                            continue
                        has_pending_checks = True
                        break
                
                if status in ("PENDING", "UNTESTED", "CHECKING") or has_pending_checks:
                    record.status = "PENDING_CORRELATION"
                    logger.info(f"Alpha {record.expression[:45]}... is still pending checks. Realignment set to PENDING_CORRELATION.")
                    db.commit()
                    continue
                    
                check_failed = False
                failure_reasons = []
                for check in checks_info:
                    if not isinstance(check, dict):
                        continue
                    name = check.get("name")
                    result = check.get("result")
                    value = check.get("value")
                    
                    if result == "FAIL" or check.get("passed") == False:
                        check_failed = True
                        failure_reasons.append(f"{name} check failed on platform")
                    elif name == "LOW_SUB_UNIVERSE_SHARPE" and value is not None:
                        try:
                            if float(value) < 0.58:
                                check_failed = True
                                failure_reasons.append(f"Sub-universe Sharpe {value} < cutoff 0.58")
                        except (ValueError, TypeError):
                            pass
                
                if check_failed:
                    record.status = "FAILED_CORRELATION"
                    record.error_log = "FAILED_CHECKS: " + "; ".join(failure_reasons)
                    logger.warning(f"Alpha {record.expression[:45]}... synced as FAILED_CORRELATION: {record.error_log}")
                elif status in ("FAIL", "FAILED", "REJECTED", "CORRELATED", "FAILED_CORRELATION"):
                    record.status = "FAILED_CORRELATION"
                    correlation_score = status_data.get("correlation", "Unknown (>0.70)")
                    record.error_log = f"FAILED_CORRELATION: Overlaps existing alphas. Score = {correlation_score}"
                    logger.warning(f"Alpha {record.expression[:45]}... synced as FAILED_CORRELATION: overlaps existing.")
                elif status in ("PASS", "PASSED", "UNCORRELATED", "SUCCESS", "UNSUBMITTED"):
                    record.status = "PASSED_CORRELATION"
                    logger.info(f"Alpha {record.expression[:45]}... synced as PASSED_CORRELATION (ready to submit).")
                else:
                    logger.info(f"Platform status '{status}' is not definitive. Keeping status as SUBMITTED.")
                    
                db.commit()
                await asyncio.sleep(2)
        except Exception as e:
            logger.exception("Error during submitted alphas synchronization.")
            db.rollback()
        finally:
            db.close()

    async def process_pending_correlations(self):
        """
        Main entry point called by the orchestrator loop.
        First, polls statuses of pending alphas.
        Second, checks for any previously passed alphas and submits them (if daily cap allows).
        """
        # 1. Poll statuses (throttled internally to 5 mins)
        await self.poll_pending_correlation_statuses()
        
        # 2. Process any PASSED_CORRELATION alphas that couldn't be submitted yet
        db = SessionLocal()
        try:
            await self.submit_ready_alphas_with_db(db)
        except Exception as e:
            logger.exception("Error during processing of ready alphas.")
        finally:
            db.close()