"""
Digital Twin - Layer 10 of IOM
Knowledge graph, timeline visualization, predictive models
"""

from typing import Dict, List, Optional, Any
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from collections import defaultdict
import json


@dataclass
class DigitalTwinNode:
    """Node in the digital twin graph"""
    goid: str
    object_type: str
    domain: str
    system: str
    subsystem: str
    
    # State
    condition: int = 3  # 1-5
    risk_level: float = 0.0
    function_status: str = "operational"
    
    # Location
    lat: Optional[float] = None
    lng: Optional[float] = None
    elevation: Optional[float] = None
    
    # Visual Geolocation (NEW)
    visual_evidence: List[Dict] = field(default_factory=list)
    evidence_history: List[Dict] = field(default_factory=list)
    geolocation_estimates: List[Dict] = field(default_factory=list)
    
    # Metadata
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    # Timeline
    observations: List[Dict] = field(default_factory=list)
    events: List[Dict] = field(default_factory=list)
    
    # Relations
    relations: List[Dict] = field(default_factory=list)
    
    def add_observation(self, observation: Dict):
        """Add observation to timeline"""
        self.observations.append({
            "timestamp": observation.get("timestamp", datetime.utcnow().isoformat()),
            "condition": observation.get("overall_condition", 3),
            "findings": observation.get("findings", []),
            "risk_level": observation.get("risk_level", 0.0)
        })
        
        # Update current state
        if "overall_condition" in observation:
            self.condition = observation["overall_condition"]
        if "risk_level" in observation:
            self.risk_level = observation["risk_level"]
    
    def add_visual_evidence(self, evidence_package: Dict):
        """Add visual geolocation evidence (NEW)"""
        self.visual_evidence.append({
            "timestamp": datetime.utcnow().isoformat(),
            "evidence": evidence_package
        })
        
        # Update evidence history
        self.evidence_history.append({
            "timestamp": datetime.utcnow().isoformat(),
            "visual_objects": len(evidence_package.get("visual_objects", [])),
            "text_detections": len(evidence_package.get("text_detections", [])),
            "geometric_features": len(evidence_package.get("geometric_features", [])),
            "environmental_signals": len(evidence_package.get("environmental_signals", []))
        })
    
    def add_geolocation_estimate(self, estimate: Dict):
        """Add geolocation estimate (NEW)"""
        self.geolocation_estimates.append({
            "timestamp": datetime.utcnow().isoformat(),
            "lat": estimate.get("lat"),
            "lng": estimate.get("lng"),
            "accuracy": estimate.get("accuracy"),
            "confidence": estimate.get("confidence"),
            "method": estimate.get("method")
        })
    
    def add_event(self, event_type: str, description: str, metadata: Dict = None):
        """Add event to timeline"""
        self.events.append({
            "timestamp": datetime.utcnow().isoformat(),
            "type": event_type,
            "description": description,
            "metadata": metadata or {}
        })
    
    def add_relation(self, relation_type: str, target_goid: str, metadata: Dict = None):
        """Add relation to another node"""
        self.relations.append({
            "type": relation_type,
            "target_goid": target_goid,
            "metadata": metadata or {}
        })
    
    def get_timeline(self, months: int = 12) -> List[Dict]:
        """Get timeline for last N months"""
        from datetime import timezone
        cutoff = datetime.now(timezone.utc) - timedelta(days=30*months)
        
        timeline = []
        
        # Add observations
        for obs in self.observations:
            ts = obs["timestamp"]
            if ts.endswith('Z'):
                ts = ts[:-1] + '+00:00'
            obs_time = datetime.fromisoformat(ts)
            if obs_time >= cutoff:
                timeline.append({
                    "date": obs["timestamp"],
                    "type": "observation",
                    "condition": obs["condition"],
                    "findings": obs["findings"]
                })
        
        # Add events
        for event in self.events:
            ts = event["timestamp"]
            if ts.endswith('Z'):
                ts = ts[:-1] + '+00:00'
            event_time = datetime.fromisoformat(ts)
            if event_time >= cutoff:
                timeline.append({
                    "date": event["timestamp"],
                    "type": event["type"],
                    "description": event["description"]
                })
        
        # Sort by date
        timeline.sort(key=lambda x: x["date"])
        
        return timeline
    
    def get_trend(self, months: int = 6) -> Dict:
        """Get trend analysis"""
        from datetime import timezone
        cutoff = datetime.now(timezone.utc) - timedelta(days=30*months)
        
        recent_obs = []
        for obs in self.observations:
            ts = obs["timestamp"]
            if ts.endswith('Z'):
                ts = ts[:-1] + '+00:00'
            obs_time = datetime.fromisoformat(ts)
            if obs_time >= cutoff:
                recent_obs.append(obs)
        
        if len(recent_obs) < 2:
            return {"trend": "insufficient_data", "change": 0}
        
        conditions = [obs["condition"] for obs in recent_obs]
        first = conditions[0]
        last = conditions[-1]
        
        change = last - first
        
        if change < 0:
            trend = "improving"
        elif change > 0:
            trend = "degrading"
        else:
            trend = "stable"
        
        return {
            "trend": trend,
            "change": change,
            "first_condition": first,
            "last_condition": last,
            "observation_count": len(recent_obs)
        }


class DigitalTwin:
    """Digital twin for a city/region"""
    
    def __init__(self, name: str):
        self.name = name
        self.nodes: Dict[str, DigitalTwinNode] = {}
        self.created_at = datetime.utcnow()
    
    def add_node(self, node: DigitalTwinNode):
        """Add node to digital twin"""
        self.nodes[node.goid] = node
    
    def get_node(self, goid: str) -> Optional[DigitalTwinNode]:
        """Get node by GOID"""
        return self.nodes.get(goid)
    
    def find_by_location(
        self,
        lat: float,
        lng: float,
        radius_km: float
    ) -> List[DigitalTwinNode]:
        """Find nodes near location"""
        from math import radians, sin, cos, sqrt, atan2
        
        results = []
        
        for node in self.nodes.values():
            if node.lat is None or node.lng is None:
                continue
            
            # Haversine formula
            R = 6371  # Earth radius in km
            
            lat1, lon1 = radians(lat), radians(lng)
            lat2, lon2 = radians(node.lat), radians(node.lng)
            
            dlat = lat2 - lat1
            dlon = lon2 - lon1
            
            a = sin(dlat/2)**2 + cos(lat1) * cos(lat2) * sin(dlon/2)**2
            c = 2 * atan2(sqrt(a), sqrt(1-a))
            
            distance = R * c
            
            if distance <= radius_km:
                results.append((node, distance))
        
        # Sort by distance
        results.sort(key=lambda x: x[1])
        
        return [node for node, _ in results]
    
    def find_by_condition(
        self,
        min_condition: int = 1,
        max_condition: int = 5
    ) -> List[DigitalTwinNode]:
        """Find nodes by condition range"""
        return [
            node for node in self.nodes.values()
            if min_condition <= node.condition <= max_condition
        ]
    
    def find_by_risk(
        self,
        min_risk: float = 0.0,
        max_risk: float = 10.0
    ) -> List[DigitalTwinNode]:
        """Find nodes by risk range"""
        return [
            node for node in self.nodes.values()
            if min_risk <= node.risk_level <= max_risk
        ]
    
    def find_degrading(self, threshold: float = 1.0) -> List[Dict]:
        """Find nodes that are degrading"""
        results = []
        
        for node in self.nodes.values():
            trend = node.get_trend()
            if trend["trend"] == "degrading" and abs(trend["change"]) >= threshold:
                results.append({
                    "goid": node.goid,
                    "object_type": node.object_type,
                    "condition": node.condition,
                    "risk_level": node.risk_level,
                    "trend": trend
                })
        
        # Sort by risk level
        results.sort(key=lambda x: x["risk_level"], reverse=True)
        
        return results
    
    def get_statistics(self) -> Dict:
        """Get statistics for digital twin"""
        total = len(self.nodes)
        
        if total == 0:
            return {"total": 0}
        
        # Condition distribution
        conditions = defaultdict(int)
        for node in self.nodes.values():
            conditions[node.condition] += 1
        
        # Risk distribution
        risk_ranges = {
            "minimal": 0,
            "low": 0,
            "medium": 0,
            "high": 0,
            "critical": 0
        }
        
        for node in self.nodes.values():
            if node.risk_level < 2:
                risk_ranges["minimal"] += 1
            elif node.risk_level < 4:
                risk_ranges["low"] += 1
            elif node.risk_level < 6:
                risk_ranges["medium"] += 1
            elif node.risk_level < 8:
                risk_ranges["high"] += 1
            else:
                risk_ranges["critical"] += 1
        
        # Domain distribution
        domains = defaultdict(int)
        for node in self.nodes.values():
            domains[node.domain] += 1
        
        return {
            "total": total,
            "conditions": dict(conditions),
            "risk_ranges": risk_ranges,
            "domains": dict(domains),
            "degrading": len(self.find_degrading()),
            "high_risk": len(self.find_by_risk(6.0, 10.0)),
            "critical": len(self.find_by_risk(8.0, 10.0))
        }
    
    def get_contradictions(self, official_data: Dict) -> List[Dict]:
        """
        Find contradictions between official data and observations
        
        Args:
            official_data: Dict with official reports
            
        Returns:
            List of contradictions
        """
        contradictions = []
        
        # Example: Official says "all bridges in good condition"
        # But observations show structural defects
        
        if "bridges_good_condition" in official_data:
            official_count = official_data["bridges_good_condition"]
            
            # Find bridges with condition > 3 (not good)
            bad_bridges = [
                node for node in self.nodes.values()
                if node.object_type in ["abutment", "pier", "deck"]
                and node.condition > 3
            ]
            
            if len(bad_bridges) > 0:
                contradictions.append({
                    "type": "condition_discrepancy",
                    "severity": "high",
                    "description": f"Official: {official_count} bridges good. Observed: {len(bad_bridges)} with defects.",
                    "affected_objects": [node.goid for node in bad_bridges[:10]],
                    "estimated_underreporting": len(bad_bridges) / max(official_count, 1)
                })
        
        return contradictions
    
    def export_geojson(self) -> Dict:
        """Export all nodes as GeoJSON"""
        features = []
        
        for node in self.nodes.values():
            if node.lat is None or node.lng is None:
                continue
            
            features.append({
                "type": "Feature",
                "geometry": {
                    "type": "Point",
                    "coordinates": [node.lng, node.lat]
                },
                "properties": {
                    "goid": node.goid,
                    "object_type": node.object_type,
                    "condition": node.condition,
                    "risk_level": node.risk_level,
                    "function_status": node.function_status
                }
            })
        
        return {
            "type": "FeatureCollection",
            "features": features
        }
    
    def to_dict(self) -> Dict:
        """Convert to dictionary"""
        return {
            "name": self.name,
            "created_at": self.created_at.isoformat(),
            "node_count": len(self.nodes),
            "statistics": self.get_statistics()
        }


if __name__ == '__main__':
    # Example: Stockholm Digital Twin
    twin = DigitalTwin("Stockholm")
    
    # Add some nodes
    node1 = DigitalTwinNode(
        goid="TRN-BRG-ABT-CON-001",
        object_type="abutment",
        domain="TRN",
        system="BRG",
        subsystem="ABT",
        lat=59.3293,
        lng=18.0686,
        condition=3,
        risk_level=5.9
    )
    
    node1.add_observation({
        "timestamp": "2026-06-01T09:00:00Z",
        "overall_condition": 3,
        "findings": [{"type": "crack", "code": "1300"}],
        "risk_level": 5.9
    })
    
    node1.add_observation({
        "timestamp": "2026-06-15T09:00:00Z",
        "overall_condition": 4,
        "findings": [{"type": "crack", "code": "1300"}, {"type": "corrosion", "code": "1100"}],
        "risk_level": 7.2
    })
    
    twin.add_node(node1)
    
    # Add another node
    node2 = DigitalTwinNode(
        goid="BYG-FAC-WIN-GLA-0001",
        object_type="window_glass",
        domain="BYG",
        system="FAC",
        subsystem="WIN",
        lat=59.3300,
        lng=18.0700,
        condition=2,
        risk_level=2.2
    )
    
    node2.add_observation({
        "timestamp": "2026-06-10T10:00:00Z",
        "overall_condition": 2,
        "findings": [{"type": "dirt_accumulation", "code": "2100"}],
        "risk_level": 2.2
    })
    
    twin.add_node(node2)
    
    # Statistics
    print("=== Stockholm Digital Twin ===")
    print(json.dumps(twin.to_dict(), indent=2))
    
    # Find degrading
    print("\n=== Degrading Objects ===")
    degrading = twin.find_degrading()
    for obj in degrading:
        print(f"{obj['goid']}: {obj['trend']['trend']} (risk: {obj['risk_level']})")
    
    # Find by location
    print("\n=== Objects near city center ===")
    nearby = twin.find_by_location(59.3293, 18.0686, 1.0)
    for node in nearby:
        print(f"{node.goid}: {node.object_type}")
    
    # Contradictions
    print("\n=== Contradictions ===")
    contradictions = twin.get_contradictions({
        "bridges_good_condition": 12
    })
    for c in contradictions:
        print(f"{c['type']}: {c['description']}")
    
    # GeoJSON
    print("\n=== GeoJSON ===")
    geojson = twin.export_geojson()
    print(f"Features: {len(geojson['features'])}")