model.py

"""
ML-Enhanced Route Optimization for CMU Landmarks

This model combines traditional routing algorithms with machine learning
to optimize routes based on user preferences and geographic constraints.
"""

import numpy as np
import json
import torch
import torch.nn as nn
from typing import List, Dict, Tuple, Any
from sklearn.preprocessing import StandardScaler


class RouteOptimizer(nn.Module):
    """Neural network for route optimization"""
    
    def __init__(self, input_dim: int = 4):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, 64)
        self.fc2 = nn.Linear(64, 32)
        self.fc3 = nn.Linear(32, 1)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.2)
        
    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.relu(self.fc2(x))
        x = self.dropout(x)
        x = self.fc3(x)
        return x


class MLRouteOptimizer:
    """
    ML-Enhanced Route Optimization (Fine-tuned approach)
    
    Uses a neural network to optimize routes considering user preferences
    and geographic constraints. This builds on the existing routing system.
    """
    
    def __init__(self, landmarks_data: List[Dict] = None, input_dim: int = 4):
        self.landmarks = landmarks_data or []
        self.id_to_idx = {lm['id']: i for i, lm in enumerate(self.landmarks)}
        self.model = RouteOptimizer(input_dim)
        self.scaler = StandardScaler()
        self.is_trained = False
        
    def _extract_route_features(self, landmark_idx: int, current_pos: int, 
                              time_budget: float, preferences: Dict) -> np.ndarray:
        """Extract features for route optimization"""
        lm = self.landmarks[landmark_idx]
        
        # Calculate distance to current position
        current_lm = self.landmarks[current_pos]
        lat1, lon1 = current_lm['geocoord']['lat'], current_lm['geocoord']['lon']
        lat2, lon2 = lm['geocoord']['lat'], lm['geocoord']['lon']
        
        # Simple distance calculation
        distance = np.sqrt((lat2 - lat1)**2 + (lon2 - lon1)**2)
        
        # Time cost
        time_cost = lm.get('time taken to explore', 30)
        
        # Preference alignment
        class_alignment = 0.0
        if 'selected_classes' in preferences:
            landmark_classes = lm.get('Class', [])
            for cls in landmark_classes:
                if cls in preferences['selected_classes']:
                    class_alignment += 0.2
        
        # Rating factor
        rating_factor = lm.get('rating', 0.0) / 5.0
        
        return np.array([distance, time_cost, class_alignment, rating_factor])
    
    def _distance(self, idx1: int, idx2: int) -> float:
        """Calculate distance between two landmarks"""
        lm1, lm2 = self.landmarks[idx1], self.landmarks[idx2]
        lat1, lon1 = lm1['geocoord']['lat'], lm1['geocoord']['lon']
        lat2, lon2 = lm2['geocoord']['lat'], lm2['geocoord']['lon']
        return np.sqrt((lat2 - lat1)**2 + (lon2 - lon1)**2)
    
    def optimize_route(self, selected_indices: List[int], start_idx: int,
                      time_budget: float, preferences: Dict) -> List[int]:
        """
        Optimize route using ML model
        
        This is a simplified version that can be extended with more sophisticated
        neural network training and fine-tuning.
        """
        if not selected_indices:
            return []
        
        # For now, use a heuristic approach with ML-inspired scoring
        # In a full implementation, this would use the trained neural network
        
        route_scores = []
        for idx in selected_indices:
            features = self._extract_route_features(idx, start_idx, time_budget, preferences)
            
            # Simple scoring function (in practice, this would be the neural network output)
            score = (
                -features[0] * 0.4 +  # Distance penalty
                features[2] * 0.3 +   # Class alignment bonus
                features[3] * 0.3     # Rating bonus
            )
            route_scores.append((idx, score))
        
        # Sort by score and use nearest neighbor for ordering
        sorted_indices = [idx for idx, _ in sorted(route_scores, key=lambda x: x[1], reverse=True)]
        
        # Apply nearest neighbor ordering
        ordered_route = []
        remaining = sorted_indices.copy()
        current = start_idx
        
        while remaining:
            next_idx = min(remaining, key=lambda x: self._distance(current, x))
            ordered_route.append(next_idx)
            remaining.remove(next_idx)
            current = next_idx
        
        return ordered_route
    
    def compare_routing_methods(self, selected_indices: List[int], 
                               start_idx: int,
                               preferences: Dict,
                               time_budget: float) -> Dict[str, Any]:
        """Compare ML-enhanced vs traditional routing methods"""
        
        # ML-enhanced route
        ml_route = self.optimize_route(selected_indices, start_idx, time_budget, preferences)
        
        # Traditional route (nearest neighbor)
        traditional_route = self._nearest_neighbor_route(selected_indices, start_idx)
        
        # Calculate metrics
        ml_distance = self._calculate_route_distance(ml_route)
        traditional_distance = self._calculate_route_distance(traditional_route)
        
        # Calculate preference satisfaction
        ml_preference_score = self._calculate_preference_satisfaction(ml_route, preferences)
        traditional_preference_score = self._calculate_preference_satisfaction(traditional_route, preferences)
        
        comparison = {
            "ml_route": {
                "route": ml_route,
                "distance": ml_distance,
                "preference_score": ml_preference_score,
                "info": {"method": "ml_enhanced", "route_length": len(ml_route)}
            },
            "traditional_route": {
                "route": traditional_route,
                "distance": traditional_distance,
                "preference_score": traditional_preference_score,
                "info": {"method": "traditional", "route_length": len(traditional_route)}
            },
            "comparison": {
                "distance_improvement": (traditional_distance - ml_distance) / max(traditional_distance, 1e-6),
                "preference_improvement": ml_preference_score - traditional_preference_score,
                "better_method": "ml" if ml_distance < traditional_distance else "traditional"
            }
        }
        
        return comparison
    
    def _nearest_neighbor_route(self, selected_indices: List[int], start_idx: int) -> List[int]:
        """Traditional nearest neighbor routing"""
        if not selected_indices:
            return []
        
        route = []
        remaining = selected_indices.copy()
        current = start_idx
        
        while remaining:
            next_idx = min(remaining, key=lambda x: self._distance(current, x))
            route.append(next_idx)
            remaining.remove(next_idx)
            current = next_idx
        
        return route
    
    def _calculate_route_distance(self, route: List[int]) -> float:
        """Calculate total distance for a route"""
        if len(route) <= 1:
            return 0.0
        
        total_distance = 0.0
        for i in range(len(route) - 1):
            total_distance += self._distance(route[i], route[i+1])
        
        return total_distance
    
    def _calculate_preference_satisfaction(self, route: List[int], 
                                         preferences: Dict) -> float:
        """Calculate how well a route satisfies user preferences"""
        if not route:
            return 0.0
        
        selected_classes = preferences.get('selected_classes', [])
        indoor_pref = preferences.get('indoor_pref')
        
        satisfaction_score = 0.0
        
        for idx in route:
            landmark = self.landmarks[idx]
            
            # Class preference satisfaction
            landmark_classes = landmark.get('Class', [])
            for cls in landmark_classes:
                if cls in selected_classes:
                    satisfaction_score += 0.2
            
            # Indoor/outdoor preference satisfaction
            io_type = landmark.get('indoor/outdoor', 'outdoor')
            if indoor_pref == 'indoor' and io_type == 'indoor':
                satisfaction_score += 0.1
            elif indoor_pref == 'outdoor' and io_type == 'outdoor':
                satisfaction_score += 0.1
            
            # Rating satisfaction
            rating = landmark.get('rating', 0.0)
            satisfaction_score += rating / 50.0  # Normalize rating contribution
        
        return satisfaction_score / len(route)  # Average satisfaction per landmark
    
    def save_model(self, filepath: str):
        """Save the model"""
        model_data = {
            'landmarks': self.landmarks,
            'id_to_idx': self.id_to_idx,
            'model_state': self.model.state_dict(),
            'scaler_mean': self.scaler.mean_.tolist() if hasattr(self.scaler, 'mean_') else None,
            'scaler_scale': self.scaler.scale_.tolist() if hasattr(self.scaler, 'scale_') else None,
            'is_trained': self.is_trained
        }
        
        with open(filepath, 'w') as f:
            json.dump(model_data, f)
    
    def load_model(self, filepath: str):
        """Load the model"""
        with open(filepath, 'r') as f:
            model_data = json.load(f)
        
        self.landmarks = model_data['landmarks']
        self.id_to_idx = model_data['id_to_idx']
        self.is_trained = model_data['is_trained']
        
        if model_data['scaler_mean']:
            self.scaler.mean_ = np.array(model_data['scaler_mean'])
            self.scaler.scale_ = np.array(model_data['scaler_scale'])
        
        # Load model state
        state_dict = {k: torch.tensor(v) for k, v in model_data['model_state'].items()}
        self.model.load_state_dict(state_dict)


def load_model_from_data(data_path: str) -> MLRouteOptimizer:
    """Load model from landmarks data"""
    with open(data_path, 'r') as f:
        landmarks = json.load(f)
    
    optimizer = MLRouteOptimizer(landmarks)
    return optimizer


# Example usage
if __name__ == "__main__":
    # Load landmarks data
    with open('data/landmarks.json', 'r') as f:
        landmarks = json.load(f)
    
    # Initialize optimizer
    optimizer = MLRouteOptimizer(landmarks)
    
    # Example route optimization
    selected_indices = [0, 5, 10, 15, 20]  # Example landmark indices
    start_idx = 0
    time_budget = 120.0
    preferences = {
        'selected_classes': ['Culture', 'Research'],
        'indoor_pref': 'indoor'
    }
    
    # Optimize route
    optimized_route = optimizer.optimize_route(selected_indices, start_idx, time_budget, preferences)
    
    # Compare methods
    comparison = optimizer.compare_routing_methods(selected_indices, start_idx, preferences, time_budget)
    
    print("ML Route:", comparison['ml_route']['route'])
    print("Traditional Route:", comparison['traditional_route']['route'])
    print("Distance Improvement:", f"{comparison['comparison']['distance_improvement']:.1%}")
    print("Preference Improvement:", f"{comparison['comparison']['preference_improvement']:.3f}")