institutional-trader/backend/python_service/services/alert_service.py

"""
Alert Service
Handles alert stream matching for options flow
"""
import pandas as pd
import asyncpg
from datetime import datetime, timedelta
from typing import Dict, List, Optional
import pytz
import re


class AlertService:
    """Service for matching alerts with flow data"""
    
    def __init__(self, pool: asyncpg.Pool):
        self.pool = pool
        self.ct_tz = pytz.timezone('America/Chicago')
        self.utc_tz = pytz.UTC
    
    def parse_alert_timestamp(self, date_str: str, time_str: str) -> Optional[datetime]:
        """Parse alert timestamp from date and time strings"""
        if pd.isna(date_str) or pd.isna(time_str):
            return None
        
        date_str = str(date_str).strip()
        time_str = str(time_str).strip()
        
        # Try YYYY-MM-DD format
        if re.match(r'^\d{4}-\d{2}-\d{2}$', date_str):
            try:
                # Try with seconds
                full_str = f"{date_str} {time_str}"
                dt = datetime.strptime(full_str, '%Y-%m-%d %I:%M:%S %p')
                return self.ct_tz.localize(dt).astimezone(self.utc_tz)
            except ValueError:
                try:
                    # Try without seconds
                    full_str = f"{date_str} {time_str}"
                    dt = datetime.strptime(full_str, '%Y-%m-%d %I:%M %p')
                    return self.ct_tz.localize(dt).astimezone(self.utc_tz)
                except ValueError:
                    pass
        
        # Try MM/DD/YYYY format
        if re.match(r'^\d{1,2}/\d{1,2}/\d{2,4}$', date_str):
            try:
                full_str = f"{date_str} {time_str}"
                dt = datetime.strptime(full_str, '%m/%d/%Y %I:%M:%S %p')
                return self.ct_tz.localize(dt).astimezone(self.utc_tz)
            except ValueError:
                try:
                    full_str = f"{date_str} {time_str}"
                    dt = datetime.strptime(full_str, '%m/%d/%Y %I:%M %p')
                    return self.ct_tz.localize(dt).astimezone(self.utc_tz)
                except ValueError:
                    pass
        
        return None
    
    async def get_alerts_for_symbols(
        self,
        symbols: List[str],
        time_window_start: datetime,
        time_window_end: datetime
    ) -> pd.DataFrame:
        """Batch fetch alerts for multiple symbols within time window"""
        if not symbols:
            return pd.DataFrame()
        
        async with self.pool.acquire() as conn:
            # Build query with IN clause
            placeholders = ','.join([f'${i+1}' for i in range(len(symbols))])
            query = f"""
                SELECT 
                    "type",
                    "message",
                    "price",
                    "volume",
                    "date",
                    "timestamp",
                    "ticker"
                FROM "AlertStream_monthly"
                WHERE UPPER("ticker") IN ({placeholders})
            """
            
            rows = await conn.fetch(query, *[s.upper() for s in symbols])
            
            if not rows:
                return pd.DataFrame()
            
            # Convert to DataFrame
            df = pd.DataFrame([dict(row) for row in rows])
            
            # Parse timestamps
            df['event_ts_utc'] = df.apply(
                lambda row: self.parse_alert_timestamp(row['date'], row['timestamp']),
                axis=1
            )
            
            # Filter by time window
            df = df[
                (df['event_ts_utc'] >= time_window_start) &
                (df['event_ts_utc'] <= time_window_end)
            ].copy()
            
            return df
    
    async def match_alerts_to_flows(
        self,
        flow_df: pd.DataFrame
    ) -> pd.DataFrame:
        """Match alerts to flow data within ±15 minutes"""
        df = flow_df.copy()
        
        if df.empty:
            return df
        
        # Get unique symbols and time range
        symbols = df['symbol_norm'].unique().tolist()
        min_time = df['flow_ts_utc'].min() - timedelta(minutes=15)
        max_time = df['flow_ts_utc'].max() + timedelta(minutes=15)
        
        # Fetch all relevant alerts
        alerts_df = await self.get_alerts_for_symbols(symbols, min_time, max_time)
        
        if alerts_df.empty:
            df['near_alert_type'] = None
            df['near_alert_msg'] = None
            df['near_alert_price'] = None
            df['near_alert_volume'] = None
            df['catalyst_flag'] = 0
            return df
        
        # Match alerts to flows
        alert_data = []
        
        for idx, flow_row in df.iterrows():
            symbol = flow_row['symbol_norm']
            flow_ts_utc = flow_row['flow_ts_utc']
            
            if pd.isna(flow_ts_utc):
                alert_data.append({
                    'near_alert_type': None,
                    'near_alert_msg': None,
                    'near_alert_price': None,
                    'near_alert_volume': None,
                })
                continue
            
            # Filter alerts for this symbol within ±15 minutes
            symbol_alerts = alerts_df[
                (alerts_df['ticker'].str.upper() == symbol.upper()) &
                (alerts_df['event_ts_utc'] >= flow_ts_utc - timedelta(minutes=15)) &
                (alerts_df['event_ts_utc'] <= flow_ts_utc + timedelta(minutes=15))
            ]
            
            if symbol_alerts.empty:
                alert_data.append({
                    'near_alert_type': None,
                    'near_alert_msg': None,
                    'near_alert_price': None,
                    'near_alert_volume': None,
                })
                continue
            
            # Find nearest alert by time difference
            symbol_alerts['time_diff'] = abs(
                (symbol_alerts['event_ts_utc'] - flow_ts_utc).dt.total_seconds()
            )
            nearest = symbol_alerts.loc[symbol_alerts['time_diff'].idxmin()]
            
            alert_data.append({
                'near_alert_type': nearest['type'],
                'near_alert_msg': nearest['message'],
                'near_alert_price': nearest['price'],
                'near_alert_volume': nearest['volume'],
            })
        
        # Merge alert data
        alert_df = pd.DataFrame(alert_data, index=df.index)
        df = pd.concat([df, alert_df], axis=1)
        
        # Add catalyst flag
        df['catalyst_flag'] = df['near_alert_type'].notna().astype(int)
        
        return df