#!/usr/bin/env python3
"""
ML-LLM Integration
==================

Orchestrates the complete flow from ML prediction to trading decision:
1. Get predictions from metamodel
2. Get predictions from individual strategies
3. Format signals for LLM consumption
4. Query LLM for trading decision
5. Parse LLM response
6. Log everything for feedback and fine-tuning

This is the main integration point between the ML pipeline and LLM agent.

Author: ML-Specialist (NEXUS v4.0)
Version: 1.0.0
Created: 2026-01-25
"""

import pandas as pd
import numpy as np
from datetime import datetime
from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass, field
from loguru import logger

# Local imports
from .prompts.trading_decision import (
    TradingDecisionPrompt,
    MetamodelOutput,
    StrategyOutput,
    MarketContext,
    AccountInfo,
    SYSTEM_PROMPT,
)
from .signal_formatter import SignalFormatter, FormattedSignal
from .decision_parser import DecisionParser, TradingDecision
from .signal_logger import LLMSignalLogger, TradeResult
from .llm_client import LLMClient, LLMClientConfig


# ============================================================
# DATA CLASSES
# ============================================================

@dataclass
class IntegrationConfig:
    """Configuration for ML-LLM integration"""
    # LLM settings
    llm_provider: str = "ollama"
    llm_model: str = "llama3:8b"
    llm_temperature: float = 0.3

    # Decision thresholds
    min_confidence: float = 0.60  # Minimum ML confidence to consider trading
    min_risk_reward: float = 1.5  # Minimum acceptable R:R ratio

    # Risk settings
    default_position_size: float = 1.0  # Default risk % per trade
    max_position_size: float = 2.0  # Maximum risk % per trade

    # Logging
    enable_logging: bool = True

    # Validation
    strict_validation: bool = False


@dataclass
class TradingRecommendation:
    """
    Complete trading recommendation with all context.

    This is the final output of the integration, containing:
    - The ML predictions
    - The LLM decision
    - All context needed for execution
    """
    # Core recommendation
    action: str  # "TRADE", "NO_TRADE", "WAIT"
    direction: Optional[str] = None  # "LONG", "SHORT"

    # Trade parameters (if action == "TRADE")
    entry: Optional[float] = None
    stop_loss: Optional[float] = None
    take_profit: List[float] = field(default_factory=list)
    position_size: float = 1.0

    # ML context
    ml_direction: str = "NEUTRAL"
    ml_confidence: float = 0.0
    ml_magnitude: float = 0.0

    # Strategy alignment
    aligned_strategies: int = 0
    total_strategies: int = 0

    # Reasoning
    reasoning: str = ""

    # Metadata
    symbol: str = ""
    timestamp: str = ""
    signal_id: Optional[str] = None

    # Quality indicators
    is_high_confidence: bool = False
    risk_reward_ratio: Optional[float] = None

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary"""
        return {
            "action": self.action,
            "direction": self.direction,
            "entry": self.entry,
            "stop_loss": self.stop_loss,
            "take_profit": self.take_profit,
            "position_size": self.position_size,
            "ml_context": {
                "direction": self.ml_direction,
                "confidence": self.ml_confidence,
                "magnitude": self.ml_magnitude,
            },
            "strategy_alignment": {
                "aligned": self.aligned_strategies,
                "total": self.total_strategies,
            },
            "reasoning": self.reasoning,
            "metadata": {
                "symbol": self.symbol,
                "timestamp": self.timestamp,
                "signal_id": self.signal_id,
            },
            "quality": {
                "is_high_confidence": self.is_high_confidence,
                "risk_reward_ratio": self.risk_reward_ratio,
            }
        }


# ============================================================
# ML-LLM INTEGRATION CLASS
# ============================================================

class MLLLMIntegration:
    """
    Main integration class between ML pipeline and LLM agent.

    Orchestrates the complete flow:
    1. ML Predict: Get predictions from metamodel and strategies
    2. Format Signal: Convert to LLM-friendly format
    3. LLM Decide: Query LLM for trading decision
    4. Parse Response: Extract structured decision
    5. Log: Save for feedback and fine-tuning

    Usage:
        integration = MLLLMIntegration(metamodel=my_metamodel)

        # Full flow with LLM decision
        recommendation = integration.analyze_and_decide(
            df=price_data,
            symbol="XAUUSD",
            account_info=account
        )

        # Just get recommendation without LLM (ML-only)
        recommendation = integration.get_recommendation(
            df=price_data,
            symbol="XAUUSD"
        )
    """

    def __init__(
        self,
        metamodel: Optional[Any] = None,
        strategy_models: Optional[Dict[str, Any]] = None,
        llm_client: Optional[LLMClient] = None,
        signal_logger: Optional[LLMSignalLogger] = None,
        config: Optional[IntegrationConfig] = None
    ):
        """
        Initialize the ML-LLM integration.

        Args:
            metamodel: Trained AssetMetamodel or similar
            strategy_models: Dict of strategy name -> model (optional)
            llm_client: Pre-configured LLM client (creates default if None)
            signal_logger: Pre-configured logger (creates default if None)
            config: Integration configuration
        """
        self.config = config or IntegrationConfig()
        self.metamodel = metamodel
        self.strategy_models = strategy_models or {}

        # Initialize components
        self.signal_formatter = SignalFormatter()
        self.decision_parser = DecisionParser(
            strict_mode=self.config.strict_validation,
            min_risk_reward=self.config.min_risk_reward,
            max_position_size=self.config.max_position_size
        )
        self.prompt_builder = TradingDecisionPrompt()

        # LLM client
        if llm_client:
            self.llm_client = llm_client
        else:
            self.llm_client = LLMClient(
                provider=self.config.llm_provider,
                model=self.config.llm_model,
                enable_fallback=True
            )

        # Signal logger
        if signal_logger:
            self.signal_logger = signal_logger
        elif self.config.enable_logging:
            try:
                self.signal_logger = LLMSignalLogger()
            except Exception as e:
                logger.warning(f"Failed to initialize signal logger: {e}")
                self.signal_logger = None
        else:
            self.signal_logger = None

        logger.info("MLLLMIntegration initialized")

    def analyze_and_decide(
        self,
        df: pd.DataFrame,
        symbol: str,
        account_info: Optional[Dict[str, Any]] = None,
        strategy_predictions: Optional[List[Dict[str, Any]]] = None,
        current_price: Optional[float] = None,
        atr: Optional[float] = None
    ) -> TradingRecommendation:
        """
        Complete flow: ML predict -> format -> LLM decide -> parse -> log.

        This is the main entry point for getting a trading decision.

        Args:
            df: OHLCV DataFrame with price data
            symbol: Trading symbol (e.g., "XAUUSD")
            account_info: Account/risk information (optional)
            strategy_predictions: Pre-computed strategy predictions (optional)
            current_price: Current price (uses last close if None)
            atr: ATR value (calculates if None)

        Returns:
            TradingRecommendation with complete decision context
        """
        timestamp = datetime.utcnow().isoformat()

        # Get current price and ATR if not provided
        if current_price is None:
            current_price = float(df["close"].iloc[-1])

        if atr is None:
            atr = self._calculate_atr(df)

        # Step 1: Get ML predictions
        ml_prediction = self._get_ml_prediction(df, symbol)

        if ml_prediction is None:
            logger.warning("Failed to get ML prediction")
            return TradingRecommendation(
                action="NO_TRADE",
                reasoning="ML prediction unavailable",
                symbol=symbol,
                timestamp=timestamp
            )

        # Step 2: Get strategy predictions (if not provided)
        if strategy_predictions is None:
            strategy_predictions = self._get_strategy_predictions(df, symbol)

        # Step 3: Format signal for LLM
        formatted_signal = self.signal_formatter.format_for_llm(
            metamodel_prediction=ml_prediction,
            strategy_predictions=strategy_predictions,
            market_context={
                "current_price": current_price,
                "atr": atr,
                "market_phase": ml_prediction.get("market_phase", "unknown"),
                "timestamp": timestamp
            },
            symbol=symbol
        )

        # Add risk context if provided
        if account_info:
            self.signal_formatter.add_risk_context(formatted_signal, account_info)

        # Check minimum confidence threshold
        if formatted_signal.confidence_value < self.config.min_confidence:
            logger.info(f"ML confidence {formatted_signal.confidence_value:.1%} below threshold")
            return TradingRecommendation(
                action="NO_TRADE",
                reasoning=f"ML confidence ({formatted_signal.confidence_value:.1%}) below minimum threshold ({self.config.min_confidence:.1%})",
                ml_direction=formatted_signal.direction,
                ml_confidence=formatted_signal.confidence_value,
                symbol=symbol,
                timestamp=timestamp
            )

        # Step 4: Build prompt and query LLM
        prompt = self._build_prompt(
            symbol=symbol,
            formatted_signal=formatted_signal,
            account_info=account_info or self._default_account_info()
        )

        try:
            llm_response = self.llm_client.get_trading_decision(
                prompt=prompt,
                temperature=self.config.llm_temperature
            )
        except Exception as e:
            logger.error(f"LLM query failed: {e}")
            return TradingRecommendation(
                action="NO_TRADE",
                reasoning=f"LLM query failed: {e}",
                ml_direction=formatted_signal.direction,
                ml_confidence=formatted_signal.confidence_value,
                symbol=symbol,
                timestamp=timestamp
            )

        # Step 5: Parse LLM response
        decision = self.decision_parser.parse_response(llm_response)

        # Step 6: Log signal and decision
        signal_id = None
        if self.signal_logger:
            try:
                signal_id = self.signal_logger.log_signal(
                    signal=formatted_signal,
                    decision=decision,
                    llm_prompt=prompt,
                    llm_response=llm_response
                )
            except Exception as e:
                logger.warning(f"Failed to log signal: {e}")

        # Step 7: Build recommendation
        recommendation = self._build_recommendation(
            decision=decision,
            formatted_signal=formatted_signal,
            symbol=symbol,
            timestamp=timestamp,
            signal_id=signal_id,
            strategy_predictions=strategy_predictions
        )

        return recommendation

    def get_recommendation(
        self,
        df: pd.DataFrame,
        symbol: str,
        current_price: Optional[float] = None,
        atr: Optional[float] = None
    ) -> TradingRecommendation:
        """
        Get ML-only recommendation without LLM decision.

        Useful for quick analysis or when LLM is unavailable.

        Args:
            df: OHLCV DataFrame
            symbol: Trading symbol
            current_price: Current price (optional)
            atr: ATR value (optional)

        Returns:
            TradingRecommendation based on ML signals only
        """
        timestamp = datetime.utcnow().isoformat()

        if current_price is None:
            current_price = float(df["close"].iloc[-1])

        if atr is None:
            atr = self._calculate_atr(df)

        # Get ML prediction
        ml_prediction = self._get_ml_prediction(df, symbol)

        if ml_prediction is None:
            return TradingRecommendation(
                action="NO_TRADE",
                reasoning="ML prediction unavailable",
                symbol=symbol,
                timestamp=timestamp
            )

        # Format signal
        formatted_signal = self.signal_formatter.format_for_llm(
            metamodel_prediction=ml_prediction,
            strategy_predictions=[],
            market_context={
                "current_price": current_price,
                "atr": atr,
                "timestamp": timestamp
            },
            symbol=symbol
        )

        # Simple ML-based decision
        if formatted_signal.confidence_value >= self.config.min_confidence:
            if formatted_signal.direction == "BULLISH":
                action = "TRADE"
                direction = "LONG"
                entry = current_price
                stop_loss = current_price - (atr * 1.5)
                take_profit = [current_price + (atr * 2)]
            elif formatted_signal.direction == "BEARISH":
                action = "TRADE"
                direction = "SHORT"
                entry = current_price
                stop_loss = current_price + (atr * 1.5)
                take_profit = [current_price - (atr * 2)]
            else:
                action = "NO_TRADE"
                direction = None
                entry = None
                stop_loss = None
                take_profit = []
        else:
            action = "NO_TRADE"
            direction = None
            entry = None
            stop_loss = None
            take_profit = []

        return TradingRecommendation(
            action=action,
            direction=direction,
            entry=entry,
            stop_loss=stop_loss,
            take_profit=take_profit,
            position_size=self.config.default_position_size,
            ml_direction=formatted_signal.direction,
            ml_confidence=formatted_signal.confidence_value,
            ml_magnitude=formatted_signal.magnitude_value,
            reasoning=f"ML-only decision: {formatted_signal.direction} with {formatted_signal.confidence_value:.1%} confidence",
            symbol=symbol,
            timestamp=timestamp,
            is_high_confidence=formatted_signal.confidence_value >= 0.7
        )

    def log_trade_result(
        self,
        signal_id: str,
        result: TradeResult
    ) -> bool:
        """
        Log the result of a trade for feedback.

        Args:
            signal_id: The signal_id from the recommendation
            result: TradeResult with outcome data

        Returns:
            True if logged successfully
        """
        if self.signal_logger:
            return self.signal_logger.log_trade_result(signal_id, result)
        return False

    def _get_ml_prediction(
        self,
        df: pd.DataFrame,
        symbol: str
    ) -> Optional[Dict[str, Any]]:
        """Get prediction from metamodel"""
        if self.metamodel is None:
            # Return mock prediction for testing
            logger.warning("No metamodel configured, using mock prediction")
            return {
                "delta_high_final": 0.5,
                "delta_low_final": 0.3,
                "confidence": 1,
                "confidence_proba": 0.65,
                "strategy_weights": {}
            }

        try:
            # This assumes metamodel has a predict method
            # Adjust based on actual metamodel interface
            prediction = self.metamodel.predict(df)

            if hasattr(prediction, "to_dict"):
                return prediction.to_dict()
            elif hasattr(prediction, "delta_high_final"):
                return {
                    "delta_high_final": prediction.delta_high_final,
                    "delta_low_final": prediction.delta_low_final,
                    "confidence": prediction.confidence,
                    "confidence_proba": prediction.confidence_proba
                }
            else:
                return dict(prediction)

        except Exception as e:
            logger.error(f"Metamodel prediction failed: {e}")
            return None

    def _get_strategy_predictions(
        self,
        df: pd.DataFrame,
        symbol: str
    ) -> List[Dict[str, Any]]:
        """Get predictions from individual strategies"""
        predictions = []

        # Default strategy names if no models configured
        default_strategies = ["PVA", "MRD", "VBP", "MTS", "MSA"]

        if not self.strategy_models:
            # Return placeholder predictions
            for name in default_strategies:
                predictions.append({
                    "name": name,
                    "direction": 0,
                    "confidence": 0.5,
                    "weight": 1.0 / len(default_strategies)
                })
            return predictions

        for name, model in self.strategy_models.items():
            try:
                result = model.predict(df)
                predictions.append({
                    "name": name,
                    "direction": getattr(result, "direction", 0),
                    "confidence": getattr(result, "confidence", 0.5),
                    "weight": getattr(result, "weight", 1.0 / len(self.strategy_models))
                })
            except Exception as e:
                logger.warning(f"Strategy {name} prediction failed: {e}")

        return predictions

    def _calculate_atr(self, df: pd.DataFrame, period: int = 14) -> float:
        """Calculate ATR from DataFrame"""
        if len(df) < period:
            return 0.0

        high = df["high"]
        low = df["low"]
        close = df["close"]

        tr1 = high - low
        tr2 = abs(high - close.shift(1))
        tr3 = abs(low - close.shift(1))

        tr = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)
        atr = tr.rolling(window=period).mean().iloc[-1]

        return float(atr) if not pd.isna(atr) else 0.0

    def _build_prompt(
        self,
        symbol: str,
        formatted_signal: FormattedSignal,
        account_info: Dict[str, Any]
    ) -> str:
        """Build the complete prompt for the LLM"""
        # Convert formatted signal to prompt builder format
        metamodel_output = MetamodelOutput(
            direction=formatted_signal.direction,
            magnitude=formatted_signal.magnitude_value,
            confidence=formatted_signal.confidence_value,
            strategy_weights=formatted_signal.strategy_weights,
            delta_high=formatted_signal.delta_high,
            delta_low=formatted_signal.delta_low
        )

        strategy_outputs = [
            StrategyOutput(
                name=s.name,
                direction=s.direction,
                confidence=s.confidence_value,
                weight=s.weight
            )
            for s in formatted_signal.strategies
        ]

        market_context = MarketContext(
            current_price=formatted_signal.current_price,
            atr=formatted_signal.atr,
            volatility_regime=formatted_signal.volatility_regime,
            market_phase=formatted_signal.market_phase,
            trading_session=formatted_signal.trading_session,
            timestamp=formatted_signal.timestamp,
            timeframe=formatted_signal.timeframe
        )

        account = AccountInfo(
            balance=account_info.get("balance", 1000.0),
            available_margin=account_info.get("available_margin", 1000.0),
            current_drawdown=account_info.get("current_drawdown", 0.0),
            open_positions=account_info.get("open_positions", 0),
            max_risk_per_trade=account_info.get("max_risk_per_trade", 2.0),
            daily_pnl=account_info.get("daily_pnl", 0.0),
            daily_pnl_pct=account_info.get("daily_pnl_pct", 0.0)
        )

        return self.prompt_builder.build_full_prompt(
            symbol=symbol,
            metamodel_output=metamodel_output,
            strategy_outputs=strategy_outputs,
            market_context=market_context,
            account_info=account
        )

    def _build_recommendation(
        self,
        decision: TradingDecision,
        formatted_signal: FormattedSignal,
        symbol: str,
        timestamp: str,
        signal_id: Optional[str],
        strategy_predictions: List[Dict[str, Any]]
    ) -> TradingRecommendation:
        """Build the final recommendation from all components"""
        # Count aligned strategies
        aligned = sum(
            1 for s in formatted_signal.strategies
            if s.direction == formatted_signal.direction and s.confidence_value > 0.5
        )

        return TradingRecommendation(
            action=decision.action,
            direction=decision.direction,
            entry=decision.entry if isinstance(decision.entry, (int, float)) else None,
            stop_loss=decision.stop_loss,
            take_profit=decision.take_profit,
            position_size=decision.position_size,
            ml_direction=formatted_signal.direction,
            ml_confidence=formatted_signal.confidence_value,
            ml_magnitude=formatted_signal.magnitude_value,
            aligned_strategies=aligned,
            total_strategies=len(formatted_signal.strategies),
            reasoning=decision.reasoning,
            symbol=symbol,
            timestamp=timestamp,
            signal_id=signal_id,
            is_high_confidence=formatted_signal.confidence_value >= 0.7,
            risk_reward_ratio=decision.get_risk_reward_ratio()
        )

    def _default_account_info(self) -> Dict[str, Any]:
        """Return default account info"""
        return {
            "balance": 1000.0,
            "available_margin": 1000.0,
            "current_drawdown": 0.0,
            "open_positions": 0,
            "max_risk_per_trade": 2.0,
            "daily_pnl": 0.0,
            "daily_pnl_pct": 0.0
        }


# ============================================================
# CONVENIENCE FUNCTIONS
# ============================================================

def create_integration(
    llm_provider: str = "ollama",
    llm_model: str = "llama3:8b",
    metamodel: Optional[Any] = None
) -> MLLLMIntegration:
    """
    Create a configured ML-LLM integration instance.

    Args:
        llm_provider: LLM provider ("ollama" or "claude")
        llm_model: Model name
        metamodel: Trained metamodel (optional)

    Returns:
        Configured MLLLMIntegration instance
    """
    config = IntegrationConfig(
        llm_provider=llm_provider,
        llm_model=llm_model
    )

    return MLLLMIntegration(
        metamodel=metamodel,
        config=config
    )


# ============================================================
# EXPORTS
# ============================================================

__all__ = [
    'MLLLMIntegration',
    'IntegrationConfig',
    'TradingRecommendation',
    'create_integration',
]
