"""
Example 1: TSP (Traveling Salesman Problem)
Solve TSP using DPO's combinatoric optimization.
"""

import numpy as np
from dpo import dpo_solve_tsp, dpo
from dpo.core.problem import CombinatoricOptimizationProblem
import matplotlib.pyplot as plt


def example_tsp_quick():
    """Quick TSP example with small problem."""
    # Create 10-city TSP
    np.random.seed(42)
    n_cities = 10
    coords = np.random.rand(n_cities, 2) * 100
    
    # Distance matrix
    dist_matrix = np.zeros((n_cities, n_cities))
    for i in range(n_cities):
        for j in range(n_cities):
            dist_matrix[i, j] = np.linalg.norm(coords[i] - coords[j])
    
    # Solve with DPO
    print("="*60)
    print("TSP Example: 10 Cities")
    print("="*60)
    
    result = dpo_solve_tsp(
        distance_matrix=dist_matrix,
        preset='balanced',
        max_iterations=100,
        population_size=50
    )
    
    print(f"\nBest Tour Length: {result['best_fitness']:.2f}")
    print(f"Best Tour: {result['best_solution']['sequence']}")
    print(f"Total Evaluations: {result['total_evaluations']}")
    print(f"Elapsed Time: {result['elapsed_time']:.2f}s")
    
    return result


def example_tsp_custom():
    """TSP example using custom Problem class."""
    print("\n" + "="*60)
    print("TSP Example: Custom Problem Class")
    print("="*60)
    
    n_cities = 150
    np.random.seed(42)
    coords = np.random.rand(n_cities, 2) * 100
    
    # Precompute distance matrix
    dist_matrix = np.zeros((n_cities, n_cities))
    for i in range(n_cities):
        for j in range(n_cities):
            dist_matrix[i, j] = np.linalg.norm(coords[i] - coords[j])
    
    def tsp_objective(seq_dict):
        """Objective function for TSP."""
        tour = seq_dict['sequence']
        total_dist = sum(
            dist_matrix[tour[i], tour[(i+1) % len(tour)]]
            for i in range(len(tour))
        )
        
        # Calculate gap to optimal (if known, else use heuristic)
        greedy_dist = calculate_greedy_tsp(dist_matrix)
        gap = (total_dist - greedy_dist) / greedy_dist
        
        metrics = {
            'tour_length': total_dist,
            'gap_to_greedy': gap,
            'accuracy': 1.0 / (1.0 + total_dist / 100.0)  # Normalized accuracy
        }
        
        return total_dist, metrics
    
    # Use universal dpo() interface
    result = dpo(
        problem_type='combinatoric',
        objective_fn=tsp_objective,
        problem_size=n_cities,
        preset='thorough',
        max_iterations=150,
        population_size=60
    )
    
    print(f"\nBest Tour Length: {result['best_fitness']:.2f}")
    print(f"Gap to Greedy: {result['best_metrics'].get('gap_to_greedy', 'N/A'):.4f}")
    print(f"Elapsed Time: {result['elapsed_time']:.2f}s")
    
    return result


def calculate_greedy_tsp(dist_matrix):
    """Quick greedy TSP solution for reference."""
    n = len(dist_matrix)
    unvisited = set(range(1, n))
    current = 0
    total_dist = 0
    
    while unvisited:
        nearest = min(unvisited, key=lambda j: dist_matrix[current, j])
        total_dist += dist_matrix[current, nearest]
        unvisited.remove(nearest)
        current = nearest
    
    total_dist += dist_matrix[current, 0]  # Return to start
    return total_dist


if __name__ == "__main__":
    result1 = example_tsp_quick()
    result2 = example_tsp_custom()
    
    print("\n" + "="*60)
    print("Complete! Check results above.")
    print("="*60)
