"""
Outcome Engine
Self-learning system where every recommendation becomes an experiment
Closed feedback loop: Observation → Analysis → Recommendation → Action → New Observation → Change → AI learns
"""

from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
import json


class OutcomeStatus(str, Enum):
    """Status of an intervention outcome"""
    SUCCESS = "success"
    PARTIAL = "partial"
    UNDERPERFORMED = "underperformed"
    FAILED = "failed"
    INCONCLUSIVE = "inconclusive"


@dataclass
class InterventionOutcome:
    """Outcome of an intervention"""
    intervention_id: str
    recommendation_id: str
    location_id: str
    action_type: str
    baseline_state: Dict[str, float]
    post_state: Dict[str, float]
    expected_change: Dict[str, float]
    actual_change: Dict[str, float]
    variance: Dict[str, float]
    status: OutcomeStatus
    confidence: float
    timestamp: str
    learnings: List[str]
    
    def to_dict(self) -> Dict:
        return {
            "intervention_id": self.intervention_id,
            "recommendation_id": self.recommendation_id,
            "location_id": self.location_id,
            "action_type": self.action_type,
            "expected_change": {k: round(v, 2) for k, v in self.expected_change.items()},
            "actual_change": {k: round(v, 2) for k, v in self.actual_change.items()},
            "variance": {k: round(v, 2) for k, v in self.variance.items()},
            "status": self.status.value,
            "confidence": round(self.confidence, 2),
            "timestamp": self.timestamp,
            "learnings": self.learnings
        }


class EvidenceEngine:
    """Tracks evidence for each recommendation type"""
    
    def __init__(self):
        self.evidence_db: Dict[str, Dict] = {}
    
    def record_outcome(self, outcome: InterventionOutcome):
        """Record outcome and update evidence"""
        action_type = outcome.action_type
        
        if action_type not in self.evidence_db:
            self.evidence_db[action_type] = {
                "action_type": action_type,
                "total_tests": 0,
                "successes": 0,
                "partials": 0,
                "failures": 0,
                "average_effect": {},
                "confidence": 0,
                "history": []
            }
        
        evidence = self.evidence_db[action_type]
        evidence["total_tests"] += 1
        evidence["history"].append(outcome.to_dict())
        
        # Update success counts
        if outcome.status == OutcomeStatus.SUCCESS:
            evidence["successes"] += 1
        elif outcome.status == OutcomeStatus.PARTIAL:
            evidence["partials"] += 1
        else:
            evidence["failures"] += 1
        
        # Update average effects
        for metric, change in outcome.actual_change.items():
            if metric not in evidence["average_effect"]:
                evidence["average_effect"][metric] = []
            evidence["average_effect"][metric].append(change)
        
        # Calculate confidence
        success_rate = (evidence["successes"] + evidence["partials"] * 0.5) / evidence["total_tests"]
        evidence["confidence"] = min(0.99, success_rate * (1 - 1 / evidence["total_tests"]))
        
        # Keep only last 100 outcomes
        evidence["history"] = evidence["history"][-100:]
    
    def get_evidence(self, action_type: str) -> Dict:
        """Get evidence for an action type"""
        evidence = self.evidence_db.get(action_type, {
            "action_type": action_type,
            "total_tests": 0,
            "confidence": 0,
            "message": "No evidence yet"
        })
        
        # Calculate averages
        avg_effects = {}
        for metric, values in evidence.get("average_effect", {}).items():
            if values:
                avg_effects[metric] = round(sum(values) / len(values), 2)
        
        return {
            "action_type": action_type,
            "total_tests": evidence["total_tests"],
            "successes": evidence.get("successes", 0),
            "partials": evidence.get("partials", 0),
            "failures": evidence.get("failures", 0),
            "success_rate": round(evidence.get("successes", 0) / max(evidence["total_tests"], 1) * 100, 1),
            "average_effects": avg_effects,
            "confidence": round(evidence.get("confidence", 0), 2),
            "reliability": self._classify_reliability(evidence["total_tests"], evidence.get("confidence", 0))
        }
    
    def _classify_reliability(self, n_tests: int, confidence: float) -> str:
        """Classify reliability of evidence"""
        if n_tests < 10:
            return "insufficient_data"
        elif confidence < 0.5:
            return "unreliable"
        elif confidence < 0.7:
            return "moderate"
        elif confidence < 0.9:
            return "reliable"
        else:
            return "highly_reliable"
    
    def compare_actions(self, action_types: List[str]) -> List[Dict]:
        """Compare evidence for multiple actions"""
        comparisons = []
        
        for action_type in action_types:
            evidence = self.get_evidence(action_type)
            comparisons.append(evidence)
        
        # Sort by confidence
        comparisons.sort(key=lambda x: x["confidence"], reverse=True)
        
        return comparisons


class BenchmarkEngine:
    """Benchmarks performance across locations and interventions"""
    
    def __init__(self):
        self.outcomes: List[InterventionOutcome] = []
        self.location_performance: Dict[str, List[Dict]] = {}
    
    def add_outcome(self, outcome: InterventionOutcome):
        """Add outcome to benchmark database"""
        self.outcomes.append(outcome)
        
        # Track by location
        loc = outcome.location_id
        if loc not in self.location_performance:
            self.location_performance[loc] = []
        self.location_performance[loc].append(outcome.to_dict())
    
    def benchmark_locations(self, metric: str = "safety_index") -> List[Dict]:
        """Benchmark locations by improvement in metric"""
        results = []
        
        for location_id, outcomes in self.location_performance.items():
            if not outcomes:
                continue
            
            # Calculate average improvement
            improvements = []
            for outcome in outcomes:
                actual = outcome.get("actual_change", {})
                if metric in actual:
                    improvements.append(actual[metric])
            
            if improvements:
                avg_improvement = sum(improvements) / len(improvements)
                results.append({
                    "location_id": location_id,
                    "interventions_count": len(outcomes),
                    "average_improvement": round(avg_improvement, 2),
                    "best_improvement": round(max(improvements), 2),
                    "success_rate": round(
                        sum(1 for o in outcomes if o.get("status") in ["success", "partial"]) / len(outcomes) * 100, 1
                    )
                })
        
        # Sort by improvement
        results.sort(key=lambda x: x["average_improvement"], reverse=True)
        
        return results
    
    def best_practices(self, area_type: str = "all") -> List[Dict]:
        """Identify best practices across all interventions"""
        # Group by action type
        by_action = {}
        for outcome in self.outcomes:
            action = outcome.action_type
            if action not in by_action:
                by_action[action] = []
            by_action[action].append(outcome)
        
        practices = []
        for action, outcomes in by_action.items():
            if len(outcomes) < 3:  # Need minimum data
                continue
            
            success_count = sum(1 for o in outcomes if o.status in [OutcomeStatus.SUCCESS, OutcomeStatus.PARTIAL])
            avg_variance = sum(
                sum(o.variance.values()) / max(len(o.variance), 1) 
                for o in outcomes
            ) / len(outcomes)
            
            practices.append({
                "action_type": action,
                "tested_count": len(outcomes),
                "success_rate": round(success_count / len(outcomes) * 100, 1),
                "average_variance": round(avg_variance, 2),
                "effectiveness": "high" if success_count / len(outcomes) > 0.7 else "medium" if success_count / len(outcomes) > 0.4 else "low"
            })
        
        # Sort by success rate
        practices.sort(key=lambda x: x["success_rate"], reverse=True)
        
        return practices


class LearningGraph:
    """Learning graph: Signal → Decision → Action → Result → Learning"""
    
    def __init__(self):
        self.learnings: List[Dict] = []
        self.signal_learnings: Dict[str, List[Dict]] = {}
    
    def add_learning(self, outcome: InterventionOutcome):
        """Extract learning from outcome"""
        learning = {
            "timestamp": outcome.timestamp,
            "action_type": outcome.action_type,
            "location_id": outcome.location_id,
            "expected": outcome.expected_change,
            "actual": outcome.actual_change,
            "variance": outcome.variance,
            "status": outcome.status.value,
            "learnings": outcome.learnings
        }
        
        self.learnings.append(learning)
        
        # Index by affected signals
        for signal in outcome.actual_change.keys():
            if signal not in self.signal_learnings:
                self.signal_learnings[signal] = []
            self.signal_learnings[signal].append(learning)
    
    def get_learnings_for_signal(self, signal_type: str) -> List[Dict]:
        """Get all learnings related to a signal type"""
        return self.signal_learnings.get(signal_type, [])
    
    def get_learnings_for_action(self, action_type: str) -> List[Dict]:
        """Get all learnings for an action type"""
        return [l for l in self.learnings if l["action_type"] == action_type]
    
    def get_insights(self) -> List[Dict]:
        """Generate insights from accumulated learnings"""
        insights = []
        
        # Find patterns in successful interventions
        successful = [l for l in self.learnings if l["status"] in ["success", "partial"]]
        
        if successful:
            # Most effective actions
            action_success = {}
            for learning in successful:
                action = learning["action_type"]
                if action not in action_success:
                    action_success[action] = []
                action_success[action].append(learning)
            
            for action, learnings in action_success.items():
                avg_improvement = sum(
                    sum(l["actual"].values()) / max(len(l["actual"]), 1)
                    for l in learnings
                ) / len(learnings)
                
                insights.append({
                    "type": "effective_action",
                    "action": action,
                    "evidence_count": len(learnings),
                    "average_improvement": round(avg_improvement, 2),
                    "insight": f"{action} has shown consistent positive results ({len(learnings)} interventions)"
                })
        
        # Find common failure patterns
        failed = [l for l in self.learnings if l["status"] in ["underperformed", "failed"]]
        if len(failed) > 5:
            insights.append({
                "type": "failure_pattern",
                "count": len(failed),
                "insight": f"{len(failed)} interventions underperformed. Consider reviewing implementation quality."
            })
        
        return insights


class InterventionLibrary:
    """Global library of interventions with proven results"""
    
    def __init__(self):
        self.interventions: Dict[str, Dict] = {}
    
    def register_intervention_type(
        self,
        intervention_type: str,
        description: str,
        typical_cost_usd: float,
        typical_timeline_months: int,
        expected_effects: Dict[str, float],
        required_resources: List[str]
    ):
        """Register an intervention type"""
        self.interventions[intervention_type] = {
            "type": intervention_type,
            "description": description,
            "typical_cost_usd": typical_cost_usd,
            "typical_timeline_months": typical_timeline_months,
            "expected_effects": expected_effects,
            "required_resources": required_resources,
            "evidence": {
                "tested_count": 0,
                "success_rate": 0,
                "average_actual_effects": {}
            }
        }
    
    def update_evidence(self, intervention_type: str, outcome: InterventionOutcome):
        """Update evidence with new outcome"""
        if intervention_type not in self.interventions:
            return
        
        intervention = self.interventions[intervention_type]
        evidence = intervention["evidence"]
        
        evidence["tested_count"] += 1
        
        # Update success rate
        if outcome.status in [OutcomeStatus.SUCCESS, OutcomeStatus.PARTIAL]:
            current_success = evidence["success_rate"] * (evidence["tested_count"] - 1)
            evidence["success_rate"] = (current_success + 1) / evidence["tested_count"]
        
        # Update average actual effects
        for metric, change in outcome.actual_change.items():
            if metric not in evidence["average_actual_effects"]:
                evidence["average_actual_effects"][metric] = []
            # Ensure it's a list (not a float from previous averaging)
            if not isinstance(evidence["average_actual_effects"][metric], list):
                evidence["average_actual_effects"][metric] = []
            evidence["average_actual_effects"][metric].append(change)
        
        # Calculate averages (store raw values separately, show averages)
        for metric, values in evidence["average_actual_effects"].items():
            if isinstance(values, list) and values:
                evidence[f"_avg_{metric}"] = round(sum(values) / len(values), 2)
    
    def get_intervention(self, intervention_type: str) -> Optional[Dict]:
        """Get intervention details with evidence"""
        return self.interventions.get(intervention_type)
    
    def search_interventions(
        self,
        target_effect: Optional[str] = None,
        max_cost: Optional[float] = None,
        min_success_rate: Optional[float] = None
    ) -> List[Dict]:
        """Search interventions by criteria"""
        results = []
        
        for intervention in self.interventions.values():
            # Filter by target effect
            if target_effect and target_effect not in intervention["expected_effects"]:
                continue
            
            # Filter by cost
            if max_cost and intervention["typical_cost_usd"] > max_cost:
                continue
            
            # Filter by success rate
            success_rate = intervention["evidence"]["success_rate"]
            if min_success_rate and success_rate < min_success_rate:
                continue
            
            results.append(intervention)
        
        # Sort by success rate
        results.sort(key=lambda x: x["evidence"]["success_rate"], reverse=True)
        
        return results


class OutcomeEngine:
    """Main outcome engine coordinating all components"""
    
    def __init__(self):
        self.evidence = EvidenceEngine()
        self.benchmark = BenchmarkEngine()
        self.learning = LearningGraph()
        self.library = InterventionLibrary()
        
        # Register common interventions
        self._register_common_interventions()
    
    def _register_common_interventions(self):
        """Register common intervention types"""
        self.library.register_intervention_type(
            "replace_lighting",
            "Replace street lighting with LED",
            800000,
            4,
            {"safety_index": 15, "night_activity": 20, "energy_efficiency": 25},
            ["electrician", "materials", "permits"]
        )
        
        self.library.register_intervention_type(
            "repair_sidewalks",
            "Repair damaged sidewalks",
            400000,
            3,
            {"walkability": 20, "safety_index": 10, "accessibility": 15},
            ["construction_crew", "materials"]
        )
        
        self.library.register_intervention_type(
            "plant_trees",
            "Plant trees along streets",
            200000,
            12,
            {"shade_index": 25, "heat_stress": -15, "walkability": 10},
            ["landscaping_crew", "saplings", "irrigation"]
        )
        
        self.library.register_intervention_type(
            "remove_graffiti",
            "Remove graffiti and paint walls",
            50000,
            1,
            {"visual_maintenance": 30, "safety_index": 8},
            ["painting_crew", "paint"]
        )
    
    def process_outcome(self, outcome: InterventionOutcome):
        """Process a new outcome through all systems"""
        # Update evidence
        self.evidence.record_outcome(outcome)
        
        # Update benchmarks
        self.benchmark.add_outcome(outcome)
        
        # Add to learning graph
        self.learning.add_learning(outcome)
        
        # Update intervention library
        self.library.update_evidence(outcome.action_type, outcome)
        
        return {
            "status": "processed",
            "intervention_id": outcome.intervention_id,
            "evidence_updated": True,
            "learnings_extracted": len(outcome.learnings)
        }
    
    def get_recommendation_with_evidence(self, action_type: str) -> Dict:
        """Get recommendation backed by evidence"""
        intervention = self.library.get_intervention(action_type)
        evidence = self.evidence.get_evidence(action_type)
        
        return {
            "intervention": intervention,
            "evidence": evidence,
            "confidence": evidence["confidence"],
            "recommendation": self._generate_recommendation(intervention, evidence)
        }
    
    def _generate_recommendation(self, intervention: Dict, evidence: Dict) -> str:
        """Generate human-readable recommendation"""
        if evidence["total_tests"] == 0:
            return f"{intervention['description']} - No evidence yet"
        
        if evidence["reliability"] == "highly_reliable":
            return f"{intervention['description']} - Highly recommended ({evidence['success_rate']}% success rate, {evidence['total_tests']} tests)"
        elif evidence["reliability"] == "reliable":
            return f"{intervention['description']} - Recommended ({evidence['success_rate']}% success rate)"
        elif evidence["reliability"] == "moderate":
            return f"{intervention['description']} - Consider with caution ({evidence['success_rate']}% success rate)"
        else:
            return f"{intervention['description']} - Insufficient evidence"
    
    def get_system_stats(self) -> Dict:
        """Get overall system statistics"""
        return {
            "total_outcomes": len(self.benchmark.outcomes),
            "total_learnings": len(self.learning.learnings),
            "intervention_types": len(self.library.interventions),
            "evidence_entries": len(self.evidence.evidence_db),
            "insights": self.learning.get_insights()
        }


# Example usage
def example_outcome_engine():
    """Example: Outcome Engine in action"""
    engine = OutcomeEngine()
    
    print("=== Outcome Engine Demo ===")
    print(f"Registered interventions: {len(engine.library.interventions)}")
    
    # Simulate outcomes
    outcomes = [
        InterventionOutcome(
            intervention_id="INT-001",
            recommendation_id="REC-001",
            location_id="LOC-001",
            action_type="replace_lighting",
            baseline_state={"safety_index": 45, "night_activity": 30},
            post_state={"safety_index": 62, "night_activity": 55},
            expected_change={"safety_index": 15, "night_activity": 20},
            actual_change={"safety_index": 17, "night_activity": 25},
            variance={"safety_index": 2, "night_activity": 5},
            status=OutcomeStatus.SUCCESS,
            confidence=0.85,
            timestamp=datetime.utcnow().isoformat(),
            learnings=["LED lighting significantly improved safety", "Color temperature 4000K optimal"]
        ),
        InterventionOutcome(
            intervention_id="INT-002",
            recommendation_id="REC-002",
            location_id="LOC-002",
            action_type="replace_lighting",
            baseline_state={"safety_index": 40, "night_activity": 25},
            post_state={"safety_index": 52, "night_activity": 38},
            expected_change={"safety_index": 15, "night_activity": 20},
            actual_change={"safety_index": 12, "night_activity": 13},
            variance={"safety_index": -3, "night_activity": -7},
            status=OutcomeStatus.PARTIAL,
            confidence=0.75,
            timestamp=datetime.utcnow().isoformat(),
            learnings=["Partial success due to uneven coverage"]
        ),
        InterventionOutcome(
            intervention_id="INT-003",
            recommendation_id="REC-003",
            location_id="LOC-003",
            action_type="plant_trees",
            baseline_state={"shade_index": 20, "heat_stress": 80},
            post_state={"shade_index": 50, "heat_stress": 60},
            expected_change={"shade_index": 25, "heat_stress": -15},
            actual_change={"shade_index": 30, "heat_stress": -20},
            variance={"shade_index": 5, "heat_stress": -5},
            status=OutcomeStatus.SUCCESS,
            confidence=0.9,
            timestamp=datetime.utcnow().isoformat(),
            learnings=["Tree species selection critical", "Native species performed best"]
        )
    ]
    
    # Process outcomes
    print("\n=== Processing Outcomes ===")
    for outcome in outcomes:
        result = engine.process_outcome(outcome)
        print(f"Processed {outcome.intervention_id}: {outcome.status.value}")
    
    # Get evidence
    print("\n=== Evidence for Lighting Replacement ===")
    evidence = engine.evidence.get_evidence("replace_lighting")
    print(f"Tests: {evidence['total_tests']}")
    print(f"Success rate: {evidence['success_rate']}%")
    print(f"Confidence: {evidence['confidence']}")
    print(f"Reliability: {evidence['reliability']}")
    
    # Get recommendation with evidence
    print("\n=== Recommendation with Evidence ===")
    rec = engine.get_recommendation_with_evidence("replace_lighting")
    print(f"Recommendation: {rec['recommendation']}")
    
    # Benchmark locations
    print("\n=== Location Benchmarks ===")
    benchmarks = engine.benchmark.benchmark_locations("safety_index")
    for b in benchmarks:
        print(f"  {b['location_id']}: {b['average_improvement']:.1f} avg improvement")
    
    # Best practices
    print("\n=== Best Practices ===")
    practices = engine.benchmark.best_practices()
    for p in practices:
        print(f"  {p['action_type']}: {p['success_rate']}% success ({p['tested_count']} tests)")
    
    # Insights
    print("\n=== AI Insights ===")
    insights = engine.learning.get_insights()
    for insight in insights:
        print(f"  [{insight['type']}] {insight['insight']}")
    
    # System stats
    print("\n=== System Stats ===")
    stats = engine.get_system_stats()
    print(f"Total outcomes: {stats['total_outcomes']}")
    print(f"Total learnings: {stats['total_learnings']}")
    print(f"Insights generated: {len(stats['insights'])}")
    
    return engine


if __name__ == '__main__':
    example_outcome_engine()