"""
Geographic Clustering Service for Order Assignment
Uses K-means clustering to group orders by kitchen location.
"""

import logging
import numpy as np
from typing import List, Dict, Any, Tuple
from collections import defaultdict
from math import radians, cos, sin, asin, sqrt

logger = logging.getLogger(__name__)


class ClusteringService:
    """Clusters orders geographically to enable balanced rider assignment."""
    
    def __init__(self):
        self.earth_radius_km = 6371
    
    def haversine(self, lat1: float, lon1: float, lat2: float, lon2: float) -> float:
        """Calculate distance between two points in km."""
        lon1, lat1, lon2, lat2 = map(radians, [float(lon1), float(lat1), float(lon2), float(lat2)])
        dlon = lon2 - lon1
        dlat = lat2 - lat1
        a = sin(dlat/2)**2 + cos(lat1) * cos(lat2) * sin(dlon/2)**2
        c = 2 * asin(min(1.0, sqrt(a)))
        return c * self.earth_radius_km
    
    def get_kitchen_location(self, order: Dict[str, Any]) -> Tuple[float, float]:
        """Extract kitchen coordinates from order."""
        try:
            lat = float(order.get("pickuplat", 0))
            lon = float(order.get("pickuplon") or order.get("pickuplong", 0))
            if lat != 0 and lon != 0:
                return lat, lon
        except (ValueError, TypeError):
            pass
        return 0.0, 0.0
    
    def cluster_orders_by_kitchen(self, orders: List[Dict[str, Any]], max_cluster_radius_km: float = 3.0) -> List[Dict[str, Any]]:
        """
        Cluster orders by kitchen proximity.
        
        Returns list of clusters, each containing:
        - centroid: (lat, lon) of cluster center
        - orders: list of orders in this cluster
        - kitchen_names: set of kitchen names in cluster
        - total_orders: count
        """
        if not orders:
            return []
        
        # Group by kitchen location
        kitchen_groups = defaultdict(list)
        kitchen_coords = {}
        
        for order in orders:
            k_name = self._get_kitchen_name(order)
            k_lat, k_lon = self.get_kitchen_location(order)
            
            if k_lat == 0:
                # Fallback: use delivery location if pickup missing
                k_lat = float(order.get("deliverylat", 0))
                k_lon = float(order.get("deliverylong", 0))
            
            if k_lat != 0:
                kitchen_groups[k_name].append(order)
                kitchen_coords[k_name] = (k_lat, k_lon)
        
        # Now cluster kitchens that are close together
        clusters = []
        processed_kitchens = set()
        
        for k_name, k_orders in kitchen_groups.items():
            if k_name in processed_kitchens:
                continue
            
            # Start a new cluster with this kitchen
            cluster_kitchens = [k_name]
            cluster_orders = k_orders[:]
            processed_kitchens.add(k_name)
            
            k_lat, k_lon = kitchen_coords[k_name]
            
            # Find nearby kitchens to merge into this cluster
            for other_name, other_coords in kitchen_coords.items():
                if other_name in processed_kitchens:
                    continue
                
                other_lat, other_lon = other_coords
                dist = self.haversine(k_lat, k_lon, other_lat, other_lon)
                
                if dist <= max_cluster_radius_km:
                    cluster_kitchens.append(other_name)
                    cluster_orders.extend(kitchen_groups[other_name])
                    processed_kitchens.add(other_name)
            
            # Calculate cluster centroid
            lats = []
            lons = []
            for order in cluster_orders:
                lat, lon = self.get_kitchen_location(order)
                if lat != 0:
                    lats.append(lat)
                    lons.append(lon)
            
            if lats:
                centroid_lat = sum(lats) / len(lats)
                centroid_lon = sum(lons) / len(lons)
            else:
                centroid_lat, centroid_lon = k_lat, k_lon
            
            clusters.append({
                'centroid': (centroid_lat, centroid_lon),
                'orders': cluster_orders,
                'kitchen_names': set(cluster_kitchens),
                'total_orders': len(cluster_orders)
            })
        
        # Sort clusters by order count (largest first)
        clusters.sort(key=lambda x: x['total_orders'], reverse=True)
        
        logger.info(f"Created {len(clusters)} clusters from {len(kitchen_groups)} kitchens")
        return clusters
    
    def _get_kitchen_name(self, order: Dict[str, Any]) -> str:
        """Extract kitchen name from order."""
        possible_keys = ['storename', 'restaurantname', 'kitchenname', 'partnername', 'store_name']
        for key in possible_keys:
            if key in order and order[key]:
                return str(order[key]).strip()
        return "Unknown"