"""
Simulation Engine (Digital Twin 2.0)
Simulates "what if" scenarios
"""

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


@dataclass
class SimulationResult:
    """Result of a simulation"""
    scenario_name: str
    baseline: Dict
    simulated: Dict
    changes: Dict
    impacted_indexes: List[Dict]
    
    def to_dict(self) -> Dict:
        return {
            "scenario": self.scenario_name,
            "baseline": self.baseline,
            "simulated": self.simulated,
            "changes": self.changes,
            "impacted_indexes": self.impacted_indexes
        }


class SimulationEngine:
    """Simulates urban scenarios"""
    
    def __init__(self):
        self.impact_models = self._load_impact_models()
    
    def _load_impact_models(self) -> Dict:
        """Load models for how changes affect indexes"""
        return {
            # Add bike lanes
            "add_bike_lanes": {
                "description": "Add dedicated bicycle lanes",
                "changes": {
                    "bicycle_friendliness": +40,
                    "walkability": +10,
                    "traffic_intensity": -5,
                    "noise_level": -3
                },
                "cost_estimate_usd": 150000,
                "implementation_months": 6
            },
            
            # Remove parking
            "remove_parking": {
                "description": "Remove on-street parking",
                "changes": {
                    "walkability": +15,
                    "pedestrian_flow": +20,
                    "retail_density": +10,
                    "traffic_intensity": -10
                },
                "cost_estimate_usd": 50000,
                "implementation_months": 3
            },
            
            # Plant trees
            "plant_40_trees": {
                "description": "Plant 40 new trees",
                "changes": {
                    "tree_coverage": +40,
                    "heat_stress": -15,
                    "shade_index": +25,
                    "walkability": +10,
                    "property_value": +5
                },
                "cost_estimate_usd": 80000,
                "implementation_months": 12
            },
            
            # Improve lighting
            "improve_lighting": {
                "description": "Upgrade street lighting to LED",
                "changes": {
                    "lighting": +30,
                    "safety_index": +15,
                    "night_activity": +20,
                    "energy_efficiency": +25
                },
                "cost_estimate_usd": 120000,
                "implementation_months": 4
            },
            
            # Add playground
            "add_playground": {
                "description": "Add children's playground",
                "changes": {
                    "family_presence": +35,
                    "greenery": +10,
                    "community_activity": +20,
                    "property_value": +8
                },
                "cost_estimate_usd": 200000,
                "implementation_months": 8
            },
            
            # Mixed-use development
            "mixed_use_development": {
                "description": "Convert ground floor to mixed-use",
                "changes": {
                    "retail_density": +30,
                    "public_life": +25,
                    "pedestrian_flow": +20,
                    "functional_density": +35,
                    "property_value": +15
                },
                "cost_estimate_usd": 500000,
                "implementation_months": 18
            }
        }
    
    def simulate(
        self,
        current_state: Dict[str, float],
        scenario_name: str
    ) -> SimulationResult:
        """
        Simulate a scenario
        
        Args:
            current_state: Current signal values
            scenario_name: Name of scenario to simulate
            
        Returns:
            Simulation result
        """
        model = self.impact_models.get(scenario_name)
        if not model:
            return SimulationResult(
                scenario_name=scenario_name,
                baseline=current_state,
                simulated=current_state,
                changes={},
                impacted_indexes=[]
            )
        
        # Apply changes
        simulated_state = current_state.copy()
        changes = {}
        
        for signal, delta in model["changes"].items():
            if signal in simulated_state:
                old_value = simulated_state[signal]
                new_value = min(100, max(0, old_value + delta))
                simulated_state[signal] = new_value
                changes[signal] = {
                    "old": old_value,
                    "new": new_value,
                    "delta": round(new_value - old_value, 1)
                }
        
        # Calculate impacted indexes
        impacted = self._calculate_impacted_indexes(current_state, simulated_state)
        
        return SimulationResult(
            scenario_name=scenario_name,
            baseline=current_state,
            simulated=simulated_state,
            changes=changes,
            impacted_indexes=impacted
        )
    
    def compare_scenarios(
        self,
        current_state: Dict[str, float],
        scenarios: List[str]
    ) -> Dict:
        """Compare multiple scenarios"""
        results = []
        
        for scenario_name in scenarios:
            result = self.simulate(current_state, scenario_name)
            model = self.impact_models.get(scenario_name, {})
            
            # Calculate overall improvement
            total_improvement = sum(
                c["delta"] for c in result.changes.values()
                if c["delta"] > 0
            )
            
            results.append({
                "scenario": scenario_name,
                "description": model.get("description", ""),
                "total_improvement": round(total_improvement, 1),
                "cost_usd": model.get("cost_estimate_usd", 0),
                "implementation_months": model.get("implementation_months", 0),
                "changes": result.changes,
                "impacted_indexes": result.impacted_indexes
            })
        
        # Sort by improvement
        results.sort(key=lambda x: x["total_improvement"], reverse=True)
        
        return {
            "baseline": current_state,
            "scenarios": results,
            "best_scenario": results[0]["scenario"] if results else None
        }
    
    def optimize(
        self,
        current_state: Dict[str, float],
        target_index: str,
        budget_usd: float
    ) -> Dict:
        """
        Find optimal combination of interventions within budget
        
        Example:
        "Maximize family-friendly within $300k budget"
        """
        # Simple greedy optimization
        affordable = []
        
        for scenario_name, model in self.impact_models.items():
            if model.get("cost_estimate_usd", float('inf')) <= budget_usd:
                result = self.simulate(current_state, scenario_name)
                
                # Calculate target improvement
                target_improvement = 0
                for idx in result.impacted_indexes:
                    if idx["index_name"] == target_index:
                        target_improvement = idx["improvement"]
                        break
                
                affordable.append({
                    "scenario": scenario_name,
                    "cost": model["cost_estimate_usd"],
                    "target_improvement": target_improvement,
                    "roi": target_improvement / max(model["cost_estimate_usd"], 1) * 100000
                })
        
        # Sort by ROI
        affordable.sort(key=lambda x: x["roi"], reverse=True)
        
        # Select best combination within budget
        selected = []
        remaining_budget = budget_usd
        
        for option in affordable:
            if option["cost"] <= remaining_budget:
                selected.append(option)
                remaining_budget -= option["cost"]
        
        total_cost = sum(s["cost"] for s in selected)
        total_improvement = sum(s["target_improvement"] for s in selected)
        
        return {
            "target_index": target_index,
            "budget_usd": budget_usd,
            "selected_scenarios": [s["scenario"] for s in selected],
            "total_cost": total_cost,
            "total_improvement": round(total_improvement, 1),
            "remaining_budget": remaining_budget,
            "efficiency": round(total_improvement / max(total_cost, 1) * 100000, 2)
        }
    
    def _calculate_impacted_indexes(
        self,
        baseline: Dict,
        simulated: Dict
    ) -> List[Dict]:
        """Calculate how indexes change"""
        impacted = []
        
        # Family Friendly Index
        family_signals = ["family_presence", "safety_index", "greenery", "playground", "noise_level"]
        family_baseline = sum(baseline.get(s, 50) for s in family_signals) / len(family_signals)
        family_simulated = sum(simulated.get(s, 50) for s in family_signals) / len(family_signals)
        
        if abs(family_simulated - family_baseline) > 1:
            impacted.append({
                "index_name": "family_friendly",
                "baseline": round(family_baseline, 1),
                "simulated": round(family_simulated, 1),
                "improvement": round(family_simulated - family_baseline, 1)
            })
        
        # Walkability Index
        walk_signals = ["sidewalk_width", "traffic_intensity", "shade_index", "bicycle_friendliness"]
        walk_baseline = sum(baseline.get(s, 50) for s in walk_signals) / len(walk_signals)
        walk_simulated = sum(simulated.get(s, 50) for s in walk_signals) / len(walk_signals)
        
        if abs(walk_simulated - walk_baseline) > 1:
            impacted.append({
                "index_name": "walkability",
                "baseline": round(walk_baseline, 1),
                "simulated": round(walk_simulated, 1),
                "improvement": round(walk_simulated - walk_baseline, 1)
            })
        
        # Property Value Index
        property_signals = ["walkability", "transit_access", "greenery", "safety_index", "retail_density"]
        prop_baseline = sum(baseline.get(s, 50) for s in property_signals) / len(property_signals)
        prop_simulated = sum(simulated.get(s, 50) for s in property_signals) / len(property_signals)
        
        if abs(prop_simulated - prop_baseline) > 1:
            impacted.append({
                "index_name": "property_value",
                "baseline": round(prop_baseline, 1),
                "simulated": round(prop_simulated, 1),
                "improvement": round(prop_simulated - prop_baseline, 1)
            })
        
        return impacted


# Example usage
def example_simulation():
    """Example: Simulate scenarios"""
    engine = SimulationEngine()
    
    # Current state
    current_state = {
        "walkability": 60,
        "bicycle_friendliness": 30,
        "traffic_intensity": 70,
        "noise_level": 75,
        "tree_coverage": 20,
        "heat_stress": 80,
        "shade_index": 25,
        "lighting": 40,
        "safety_index": 45,
        "night_activity": 35,
        "family_presence": 30,
        "playground": 10,
        "greenery": 25,
        "retail_density": 50,
        "pedestrian_flow": 55,
        "property_value": 60,
        "functional_density": 45
    }
    
    # Simulate single scenario
    print("=== Simulate: Plant 40 Trees ===")
    result = engine.simulate(current_state, "plant_40_trees")
    
    print(f"Changes:")
    for signal, change in result.changes.items():
        print(f"  {signal}: {change['old']} → {change['new']} ({change['delta']:+.1f})")
    
    print(f"\nImpacted Indexes:")
    for idx in result.impacted_indexes:
        print(f"  {idx['index_name']}: {idx['baseline']} → {idx['simulated']} ({idx['improvement']:+.1f})")
    
    # Compare scenarios
    print("\n=== Compare Scenarios ===")
    comparison = engine.compare_scenarios(current_state, [
        "plant_40_trees",
        "add_bike_lanes",
        "improve_lighting",
        "add_playground"
    ])
    
    for scenario in comparison["scenarios"]:
        print(f"\n{scenario['scenario']}:")
        print(f"  Improvement: {scenario['total_improvement']}")
        print(f"  Cost: ${scenario['cost_usd']:,}")
        print(f"  Time: {scenario['implementation_months']} months")
    
    print(f"\nBest scenario: {comparison['best_scenario']}")
    
    # Optimize
    print("\n=== Optimize: Family Friendly ($300k budget) ===")
    optimization = engine.optimize(current_state, "family_friendly", 300000)
    
    print(f"Selected: {optimization['selected_scenarios']}")
    print(f"Total cost: ${optimization['total_cost']:,}")
    print(f"Total improvement: {optimization['total_improvement']}")
    print(f"Efficiency: {optimization['efficiency']}")
    
    return result


if __name__ == '__main__':
    example_simulation()