reasoning/bayesian.py

"""Advanced Bayesian reasoning for probabilistic analysis."""

import logging
from typing import Dict, Any, List, Optional, Set, Union, Type, Tuple
import json
from dataclasses import dataclass, field
from enum import Enum
from datetime import datetime
import numpy as np
from collections import defaultdict

from .base import ReasoningStrategy, StrategyResult

@dataclass
class BayesianHypothesis:
    """Bayesian hypothesis with probabilities."""
    name: str
    prior: float
    likelihood: float
    posterior: float = 0.0
    evidence: List[Dict[str, Any]] = field(default_factory=list)

class BayesianStrategy(ReasoningStrategy):
    """Advanced Bayesian reasoning that:
    1. Generates hypotheses
    2. Calculates prior probabilities
    3. Updates with evidence
    4. Computes posteriors
    5. Provides probabilistic analysis
    """
    
    def __init__(self, config: Optional[Dict[str, Any]] = None):
        """Initialize Bayesian reasoning."""
        super().__init__()
        self.config = config or {}
        
        # Configure Bayesian parameters
        self.prior_weight = self.config.get('prior_weight', 0.3)
        self.evidence_threshold = self.config.get('evidence_threshold', 0.1)
        self.min_likelihood = self.config.get('min_likelihood', 0.01)
        
        # Initialize hypothesis storage
        self.hypotheses: List[BayesianHypothesis] = []
    
    async def reason(
        self,
        query: str,
        context: Dict[str, Any]
    ) -> StrategyResult:
        """
        Apply Bayesian reasoning to analyze probabilities and update beliefs.
        
        Args:
            query: The input query to reason about
            context: Additional context and parameters
            
        Returns:
            StrategyResult containing reasoning results and confidence scores
        """
        try:
            # Generate initial hypotheses
            self.hypotheses = await self._generate_hypotheses(query, context)
            
            # Calculate prior probabilities
            priors = await self._calculate_priors(self.hypotheses, context)
            
            # Update with evidence
            posteriors = await self._update_with_evidence(self.hypotheses, priors, context)
            
            # Generate analysis
            analysis = await self._generate_analysis(posteriors, context)
            
            # Format results
            answer = self._format_analysis(analysis)
            confidence = self._calculate_confidence(posteriors)
            
            return StrategyResult(
                strategy_type="bayesian",
                success=True,
                answer=answer,
                confidence=confidence,
                reasoning_trace=[{
                    "step": "bayesian_analysis",
                    "hypotheses": [h.__dict__ for h in self.hypotheses],
                    "priors": priors,
                    "posteriors": posteriors,
                    "analysis": analysis,
                    "timestamp": datetime.now().isoformat()
                }],
                metadata={
                    "num_hypotheses": len(self.hypotheses),
                    "max_posterior": max(posteriors.values()) if posteriors else 0.0,
                    "config": self.config
                },
                performance_metrics={
                    "prior_weight": self.prior_weight,
                    "evidence_threshold": self.evidence_threshold,
                    "min_likelihood": self.min_likelihood
                }
            )
            
        except Exception as e:
            logging.error(f"Bayesian reasoning error: {str(e)}")
            return StrategyResult(
                strategy_type="bayesian",
                success=False,
                answer=None,
                confidence=0.0,
                reasoning_trace=[{
                    "step": "error",
                    "error": str(e),
                    "timestamp": datetime.now().isoformat()
                }],
                metadata={"error": str(e)},
                performance_metrics={}
            )
    
    async def _generate_hypotheses(
        self,
        query: str,
        context: Dict[str, Any]
    ) -> List[BayesianHypothesis]:
        """Generate plausible hypotheses."""
        # Extract key terms for hypothesis generation
        terms = self._extract_factors(query, set())
        
        # Generate alternative hypotheses
        alternatives = self._generate_alternative_factors(terms)
        
        # Create hypothesis objects
        hypotheses = []
        for name, prior in alternatives.items():
            hypotheses.append(BayesianHypothesis(
                name=name,
                prior=prior,
                likelihood=1.0  # Initial likelihood
            ))
        
        return hypotheses
    
    async def _calculate_priors(
        self,
        hypotheses: List[BayesianHypothesis],
        context: Dict[str, Any]
    ) -> Dict[str, float]:
        """Calculate prior probabilities."""
        priors = {}
        total_prior = sum(h.prior for h in hypotheses)
        
        if total_prior > 0:
            # Normalize priors
            for h in hypotheses:
                priors[h.name] = h.prior / total_prior
        else:
            # Equal priors if no information
            prior = 1.0 / len(hypotheses)
            for h in hypotheses:
                priors[h.name] = prior
        
        return priors
    
    async def _update_with_evidence(
        self,
        hypotheses: List[BayesianHypothesis],
        priors: Dict[str, float],
        context: Dict[str, Any]
    ) -> Dict[str, float]:
        """Update probabilities with evidence."""
        posteriors = priors.copy()
        
        # Get evidence from context
        evidence = context.get('evidence', [])
        
        for e in evidence:
            # Calculate likelihoods
            likelihoods = {}
            total_likelihood = 0.0
            
            for h in hypotheses:
                likelihood = await self._calculate_likelihood(h, e)
                likelihoods[h.name] = max(likelihood, self.min_likelihood)
                total_likelihood += likelihood * posteriors[h.name]
            
            # Update posteriors using Bayes' rule
            if total_likelihood > 0:
                for h in hypotheses:
                    posteriors[h.name] = (
                        likelihoods[h.name] * posteriors[h.name] / total_likelihood
                    )
        
        return posteriors
    
    async def _calculate_likelihood(
        self,
        hypothesis: BayesianHypothesis,
        evidence: Dict[str, Any]
    ) -> float:
        """Calculate likelihood of evidence given hypothesis."""
        # Simple likelihood calculation
        # Could be enhanced with more sophisticated methods
        base_likelihood = 0.5
        
        # Adjust based on evidence strength
        strength = evidence.get('strength', 0.0)
        likelihood = base_likelihood * (1 + strength)
        
        return min(1.0, max(self.min_likelihood, likelihood))
    
    async def _generate_analysis(
        self,
        posteriors: Dict[str, float],
        context: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Generate probabilistic analysis."""
        analysis = {
            'top_hypothesis': max(posteriors.items(), key=lambda x: x[1]),
            'confidence': self._calculate_confidence(posteriors),
            'distribution': posteriors,
            'summary': []
        }
        
        # Generate summary points
        for name, prob in sorted(posteriors.items(), key=lambda x: x[1], reverse=True):
            analysis['summary'].append({
                'hypothesis': name,
                'probability': prob,
                'strength': 'strong' if prob > 0.7 else 'moderate' if prob > 0.3 else 'weak'
            })
        
        return analysis
    
    def _format_analysis(self, analysis: Dict[str, Any]) -> str:
        """Format analysis into readable text."""
        top_hyp, top_prob = analysis['top_hypothesis']
        
        text = [
            f"Based on Bayesian analysis:",
            f"- Most likely hypothesis: {top_hyp} (probability: {top_prob:.2f})",
            "\nProbability distribution:"
        ]
        
        for item in analysis['summary']:
            text.append(
                f"- {item['hypothesis']}: {item['probability']:.2f} "
                f"({item['strength']} evidence)"
            )
        
        return "\n".join(text)
    
    def _calculate_confidence(self, posteriors: Dict[str, float]) -> float:
        """Calculate overall confidence score."""
        if not posteriors:
            return 0.0
        
        # Get top two probabilities
        probs = sorted(posteriors.values(), reverse=True)
        top_prob = probs[0]
        
        if len(probs) > 1:
            # Consider the gap between top hypotheses
            second_prob = probs[1]
            margin = top_prob - second_prob
            
            # Confidence increases with both probability and margin
            confidence = (top_prob + margin) / 2
        else:
            confidence = top_prob
        
        return min(1.0, max(0.0, confidence))
    
    def _extract_factors(self, text: str, terms: Set[str]) -> Set[str]:
        """Extract relevant factors from text."""
        # Simple word-based extraction
        # Could be enhanced with NLP techniques
        words = text.lower().split()
        return set(words).union(terms)
    
    def _generate_alternative_factors(self, terms: Set[str]) -> Dict[str, float]:
        """Generate factors for alternative hypothesis."""
        # Simple alternative generation
        # Could be enhanced with domain knowledge
        alternatives = {
            'primary': 0.6,
            'alternative': 0.3,
            'null': 0.1
        }
        
        return alternatives