#!/usr/bin/env python3
"""
Hierarchical Predictor Service
==============================

Production-ready service that wraps the 3-level hierarchical ML pipeline.
Provides async interface, caching, health checks, and integration with Data Service.

Architecture:
    Level 0: Attention Model - Determines WHEN to pay attention
    Level 1: Base Models - Symbol/timeframe specific predictions
    Level 2: Metamodel - Synthesizes 5m and 15m predictions

Author: ML Pipeline
Version: 1.0.0
Created: 2026-01-07
"""

import asyncio
import os
from datetime import datetime, timedelta
from typing import Optional, List, Dict, Any, Tuple
from dataclasses import dataclass, asdict
from enum import Enum
from pathlib import Path
import pandas as pd
import numpy as np
from loguru import logger
from cachetools import TTLCache

# Pipeline import
import sys
sys.path.insert(0, str(Path(__file__).parent.parent))
from pipelines.hierarchical_pipeline import (
    HierarchicalPipeline,
    PipelineConfig,
    PredictionResult
)

# Data imports
try:
    from ..data.data_service_client import (
        DataServiceManager,
        DataServiceClient,
        Timeframe
    )
    HAS_DATA_SERVICE = True
except ImportError:
    HAS_DATA_SERVICE = False
    logger.warning("Data service client not available, using local data loading")


class TradeDirection(Enum):
    """Trade direction enum"""
    LONG = "long"
    SHORT = "short"
    NEUTRAL = "neutral"


class SignalQuality(Enum):
    """Signal quality classification"""
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"
    SKIP = "skip"


@dataclass
class HierarchicalSignal:
    """Complete hierarchical trading signal"""

    signal_id: str
    symbol: str
    timestamp: datetime

    # Direction based on predictions
    direction: TradeDirection
    direction_confidence: float

    # Price levels
    entry_price: float
    stop_loss: float
    take_profit: float
    risk_reward_ratio: float

    # Hierarchical predictions
    attention_score: float      # Average of 5m and 15m
    attention_quality: str      # 'high', 'medium', 'low'
    delta_high: float           # Final metamodel prediction
    delta_low: float            # Final metamodel prediction
    confidence: float           # Metamodel confidence

    # Signal quality
    should_trade: bool
    quality: SignalQuality

    # Raw predictions for transparency
    raw_predictions: Dict[str, Any]

    # Validity
    valid_until: datetime

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary"""
        return {
            'signal_id': self.signal_id,
            'symbol': self.symbol,
            'timestamp': self.timestamp.isoformat() if self.timestamp else None,
            'direction': self.direction.value,
            'direction_confidence': self.direction_confidence,
            'entry_price': self.entry_price,
            'stop_loss': self.stop_loss,
            'take_profit': self.take_profit,
            'risk_reward_ratio': self.risk_reward_ratio,
            'attention_score': self.attention_score,
            'attention_quality': self.attention_quality,
            'delta_high': self.delta_high,
            'delta_low': self.delta_low,
            'confidence': self.confidence,
            'should_trade': self.should_trade,
            'quality': self.quality.value,
            'raw_predictions': self.raw_predictions,
            'valid_until': self.valid_until.isoformat() if self.valid_until else None
        }


@dataclass
class ServiceHealth:
    """Service health status"""
    healthy: bool
    models_loaded: Dict[str, bool]
    last_prediction_time: Optional[datetime]
    predictions_count: int
    errors_count: int
    avg_latency_ms: float


class HierarchicalPredictorService:
    """
    Production service for hierarchical ML predictions.

    Features:
    - Async interface for non-blocking predictions
    - In-memory caching with TTL
    - Health monitoring
    - Integration with Data Service
    - Support for batch predictions
    """

    def __init__(
        self,
        models_dir: str = "models",
        data_service_url: Optional[str] = None,
        cache_ttl: int = 60,
        cache_maxsize: int = 100
    ):
        """
        Initialize the service.

        Args:
            models_dir: Directory containing trained models
            data_service_url: URL of Data Service (optional)
            cache_ttl: Cache TTL in seconds
            cache_maxsize: Maximum cache entries
        """
        self.models_dir = Path(models_dir)
        self.data_service_url = data_service_url

        # Pipeline configuration
        self.pipeline_config = PipelineConfig(
            attention_model_path=str(self.models_dir / 'attention'),
            base_model_path=str(self.models_dir / 'symbol_timeframe_models'),
            metamodel_path=str(self.models_dir / 'metamodels'),
            attention_threshold_low=0.8,
            attention_threshold_high=2.0,
            confidence_threshold=0.5
        )

        # Pipeline instance
        self.pipeline: Optional[HierarchicalPipeline] = None

        # Data service
        self.data_manager: Optional[Any] = None

        # Caching
        self._cache = TTLCache(maxsize=cache_maxsize, ttl=cache_ttl)

        # Health metrics
        self._models_loaded: Dict[str, bool] = {}
        self._last_prediction_time: Optional[datetime] = None
        self._predictions_count: int = 0
        self._errors_count: int = 0
        self._latencies: List[float] = []

        # Supported symbols
        self.supported_symbols = ['XAUUSD', 'EURUSD', 'BTCUSD', 'GBPUSD', 'USDJPY']

        logger.info("Initialized HierarchicalPredictorService")
        logger.info(f"  Models dir: {self.models_dir}")
        logger.info(f"  Cache TTL: {cache_ttl}s")

    async def initialize(self) -> bool:
        """
        Initialize service: load models and connect to data service.

        Returns:
            True if initialization successful
        """
        logger.info("Initializing HierarchicalPredictorService...")

        try:
            # Initialize pipeline
            self.pipeline = HierarchicalPipeline(self.pipeline_config)

            # Load models for supported symbols
            for symbol in self.supported_symbols:
                try:
                    loaded = self.pipeline.load_models(symbol)
                    self._models_loaded[symbol] = loaded
                    if loaded:
                        logger.info(f"  Loaded models for {symbol}")
                    else:
                        logger.warning(f"  Failed to load models for {symbol}")
                except Exception as e:
                    logger.error(f"  Error loading models for {symbol}: {e}")
                    self._models_loaded[symbol] = False

            # Initialize data service if available
            if HAS_DATA_SERVICE and self.data_service_url:
                try:
                    client = DataServiceClient(base_url=self.data_service_url)
                    self.data_manager = DataServiceManager(client)
                    logger.info("  Connected to Data Service")
                except Exception as e:
                    logger.warning(f"  Data Service connection failed: {e}")

            logger.info("HierarchicalPredictorService initialized")
            return True

        except Exception as e:
            logger.error(f"Initialization failed: {e}")
            return False

    async def predict(
        self,
        symbol: str,
        df_5m: Optional[pd.DataFrame] = None,
        df_15m: Optional[pd.DataFrame] = None,
        use_cache: bool = True
    ) -> PredictionResult:
        """
        Generate hierarchical prediction for a symbol.

        Args:
            symbol: Trading symbol (e.g., 'XAUUSD')
            df_5m: 5-minute OHLCV data (optional if data service available)
            df_15m: 15-minute OHLCV data (optional if data service available)
            use_cache: Whether to use cached predictions

        Returns:
            PredictionResult with all hierarchical outputs
        """
        start_time = datetime.now()

        # Check cache
        cache_key = f"{symbol}_{datetime.now().strftime('%Y%m%d_%H%M')}"
        if use_cache and cache_key in self._cache:
            logger.debug(f"Cache hit for {symbol}")
            return self._cache[cache_key]

        try:
            # Get data if not provided
            if df_5m is None or df_15m is None:
                df_5m, df_15m = await self._fetch_market_data(symbol)

            # Run pipeline
            result = self.pipeline.predict(df_5m, df_15m, symbol)

            # Update metrics
            self._predictions_count += 1
            self._last_prediction_time = datetime.now()
            latency = (datetime.now() - start_time).total_seconds() * 1000
            self._latencies.append(latency)
            if len(self._latencies) > 100:
                self._latencies.pop(0)

            # Cache result
            if use_cache:
                self._cache[cache_key] = result

            logger.debug(f"Prediction for {symbol} completed in {latency:.0f}ms")
            return result

        except Exception as e:
            self._errors_count += 1
            logger.error(f"Prediction failed for {symbol}: {e}")
            raise

    async def generate_signal(
        self,
        symbol: str,
        df_5m: Optional[pd.DataFrame] = None,
        df_15m: Optional[pd.DataFrame] = None,
        risk_reward: float = 2.0,
        use_cache: bool = True
    ) -> HierarchicalSignal:
        """
        Generate complete trading signal with entry/SL/TP.

        Args:
            symbol: Trading symbol
            df_5m: 5-minute OHLCV data
            df_15m: 15-minute OHLCV data
            risk_reward: Risk/reward ratio for TP calculation
            use_cache: Whether to use cached predictions

        Returns:
            HierarchicalSignal with complete trading information
        """
        import uuid

        # Get prediction
        if df_5m is None or df_15m is None:
            df_5m, df_15m = await self._fetch_market_data(symbol)

        result = await self.predict(symbol, df_5m, df_15m, use_cache)

        # Get current price
        current_price = float(df_15m['close'].iloc[-1])

        # Determine direction
        direction, direction_conf = self._determine_direction(result, df_15m)

        # Calculate levels
        entry, sl, tp = self._calculate_levels(
            current_price,
            direction,
            result,
            risk_reward
        )

        # Calculate attention quality
        avg_attention = (result.attention_score_5m + result.attention_score_15m) / 2
        if avg_attention >= self.pipeline_config.attention_threshold_high:
            attention_quality = 'high'
        elif avg_attention >= self.pipeline_config.attention_threshold_low:
            attention_quality = 'medium'
        else:
            attention_quality = 'low'

        # Map quality
        quality_map = {
            'high': SignalQuality.HIGH,
            'medium': SignalQuality.MEDIUM,
            'low': SignalQuality.LOW,
            'skip': SignalQuality.SKIP
        }

        now = datetime.now()

        return HierarchicalSignal(
            signal_id=f"HSIG-{uuid.uuid4().hex[:8].upper()}",
            symbol=symbol,
            timestamp=result.timestamp,
            direction=direction,
            direction_confidence=direction_conf,
            entry_price=entry,
            stop_loss=sl,
            take_profit=tp,
            risk_reward_ratio=risk_reward,
            attention_score=avg_attention,
            attention_quality=attention_quality,
            delta_high=result.delta_high_final,
            delta_low=result.delta_low_final,
            confidence=result.confidence_proba,
            should_trade=result.should_trade,
            quality=quality_map.get(result.trade_quality, SignalQuality.SKIP),
            raw_predictions=result.to_dict(),
            valid_until=now + timedelta(minutes=15)
        )

    async def predict_batch(
        self,
        symbols: List[str],
        df_dict_5m: Optional[Dict[str, pd.DataFrame]] = None,
        df_dict_15m: Optional[Dict[str, pd.DataFrame]] = None
    ) -> Dict[str, PredictionResult]:
        """
        Generate predictions for multiple symbols.

        Args:
            symbols: List of symbols
            df_dict_5m: Dictionary of 5m DataFrames by symbol
            df_dict_15m: Dictionary of 15m DataFrames by symbol

        Returns:
            Dictionary of PredictionResult by symbol
        """
        results = {}

        for symbol in symbols:
            try:
                df_5m = df_dict_5m.get(symbol) if df_dict_5m else None
                df_15m = df_dict_15m.get(symbol) if df_dict_15m else None

                result = await self.predict(symbol, df_5m, df_15m)
                results[symbol] = result

            except Exception as e:
                logger.error(f"Batch prediction failed for {symbol}: {e}")
                results[symbol] = None

        return results

    async def _fetch_market_data(
        self,
        symbol: str,
        lookback_periods: int = 500
    ) -> Tuple[pd.DataFrame, pd.DataFrame]:
        """
        Fetch market data from Data Service or local database.

        Args:
            symbol: Trading symbol
            lookback_periods: Number of periods to fetch

        Returns:
            Tuple of (df_5m, df_15m)
        """
        if self.data_manager:
            # Use Data Service
            async with self.data_manager.client:
                df_5m = await self.data_manager.get_ml_features_data(
                    symbol=symbol,
                    timeframe=Timeframe('5m'),
                    lookback_periods=lookback_periods * 3
                )
                df_15m = await self.data_manager.get_ml_features_data(
                    symbol=symbol,
                    timeframe=Timeframe('15m'),
                    lookback_periods=lookback_periods
                )
            return df_5m, df_15m

        # Fallback to local MySQL
        return await self._fetch_from_mysql(symbol, lookback_periods)

    async def _fetch_from_mysql(
        self,
        symbol: str,
        lookback_periods: int = 500
    ) -> Tuple[pd.DataFrame, pd.DataFrame]:
        """Fetch data from local MySQL database."""
        import mysql.connector

        db_config = {
            'host': os.environ.get('DB_HOST', 'localhost'),
            'user': os.environ.get('DB_USER', 'orbiquantia'),
            'password': os.environ.get('DB_PASSWORD', 'orbiquantia_dev_2025'),
            'database': os.environ.get('DB_NAME', 'orbiquantia_platform'),
            'port': int(os.environ.get('DB_PORT', 5433))
        }

        # Map symbol to ticker
        ticker_map = {
            'XAUUSD': 'C:XAUUSD',
            'EURUSD': 'C:EURUSD',
            'GBPUSD': 'C:GBPUSD',
            'USDJPY': 'C:USDJPY',
            'BTCUSD': 'X:BTCUSD'
        }
        ticker = ticker_map.get(symbol, f'C:{symbol}')

        try:
            conn = mysql.connector.connect(**db_config)

            # Fetch 5m data
            query_5m = f"""
                SELECT time_start as timestamp, open, high, low, close, volume
                FROM tickers_agg_data
                WHERE ticker = %s AND timeframe = '5'
                ORDER BY time_start DESC
                LIMIT {lookback_periods * 3}
            """
            df_5m = pd.read_sql(query_5m, conn, params=[ticker])
            df_5m['timestamp'] = pd.to_datetime(df_5m['timestamp'])
            df_5m.set_index('timestamp', inplace=True)
            df_5m.sort_index(inplace=True)

            # Fetch 15m data
            query_15m = f"""
                SELECT time_start as timestamp, open, high, low, close, volume
                FROM tickers_agg_data
                WHERE ticker = %s AND timeframe = '15'
                ORDER BY time_start DESC
                LIMIT {lookback_periods}
            """
            df_15m = pd.read_sql(query_15m, conn, params=[ticker])
            df_15m['timestamp'] = pd.to_datetime(df_15m['timestamp'])
            df_15m.set_index('timestamp', inplace=True)
            df_15m.sort_index(inplace=True)

            conn.close()

            return df_5m, df_15m

        except Exception as e:
            logger.error(f"Failed to fetch data from MySQL: {e}")
            raise

    def _determine_direction(
        self,
        result: PredictionResult,
        df_15m: pd.DataFrame
    ) -> Tuple[TradeDirection, float]:
        """
        Determine trade direction from predictions.

        Args:
            result: Pipeline prediction result
            df_15m: 15-minute OHLCV data

        Returns:
            Tuple of (direction, confidence)
        """
        # Compare predicted ranges
        high_pred = result.delta_high_final
        low_pred = result.delta_low_final

        # Calculate bias
        if high_pred > low_pred * 1.2:
            # Bullish bias - expect more upside
            direction = TradeDirection.LONG
            confidence = min(0.9, high_pred / (high_pred + low_pred + 0.01))
        elif low_pred > high_pred * 1.2:
            # Bearish bias - expect more downside
            direction = TradeDirection.SHORT
            confidence = min(0.9, low_pred / (high_pred + low_pred + 0.01))
        else:
            # Neutral - use recent momentum
            recent_return = (df_15m['close'].iloc[-1] / df_15m['close'].iloc[-5] - 1)
            if recent_return > 0.001:
                direction = TradeDirection.LONG
                confidence = 0.55
            elif recent_return < -0.001:
                direction = TradeDirection.SHORT
                confidence = 0.55
            else:
                direction = TradeDirection.NEUTRAL
                confidence = 0.5

        # Adjust confidence based on attention
        avg_attention = (result.attention_score_5m + result.attention_score_15m) / 2
        if avg_attention < self.pipeline_config.attention_threshold_low:
            confidence *= 0.7  # Reduce confidence in low attention
        elif avg_attention >= self.pipeline_config.attention_threshold_high:
            confidence = min(0.95, confidence * 1.1)  # Boost in high attention

        # Adjust based on metamodel confidence
        if not result.confidence:
            confidence *= 0.8

        return direction, round(confidence, 3)

    def _calculate_levels(
        self,
        current_price: float,
        direction: TradeDirection,
        result: PredictionResult,
        risk_reward: float
    ) -> Tuple[float, float, float]:
        """
        Calculate entry, stop loss, and take profit levels.

        Args:
            current_price: Current market price
            direction: Trade direction
            result: Pipeline prediction result
            risk_reward: Risk/reward ratio

        Returns:
            Tuple of (entry, stop_loss, take_profit)
        """
        # Use predicted deltas
        delta_high = result.delta_high_final
        delta_low = result.delta_low_final

        if direction == TradeDirection.LONG:
            entry = current_price
            sl = current_price - delta_low
            risk = current_price - sl
            tp = current_price + (risk * risk_reward)
        elif direction == TradeDirection.SHORT:
            entry = current_price
            sl = current_price + delta_high
            risk = sl - current_price
            tp = current_price - (risk * risk_reward)
        else:
            # Neutral - use symmetric levels
            entry = current_price
            sl = current_price - delta_low
            tp = current_price + delta_high

        # Round based on symbol precision (simple heuristic)
        if current_price > 1000:
            # Gold-like pricing
            decimals = 2
        elif current_price > 1:
            # Forex majors
            decimals = 5
        else:
            # Crypto or other
            decimals = 6

        return (
            round(entry, decimals),
            round(sl, decimals),
            round(tp, decimals)
        )

    def get_health(self) -> ServiceHealth:
        """Get service health status."""
        avg_latency = (
            sum(self._latencies) / len(self._latencies)
            if self._latencies else 0
        )

        return ServiceHealth(
            healthy=any(self._models_loaded.values()),
            models_loaded=self._models_loaded.copy(),
            last_prediction_time=self._last_prediction_time,
            predictions_count=self._predictions_count,
            errors_count=self._errors_count,
            avg_latency_ms=round(avg_latency, 2)
        )

    def get_model_info(self, symbol: str) -> Dict[str, Any]:
        """Get information about loaded models for a symbol."""
        if self.pipeline:
            return self.pipeline.get_model_info(symbol)
        return {'symbol': symbol, 'models_loaded': False}

    def clear_cache(self):
        """Clear prediction cache."""
        self._cache.clear()
        logger.info("Prediction cache cleared")


# Singleton instance
_service: Optional[HierarchicalPredictorService] = None


def get_hierarchical_predictor() -> HierarchicalPredictorService:
    """Get or create hierarchical predictor singleton."""
    global _service
    if _service is None:
        _service = HierarchicalPredictorService()
    return _service


async def initialize_hierarchical_predictor() -> HierarchicalPredictorService:
    """Initialize the hierarchical predictor service."""
    service = get_hierarchical_predictor()
    await service.initialize()
    return service


if __name__ == "__main__":
    """Test the service"""
    import asyncio

    async def test_service():
        print("Testing HierarchicalPredictorService...")

        # Initialize
        service = HierarchicalPredictorService(models_dir="models")
        success = await service.initialize()
        print(f"\nInitialization: {'SUCCESS' if success else 'FAILED'}")

        # Check health
        health = service.get_health()
        print(f"\nService Health:")
        print(f"  Healthy: {health.healthy}")
        print(f"  Models loaded: {health.models_loaded}")

        # Create sample data for testing
        import numpy as np

        np.random.seed(42)
        n_5m = 500
        n_15m = n_5m // 3

        dates_5m = pd.date_range('2024-01-01', periods=n_5m, freq='5min')
        dates_15m = pd.date_range('2024-01-01', periods=n_15m, freq='15min')

        price = 2650 + np.cumsum(np.random.randn(n_5m) * 0.5)

        df_5m = pd.DataFrame({
            'open': price,
            'high': price + np.abs(np.random.randn(n_5m)) * 3,
            'low': price - np.abs(np.random.randn(n_5m)) * 3,
            'close': price + np.random.randn(n_5m) * 0.3,
            'volume': np.random.randint(100, 1000, n_5m)
        }, index=dates_5m)

        price_15m = 2650 + np.cumsum(np.random.randn(n_15m) * 0.8)

        df_15m = pd.DataFrame({
            'open': price_15m,
            'high': price_15m + np.abs(np.random.randn(n_15m)) * 5,
            'low': price_15m - np.abs(np.random.randn(n_15m)) * 5,
            'close': price_15m + np.random.randn(n_15m) * 0.5,
            'volume': np.random.randint(300, 3000, n_15m)
        }, index=dates_15m)

        # Test prediction
        if health.models_loaded.get('XAUUSD'):
            print("\nGenerating prediction for XAUUSD...")
            result = await service.predict('XAUUSD', df_5m, df_15m)

            print(f"\nPrediction Result:")
            print(f"  Attention 5m: {result.attention_score_5m:.2f}")
            print(f"  Attention 15m: {result.attention_score_15m:.2f}")
            print(f"  Delta High Final: {result.delta_high_final:.2f}")
            print(f"  Delta Low Final: {result.delta_low_final:.2f}")
            print(f"  Should Trade: {result.should_trade}")
            print(f"  Trade Quality: {result.trade_quality}")

            # Test signal generation
            print("\nGenerating trading signal...")
            signal = await service.generate_signal('XAUUSD', df_5m, df_15m, risk_reward=2.0)

            print(f"\nTrading Signal:")
            print(f"  ID: {signal.signal_id}")
            print(f"  Direction: {signal.direction.value}")
            print(f"  Entry: {signal.entry_price}")
            print(f"  SL: {signal.stop_loss}")
            print(f"  TP: {signal.take_profit}")
            print(f"  R:R: {signal.risk_reward_ratio}")
            print(f"  Quality: {signal.quality.value}")

        # Final health check
        health = service.get_health()
        print(f"\nFinal Health:")
        print(f"  Predictions: {health.predictions_count}")
        print(f"  Errors: {health.errors_count}")
        print(f"  Avg Latency: {health.avg_latency_ms}ms")

        print("\nTest complete!")

    asyncio.run(test_service())
