live_data_fetcher.py

"""
Live Data Fetcher for TSLA Trading Bot
Fetches real-time 5-minute OHLCV data from Polygon API
"""

import os
import time
import requests
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import pytz
from typing import Optional, Dict, List
import logging
from dataclasses import dataclass

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

@dataclass
class BarData:
    """Single bar of OHLCV data"""
    timestamp: datetime
    open: float
    high: float
    low: float
    close: float
    volume: int

class PolygonDataFetcher:
    """Fetches live and historical data from Polygon API"""
    
    def __init__(self, api_key: str):
        self.api_key = api_key
        self.base_url = "https://api.polygon.io"
        self.session = requests.Session()
        self.session.headers.update({
            'Authorization': f'Bearer {api_key}'
        })
        
    def get_historical_bars(self, symbol: str, timespan: str = "minute",
                           multiplier: int = 5, limit: int = 50000,
                           from_date: Optional[str] = None) -> List[BarData]:
        """
        Fetch historical bars for initial data loading
        
        Args:
            symbol: Stock symbol (e.g., 'TSLA')
            timespan: 'minute', 'hour', 'day'
            multiplier: Number of timespans (e.g., 5 for 5-minute bars)
            limit: Maximum number of bars to fetch
            from_date: Start date in YYYY-MM-DD format
        """
        # Use US/Eastern timezone to align with market hours
        now = datetime.now(pytz.timezone('US/Eastern'))
        if not from_date:
            # Default to 30 days ago from current time
            from_date = (now - timedelta(days=30)).strftime('%Y-%m-%d')

        to_date = now.strftime('%Y-%m-%d')

        logger.debug(
            f"Requesting historical bars for {symbol} from {from_date} to {to_date}"
        )

        url = (
            f"{self.base_url}/v2/aggs/ticker/{symbol}/range/"
            f"{multiplier}/{timespan}/{from_date}/{to_date}"
        )
        
        params = {
            'adjusted': 'true',
            'sort': 'asc',
            'limit': 50000
        }
        
        try:
            response = self.session.get(url, params=params)
            response.raise_for_status()
            
            data = response.json()
            
            if data.get('status') != 'OK':
                logger.error(f"Polygon API error: {data}")
                return []
            
            bars = []
            results = data.get('results', [])
            if not results:
                logger.warning(
                    f"No historical bars returned for {symbol} between {from_date} and {to_date}"
                )
                return []
            for result in results:
                # Convert timestamp from milliseconds to datetime
                timestamp = datetime.fromtimestamp(result['t'] / 1000, tz=pytz.timezone('US/Eastern'))
                
                bar = BarData(
                    timestamp=timestamp,
                    open=result['o'],
                    high=result['h'],
                    low=result['l'],
                    close=result['c'],
                    volume=result['v']
                )
                bars.append(bar)
            
            logger.info(f"Fetched {len(bars)} historical bars for {symbol}")
            return bars
            
        except Exception as e:
            logger.error(f"Error fetching historical data: {e}")
            return []
    
    def get_latest_bar(self, symbol: str) -> Optional[BarData]:
        """
        Fetch the most recent completed bar
        """
        # Get the last 2 bars to ensure we have the most recent completed one
        bars = self.get_historical_bars(symbol, limit=2)
        
        if not bars:
            return None
        
        # Return the most recent bar
        return bars[-1]
    
    def get_current_price(self, symbol: str) -> Optional[float]:
        """
        Get current real-time price
        """
        url = f"{self.base_url}/v2/last/trade/{symbol}"
        
        try:
            response = self.session.get(url)
            response.raise_for_status()
            
            data = response.json()
            
            if data.get('status') != 'OK':
                logger.error(f"Error getting current price: {data}")
                return None
            
            return data['results']['p']  # price
            
        except Exception as e:
            logger.error(f"Error fetching current price: {e}")
            return None

class DataManager:
    """Manages historical and live data for the trading bot"""
    
    def __init__(self, polygon_api_key: str, symbol: str = "TSLA"):
        self.symbol = symbol
        self.fetcher = PolygonDataFetcher(polygon_api_key)
        self.historical_data = pd.DataFrame()
        self.current_price = None
        self.last_update = None
        
    def initialize_historical_data(self, days_back: int = 30) -> bool:
        """
        Load initial historical data for indicator calculations
        """
        logger.info(f"Initializing historical data for {self.symbol}...")
        
        from_date = (datetime.now() - timedelta(days=days_back)).strftime('%Y-%m-%d')
        bars = self.fetcher.get_historical_bars(self.symbol, from_date=from_date)
        
        if not bars:
            logger.error("Failed to fetch historical data")
            return False
        
        # Convert to DataFrame
        data = []
        for bar in bars:
            data.append({
                'Datetime': bar.timestamp,
                'Open': bar.open,
                'High': bar.high,
                'Low': bar.low,
                'Close': bar.close,
                'Volume': bar.volume
            })
        
        self.historical_data = pd.DataFrame(data)
        self.historical_data.set_index('Datetime', inplace=True)
        self.historical_data.sort_index(inplace=True)
        
        logger.info(f"Loaded {len(self.historical_data)} historical bars")
        logger.info(f"Data range: {self.historical_data.index[0]} to {self.historical_data.index[-1]}")
        
        return True
    
    def update_current_price(self) -> bool:
        """
        Update current real-time price
        """
        price = self.fetcher.get_current_price(self.symbol)
        if price:
            self.current_price = price
            self.last_update = datetime.now()
            return True
        return False
    
    def get_latest_complete_data(self) -> pd.DataFrame:
        """
        Get the most recent complete dataset for analysis
        """
        # Check if we need to fetch a new bar
        latest_bar = self.fetcher.get_latest_bar(self.symbol)
        
        if latest_bar and not self.historical_data.empty:
            latest_timestamp = latest_bar.timestamp
            
            # Check if this is a new bar we haven't seen
            if latest_timestamp not in self.historical_data.index:
                # Add new bar to historical data
                new_row = pd.DataFrame({
                    'Open': [latest_bar.open],
                    'High': [latest_bar.high],
                    'Low': [latest_bar.low],
                    'Close': [latest_bar.close],
                    'Volume': [latest_bar.volume]
                }, index=[latest_timestamp])
                
                self.historical_data = pd.concat([self.historical_data, new_row])
                self.historical_data.sort_index(inplace=True)
                
                # Keep only last 1000 bars to manage memory
                if len(self.historical_data) > 1000:
                    self.historical_data = self.historical_data.tail(1000)
                
                logger.info(f"Added new bar: {latest_timestamp} - Close: ${latest_bar.close:.2f}")
        
        return self.historical_data.copy()
    
    def get_market_status(self) -> Dict[str, any]:
        """
        Get current market status and data freshness
        """
        now = datetime.now(pytz.timezone('US/Eastern'))
        
        # Market hours: 9:30 AM - 4:00 PM ET
        market_open = now.replace(hour=9, minute=30, second=0, microsecond=0)
        market_close = now.replace(hour=16, minute=0, second=0, microsecond=0)
        
        is_market_hours = market_open <= now <= market_close and now.weekday() < 5
        
        return {
            'current_time': now,
            'is_market_hours': is_market_hours,
            'current_price': self.current_price,
            'last_price_update': self.last_update,
            'data_bars_count': len(self.historical_data),
            'latest_bar_time': self.historical_data.index[-1] if not self.historical_data.empty else None
        }

def test_data_fetcher():
    """Test the data fetcher functionality"""
    # Use configured Polygon API key
    api_key = os.getenv('POLYGON_API_KEY', 'JlAQap9qJ8F8VrfChiPmYpticVo6SMPO')
    
    # Initialize data manager
    data_manager = DataManager(api_key, "TSLA")
    
    # Test historical data loading
    if data_manager.initialize_historical_data(days_back=7):
        print("✓ Historical data loaded successfully")
        print(f"Data shape: {data_manager.historical_data.shape}")
        print(f"Latest bar: {data_manager.historical_data.tail(1)}")
    else:
        print("✗ Failed to load historical data")
        return
    
    # Test current price fetching
    if data_manager.update_current_price():
        print(f"✓ Current TSLA price: ${data_manager.current_price:.2f}")
    else:
        print("✗ Failed to fetch current price")
    
    # Test market status
    status = data_manager.get_market_status()
    print(f"✓ Market status: {status}")

if __name__ == "__main__":
    test_data_fetcher()