"""
Outcome Tracker
Measures whether actions actually produced desired effects
Closes the loop: Observation → Decision → Action → Outcome
"""

from typing import Dict, List, Optional
from dataclasses import dataclass
from datetime import datetime


@dataclass
class Outcome:
    """Result of an intervention"""
    intervention_id: str
    location_id: str
    decision_type: str
    expected_impact: float
    actual_impact: float
    variance: float
    confidence: float
    timestamp: str
    
    @property
    def variance_percent(self) -> float:
        return round((self.variance / max(self.expected_impact, 1)) * 100, 1)
    
    def to_dict(self) -> Dict:
        return {
            "intervention_id": self.intervention_id,
            "location_id": self.location_id,
            "decision_type": self.decision_type,
            "expected_impact": round(self.expected_impact, 2),
            "actual_impact": round(self.actual_impact, 2),
            "variance": round(self.variance, 2),
            "variance_percent": self.variance_percent,
            "confidence": round(self.confidence, 2),
            "timestamp": self.timestamp
        }


class OutcomeTracker:
    """Tracks outcomes of interventions"""
    
    def __init__(self):
        self.interventions: Dict[str, Dict] = {}
        self.outcomes: Dict[str, Outcome] = {}
    
    def register_intervention(
        self,
        intervention_id: str,
        location_id: str,
        decision_type: str,
        expected_impact: float,
        baseline_signals: Dict[str, float]
    ) -> Dict:
        """
        Register an intervention for tracking
        
        Args:
            intervention_id: Unique ID
            location_id: Where
            decision_type: What type
            expected_impact: Predicted effect
            baseline_signals: Signal values before intervention
            
        Returns:
            Registration confirmation
        """
        self.interventions[intervention_id] = {
            "intervention_id": intervention_id,
            "location_id": location_id,
            "decision_type": decision_type,
            "expected_impact": expected_impact,
            "baseline_signals": baseline_signals,
            "started_at": datetime.utcnow().isoformat(),
            "status": "active"
        }
        
        return {
            "status": "registered",
            "intervention_id": intervention_id,
            "message": f"Tracking {decision_type} at {location_id}. Expected impact: {expected_impact}"
        }
    
    def measure_outcome(
        self,
        intervention_id: str,
        post_signals: Dict[str, float]
    ) -> Outcome:
        """
        Measure actual outcome of an intervention
        
        Args:
            intervention_id: ID of intervention
            post_signals: Signal values after intervention
            
        Returns:
            Outcome measurement
        """
        intervention = self.interventions.get(intervention_id)
        if not intervention:
            raise ValueError(f"Intervention {intervention_id} not found")
        
        baseline = intervention["baseline_signals"]
        
        # Calculate actual impact
        actual_impact = self._calculate_impact(baseline, post_signals, intervention["decision_type"])
        
        # Calculate variance
        expected = intervention["expected_impact"]
        variance = actual_impact - expected
        
        # Calculate confidence
        confidence = self._calculate_confidence(baseline, post_signals)
        
        outcome = Outcome(
            intervention_id=intervention_id,
            location_id=intervention["location_id"],
            decision_type=intervention["decision_type"],
            expected_impact=expected,
            actual_impact=actual_impact,
            variance=variance,
            confidence=confidence,
            timestamp=datetime.utcnow().isoformat()
        )
        
        self.outcomes[intervention_id] = outcome
        
        # Update intervention status
        intervention["status"] = "completed"
        intervention["completed_at"] = datetime.utcnow().isoformat()
        
        return outcome
    
    def get_outcome_report(self, intervention_id: str) -> Dict:
        """Get full outcome report"""
        intervention = self.interventions.get(intervention_id)
        outcome = self.outcomes.get(intervention_id)
        
        if not intervention:
            return {"error": "Intervention not found"}
        
        report = {
            "intervention": intervention,
            "outcome": outcome.to_dict() if outcome else None,
            "success": self._assess_success(intervention, outcome) if outcome else None
        }
        
        return report
    
    def get_aggregate_stats(self) -> Dict:
        """Get aggregate statistics across all interventions"""
        if not self.outcomes:
            return {"status": "no_data"}
        
        outcomes_list = list(self.outcomes.values())
        
        # Calculate statistics
        total_interventions = len(self.interventions)
        completed = len(outcomes_list)
        
        avg_expected = sum(o.expected_impact for o in outcomes_list) / len(outcomes_list)
        avg_actual = sum(o.actual_impact for o in outcomes_list) / len(outcomes_list)
        avg_variance = sum(o.variance for o in outcomes_list) / len(outcomes_list)
        
        # Success rate
        successful = sum(1 for o in outcomes_list if o.variance > -10)  # Within 10% of expected
        success_rate = successful / len(outcomes_list) * 100
        
        # By decision type
        by_type = {}
        for outcome in outcomes_list:
            dtype = outcome.decision_type
            if dtype not in by_type:
                by_type[dtype] = []
            by_type[dtype].append(outcome)
        
        type_stats = {}
        for dtype, outcomes in by_type.items():
            type_stats[dtype] = {
                "count": len(outcomes),
                "avg_expected": round(sum(o.expected_impact for o in outcomes) / len(outcomes), 2),
                "avg_actual": round(sum(o.actual_impact for o in outcomes) / len(outcomes), 2),
                "avg_variance": round(sum(o.variance for o in outcomes) / len(outcomes), 2)
            }
        
        return {
            "total_interventions": total_interventions,
            "completed": completed,
            "success_rate": round(success_rate, 1),
            "avg_expected_impact": round(avg_expected, 2),
            "avg_actual_impact": round(avg_actual, 2),
            "avg_variance": round(avg_variance, 2),
            "by_decision_type": type_stats
        }
    
    def _calculate_impact(
        self,
        baseline: Dict[str, float],
        post: Dict[str, float],
        decision_type: str
    ) -> float:
        """Calculate impact of intervention"""
        # Calculate average improvement across relevant signals
        improvements = []
        
        for signal, post_value in post.items():
            if signal in baseline:
                baseline_value = baseline[signal]
                improvement = post_value - baseline_value
                improvements.append(improvement)
        
        if improvements:
            return sum(improvements) / len(improvements)
        
        return 0
    
    def _calculate_confidence(
        self,
        baseline: Dict[str, float],
        post: Dict[str, float]
    ) -> float:
        """Calculate confidence in measurement"""
        # More signals measured = higher confidence
        measured_signals = sum(1 for s in baseline if s in post)
        total_signals = len(baseline)
        
        if total_signals == 0:
            return 0
        
        return min(0.95, measured_signals / total_signals)
    
    def _assess_success(self, intervention: Dict, outcome: Outcome) -> Dict:
        """Assess whether intervention was successful"""
        variance_pct = (outcome.variance / max(intervention["expected_impact"], 1)) * 100
        
        if variance_pct >= -10:
            status = "success"
            message = "Intervention achieved expected impact"
        elif variance_pct >= -25:
            status = "partial"
            message = "Intervention achieved partial impact"
        else:
            status = "underperformed"
            message = "Intervention underperformed expectations"
        
        return {
            "status": status,
            "message": message,
            "variance_percent": round(variance_pct, 1),
            "recommendation": self._generate_recommendation(status, intervention, outcome)
        }
    
    def _generate_recommendation(
        self,
        status: str,
        intervention: Dict,
        outcome: Outcome
    ) -> str:
        """Generate follow-up recommendation"""
        if status == "success":
            return "Consider scaling this intervention to similar locations"
        elif status == "partial":
            return "Review implementation and consider complementary measures"
        else:
            return "Reassess approach. Consider different intervention type or location"


# Example usage
def example_outcome_tracking():
    """Example: Track intervention outcomes"""
    tracker = OutcomeTracker()
    
    # Register intervention
    print("=== Register Intervention ===")
    result = tracker.register_intervention(
        intervention_id="INT-001",
        location_id="LOC-001",
        decision_type="safety_improvement",
        expected_impact=15,
        baseline_signals={
            "safety_index": 45,
            "lighting": 35,
            "crime_rate": 60,
            "pedestrian_flow": 50
        }
    )
    print(f"Status: {result['status']}")
    print(f"Message: {result['message']}")
    
    # Measure outcome (after intervention)
    print("\n=== Measure Outcome ===")
    outcome = tracker.measure_outcome(
        intervention_id="INT-001",
        post_signals={
            "safety_index": 58,
            "lighting": 65,
            "crime_rate": 45,
            "pedestrian_flow": 55
        }
    )
    
    print(f"Expected impact: {outcome.expected_impact}")
    print(f"Actual impact: {outcome.actual_impact}")
    print(f"Variance: {outcome.variance} ({outcome.variance_percent}%)")
    print(f"Confidence: {outcome.confidence}")
    
    # Get report
    print("\n=== Outcome Report ===")
    report = tracker.get_outcome_report("INT-001")
    success = report.get("success", {})
    print(f"Status: {success.get('status')}")
    print(f"Message: {success.get('message')}")
    print(f"Recommendation: {success.get('recommendation')}")
    
    # Add more interventions
    tracker.register_intervention(
        intervention_id="INT-002",
        location_id="LOC-002",
        decision_type="maintenance",
        expected_impact=20,
        baseline_signals={"maintenance_quality": 40, "infrastructure_reliability": 50}
    )
    
    tracker.measure_outcome(
        intervention_id="INT-002",
        post_signals={"maintenance_quality": 55, "infrastructure_reliability": 60}
    )
    
    # Aggregate stats
    print("\n=== Aggregate Statistics ===")
    stats = tracker.get_aggregate_stats()
    print(f"Total interventions: {stats['total_interventions']}")
    print(f"Completed: {stats['completed']}")
    print(f"Success rate: {stats['success_rate']}%")
    print(f"Average variance: {stats['avg_variance']}")
    
    return tracker


if __name__ == '__main__':
    example_outcome_tracking()