#!/usr/bin/env python3
"""
Survey route optimization script.
Finds optimal paths through cave systems.
"""

import argparse
import sys
import json
import numpy as np

sys.path.insert(0, '.')

try:
    from cavora.prediction import RouteOptimization
    from cavora.utils.logger import setup_logger
except ImportError:
    print("Error: CAVORA package not found")
    sys.exit(1)


def main():
    parser = argparse.ArgumentParser(description="Optimize survey routes")
    parser.add_argument("--nodes", "-n", required=True, help="Nodes file (JSON)")
    parser.add_argument("--start", "-s", required=True, help="Start node ID")
    parser.add_argument("--end", "-e", required=True, help="End node ID")
    parser.add_argument("--safety-weight", type=float, default=0.5, help="Safety vs distance weight (0-1)")
    parser.add_argument("--output", "-o", help="Output file (JSON)")
    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose output")
    
    args = parser.parse_args()
    
    # Setup logging
    logger = setup_logger("route_optimization", 
                         log_level="DEBUG" if args.verbose else "INFO")
    
    print("\n" + "="*60)
    print("CAVORA - Survey Route Optimization")
    print("="*60)
    
    # Load nodes data
    print(f"\nLoading nodes from: {args.nodes}")
    try:
        with open(args.nodes, 'r') as f:
            nodes_data = json.load(f)
        
        print(f"Loaded {len(nodes_data.get('nodes', []))} nodes")
        
    except Exception as e:
        print(f"Error loading nodes data: {e}")
        sys.exit(1)
    
    # Initialize route optimizer
    ro = RouteOptimization()
    
    # Add nodes
    for node in nodes_data.get('nodes', []):
        ro.add_node(
            node_id=node['id'],
            x=node['x'],
            y=node['y'],
            z=node['z'],
            stability=node.get('stability', 1.0)
        )
    
    # Add connections
    for conn in nodes_data.get('connections', []):
        ro.add_connection(
            node1_id=conn['from'],
            node2_id=conn['to'],
            distance=conn.get('distance')
        )
    
    print(f"\nNetwork built with {len(ro.nodes)} nodes")
    
    # Find optimal route
    print(f"\nFinding optimal route from {args.start} to {args.end}...")
    route = ro.find_optimal_route(args.start, args.end, safety_weight=args.safety_weight)
    
    if 'error' in route:
        print(f"Error: {route['error']}")
        sys.exit(1)
    
    # Display route
    print(f"\n{'='*60}")
    print(f"OPTIMAL ROUTE")
    print(f"{'='*60}")
    print(f"\nPath: {' → '.join(route['path'])}")
    print(f"Total Distance: {route['total_distance']} m")
    print(f"Average Stability: {route['avg_stability']}")
    print(f"Number of Nodes: {route['n_nodes']}")
    
    # Find safest route
    print(f"\n{'='*60}")
    print(f"SAFEST ROUTE")
    print(f"{'='*60}")
    
    safest = ro.find_safest_route(args.start, args.end, min_stability=0.6)
    
    if 'error' not in safest:
        print(f"\nPath: {' → '.join(safest['path'])}")
        print(f"Total Distance: {safest['total_distance']} m")
        print(f"Minimum Stability: {safest['min_stability']}")
        print(f"Average Stability: {safest['avg_stability']}")
    
    # Save results
    if args.output:
        output_data = {
            'timestamp': datetime.now().isoformat(),
            'input': {
                'start': args.start,
                'end': args.end,
                'safety_weight': args.safety_weight
            },
            'optimal_route': route,
            'safest_route': safest if 'error' not in safest else None
        }
        
        with open(args.output, 'w') as f:
            json.dump(output_data, f, indent=2)
        
        print(f"\nResults saved to: {args.output}")
    
    print("\n" + "="*60)


if __name__ == "__main__":
    from datetime import datetime
    main()
