"""
Geographic Clustering Service for Order Assignment
Uses K-means clustering to group orders by kitchen location.
Enhanced with Geohash encoding for spatial learning (idea.txt data encoding).
"""

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 GeohashEncoder:
    """
    Geohash Encoding for spatial data.
    Converts lat/lon coordinates to grid cell strings for locality capture.

    Precision levels:
    - 4 chars: ~156km x 156km (regional)
    - 5 chars: ~39km x 19km (city-level)
    - 6 chars: ~4.9km x 4.9km (neighborhood)
    - 7 chars: ~1.2km x 609m (local zone)
    """

    BASE32 = "0123456789bcdefghjkmnpqrstuvwxyz"

    @classmethod
    def encode(cls, lat: float, lon: float, precision: int = 6) -> str:
        """
        Encode lat/lon to geohash string.

        Args:
            lat: Latitude (-90 to 90)
            lon: Longitude (-180 to 180)
            precision: Number of characters (4-12)

        Returns:
            Geohash string
        """
        if lat == 0 and lon == 0:
            return "unknown"

        lat_min, lat_max = -90.0, 90.0
        lon_min, lon_max = -180.0, 180.0

        hash_chars = []
        is_lon = True

        for _ in range(precision):
            char_bits = 0
            for _ in range(5):
                if is_lon:
                    mid = (lon_min + lon_max) / 2
                    if lon >= mid:
                        char_bits = (char_bits << 1) | 1
                        lon_min = mid
                    else:
                        char_bits = char_bits << 1
                        lon_max = mid
                else:
                    mid = (lat_min + lat_max) / 2
                    if lat >= mid:
                        char_bits = (char_bits << 1) | 1
                        lat_min = mid
                    else:
                        char_bits = char_bits << 1
                        lat_max = mid
                is_lon = not is_lon
            hash_chars.append(cls.BASE32[char_bits])

        return "".join(hash_chars)

    @classmethod
    def get_zone_from_geohash(cls, geohash: str) -> Dict[str, Any]:
        """
        Extract zone metadata from geohash for ML features.

        Returns:
            Dict with zone info: zone_id, precision, cell_size, etc.
        """
        if not geohash or geohash == "unknown":
            return {"zone_id": "unknown", "precision": 0}

        precision = len(geohash)

        # Approximate cell sizes (in km)
        lat_error = 180.0 / (2 ** (precision * 5 // 2)) / 2
        lon_error = 360.0 / (2 ** (precision * 5 // 2 + 1)) / 2

        # Cell size approximation
        cell_width_km = lat_error * 111  # 1 degree lat ≈ 111km
        cell_height_km = (
            lon_error * 111 * cos(11 * 3.14159 / 180)
        )  # Adjust for Coimbatore lat

        return {
            "zone_id": geohash,
            "precision": precision,
            "cell_width_km": round(cell_width_km, 3),
            "cell_height_km": round(cell_height_km, 3),
            "prefix_4": geohash[:4] if len(geohash) >= 4 else geohash,
            "prefix_5": geohash[:5] if len(geohash) >= 5 else geohash,
            "prefix_6": geohash[:6] if len(geohash) >= 6 else geohash,
        }


class ClusteringService:
    """Clusters orders geographically to enable balanced rider assignment."""

    def __init__(self):
        self.earth_radius_km = 6371
        self.geohash_encoder = GeohashEncoder()

    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 _encode_location_features(self, lat: float, lon: float) -> Dict[str, Any]:
        """
        Encode location with multiple spatial features (Data Encoding: Geohash + Distance).

        Returns:
            Dict with geohash, zone info, distance features
        """
        features = {}

        # Primary geohash encoding (6 chars = ~5km zone for Coimbatore)
        geohash_6 = self.geohash_encoder.encode(lat, lon, 6)
        features["geohash_6"] = geohash_6

        # Fine-grained geohash (7 chars = ~1.2km zone)
        geohash_7 = self.geohash_encoder.encode(lat, lon, 7)
        features["geohash_7"] = geohash_7

        # Coarse geohash for regional grouping (4 chars = ~156km)
        geohash_4 = self.geohash_encoder.encode(lat, lon, 4)
        features["geohash_4"] = geohash_4

        # Zone metadata
        features["zone_info"] = self.geohash_encoder.get_zone_from_geohash(geohash_6)

        return features

    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
        - geohash_6: geohash encoding of centroid (NEW)
        - zone_info: zone metadata (NEW)
        - spatial_features: all location encodings (NEW)
        """
        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]

            # ── CHAIN-MERGE FIX ──────────────────────────────────────────
            # Original bug: only checked against the SEED kitchen.
            # Example: A→B = 2km, B→C = 2km, A→C = 4km (max_radius = 3km).
            # Old code: C is NOT merged (A→C > 3km).
            # New code: iteratively re-check remaining kitchens against the
            # CURRENT centroid after each merge, so chains like A→B→C
            # are correctly collapsed into one cluster.
            # ─────────────────────────────────────────────────────────────
            changed = True
            while changed:
                changed = False
                # Recompute centroid of current cluster
                c_lats = [kitchen_coords[n][0] for n in cluster_kitchens if n in kitchen_coords]
                c_lons = [kitchen_coords[n][1] for n in cluster_kitchens if n in kitchen_coords]
                c_lat = sum(c_lats) / len(c_lats) if c_lats else k_lat
                c_lon = sum(c_lons) / len(c_lons) if c_lons else k_lon

                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(c_lat, c_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)
                        changed = True  # centroid shifted — re-scan remaining

            # 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

            # ENHANCED: Add geohash encoding features
            spatial_features = self._encode_location_features(
                centroid_lat, centroid_lon
            )

            clusters.append(
                {
                    "centroid": (centroid_lat, centroid_lon),
                    "orders": cluster_orders,
                    "kitchen_names": set(cluster_kitchens),
                    "total_orders": len(cluster_orders),
                    # NEW: Geohash encoding features
                    "geohash_6": spatial_features["geohash_6"],
                    "geohash_7": spatial_features["geohash_7"],
                    "geohash_4": spatial_features["geohash_4"],
                    "zone_info": spatial_features["zone_info"],
                    "spatial_features": spatial_features,
                }
            )

        # 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 with geohash encoding"
        )
        return clusters

    def _get_kitchen_name(self, order: Dict[str, Any]) -> str:
        """Extract kitchen name from order."""
        possible_keys = [
            "pickupcustomer",   # confirmed primary field in production orders
            "locationname",     # confirmed backup field in production orders
            "storename", "store_name",
            "restaurantname", "restaurant_name",
            "kitchenname", "kitchen_name",
            "partnername", "partner_name",
            "tenantname",
        ]
        for key in possible_keys:
            if key in order and order[key]:
                return str(order[key]).strip()
        return "Unknown"
