import numpy as np
import time
import os
import re
import json

# --- Helper Functions for Data Loading ---

def read_distance_matrix_from_file(file_path: str) -> np.ndarray:
    """
    Reads a pre-computed distance matrix from a .txt file.
    This function skips non-numeric lines, such as headers or the optimal distance line.
    """
    matrix_rows = []
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                # Skip empty lines or lines that do not start with a digit
                if not line or not line[0].isdigit():
                    continue
                
                try:
                    row = [float(num) for num in line.split()]
                    if row:
                        matrix_rows.append(row)
                except ValueError:
                    # Skip lines that cannot be fully converted to numbers
                    continue
    except FileNotFoundError:
        raise

    if not matrix_rows:
        raise ValueError(f"Could not parse any matrix data from {os.path.basename(file_path)}.")

    return np.array(matrix_rows)


def parse_optimal_distance_from_file(file_path: str) -> float | None:
    """
    Parses a pre-saved optimal distance from a file using regex.
    It looks for a line formatted like 'Total Distance: 1234.56'.
    """
    pattern = re.compile(r"Total Distance:\s*([\d.]+)")
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            for line in reversed(f.readlines()):
                match = pattern.search(line)
                if match:
                    return float(match.group(1))
    except (FileNotFoundError, IOError):
        return None
    return None

# --- Main TSP Solver (Nearest Neighbor Heuristic) ---

def solve_tsp_approximate(dist_matrix):
    """
    Find an approximate solution to the TSP for a given distance matrix
    using the Nearest Neighbor heuristic.
    """
    n = dist_matrix.shape[0]
    if n < 2:
        return np.array(range(n)), 0.0

    # Start the tour at the first city (city 0)
    start_city = 0
    unvisited_cities = set(range(n))
    tour = [start_city]
    unvisited_cities.remove(start_city)
    current_city = start_city
    total_distance = 0.0

    # Greedily select the nearest unvisited city until all cities are visited
    while unvisited_cities:
        # Find the nearest city from the current city among the unvisited ones
        nearest_city = min(
            unvisited_cities, key=lambda city: dist_matrix[current_city, city]
        )
        
        # Add the distance to the nearest city to the total distance
        total_distance += dist_matrix[current_city, nearest_city]
        
        # Update the current city to the newly visited city
        current_city = nearest_city
        
        # Add the new city to our tour
        tour.append(current_city)
        
        # Remove the city from the set of unvisited cities
        unvisited_cities.remove(current_city)

    # Add the distance from the last city back to the start city to close the loop
    total_distance += dist_matrix[current_city, start_city]

    return np.array(tour), total_distance


# --- Main Execution Block ---

if __name__ == "__main__":
    # ❗ **SET YOUR FOLDER PATH HERE**
    folder_path = 'all'
    
    if not os.path.isdir(folder_path):
        print(f"❌ Error: Folder not found at '{folder_path}'")
        exit()

    errors = []
    results_data = [] # List to store results for JSON output
    total_start_time = time.time()

    # Modified to look for .txt files
    print(f"🚀 Scanning for .txt files in '{os.path.abspath(folder_path)}'...")
    
    problem_files = sorted([f for f in os.listdir(folder_path) if f.endswith(".txt")])
    
    if not problem_files:
        # Modified to mention .txt files
        print("⚠️ No .txt files found in the specified directory.")
        exit()

    # --- Main Processing Loop ---
    for filename in problem_files:
        file_path = os.path.join(folder_path, filename)
        print("\n" + "-"*50)
        print(f"Processing: {filename}")
        
        try:
            # Step 1: Get the known optimal result from the file
            optimal_result = parse_optimal_distance_from_file(file_path)
            if optimal_result is None:
                print("   -> ⚠️ Warning: No saved optimal result found in file. Skipping.")
                continue

            # Step 2: Read the distance matrix directly from the file
            dist_matrix = read_distance_matrix_from_file(file_path)
            
            # Step 3: Run your heuristic to get a new result
            _, heuristic_result = solve_tsp_approximate(dist_matrix)
            
            # Step 4: Calculate the error and store it
            error = (heuristic_result - optimal_result) / optimal_result
            errors.append(error)
            
            # Store results in a dictionary for JSON output
            result_entry = {
                "filename": filename,
                "optimal_distance": optimal_result,
                "heuristic_distance": heuristic_result,
                "error": error
            }
            results_data.append(result_entry)
            
            print(f"   -> Optimal Result (from file): {optimal_result:.4f}")
            print(f"   -> Heuristic Result (this code): {heuristic_result:.4f}")
            print(f"   -> Error: {error:.2%}")

        except Exception as e:
            print(f"   -> ❌ An error occurred: {e}")

    # --- Save results to JSON file ---
    json_output_path = os.path.join(folder_path, 'nearest_neighbor_result.json')
    print("\n" + "-"*50)
    print(f"💾 Saving detailed results to {json_output_path}...")
    try:
        with open(json_output_path, 'w', encoding='utf-8') as f:
            json.dump(results_data, f, indent=4)
        print("   -> ✅ Successfully saved.")
    except Exception as e:
        print(f"   -> ❌ Failed to save JSON file: {e}")

    # --- Final Summary ---
    total_execution_time = time.time() - total_start_time
    
    print("\n" + "="*50)
    print("🎉 All files processed. Final Report:")
    print("="*50)
    
    if errors:
        average_error = np.mean(errors)
        print(f"📂 Files Processed: {len(errors)}")
        print(f"📊 **Average Error:** **{average_error:.2%}**")
    else:
        print("No files were successfully processed to calculate an average error.")
        
    print(f"⏱️ **Total Execution Time:** **{total_execution_time:.4f} seconds**")