#!/usr/bin/env python3
"""
Visualize the circle packing arrangement from the best solution
"""
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
import json
import sys
import os

def visualize_circle_packing(results_dir):
    """
    Visualize the circle packing from a results directory
    
    Args:
        results_dir: Path to the results directory (e.g., results_with_eval_service_gen50_20260203_083049)
    """
    # Construct paths
    best_dir = os.path.join(results_dir, "best")
    metrics_path = os.path.join(best_dir, "results", "metrics.json")
    npz_path = os.path.join(best_dir, "results", "extra.npz")
    
    if not os.path.exists(metrics_path):
        print(f"Error: metrics.json not found at {metrics_path}")
        return
    
    # Load the metrics data
    with open(metrics_path, 'r') as f:
        metrics = json.load(f)
    
    # Parse the centers from the string
    centers_str = metrics['primary']['public']['centers_str']
    lines = centers_str.strip().split('\n')
    
    centers = []
    for line in lines:
        # Parse format: "  centers[i] = (x, y)"
        parts = line.split('=')[1].strip()
        x, y = parts.strip('()').split(',')
        centers.append([float(x), float(y)])
    
    centers = np.array(centers)
    n_circles = len(centers)
    
    print(f"Number of circles: {n_circles}")
    print(f"Reported sum of radii: {metrics['primary']['private']['reported_sum_of_radii']}")
    
    # Try to load radii from extra.npz if it exists
    radii = None
    if os.path.exists(npz_path):
        try:
            data = np.load(npz_path)
            if 'radii' in data:
                radii = data['radii']
                print(f"Loaded radii from extra.npz")
            else:
                print("Available keys in npz:", list(data.keys()))
        except Exception as e:
            print(f"Could not load radii from npz: {e}")
    
    # If radii not available, estimate them based on distances
    if radii is None:
        print("Estimating radii based on nearest neighbor distances...")
        radii = np.zeros(n_circles)
        for i in range(n_circles):
            # Distance to walls
            x, y = centers[i]
            r_max = min(x, y, 1-x, 1-y)
            
            # Distance to other circles
            for j in range(n_circles):
                if i != j:
                    dist = np.linalg.norm(centers[i] - centers[j])
                    r_max = min(r_max, dist / 2.0)
            
            radii[i] = r_max * 0.95  # Scale down slightly to avoid overlap
    
    print(f"Radii range: [{radii.min():.4f}, {radii.max():.4f}]")
    print(f"Calculated sum of radii: {radii.sum():.4f}")
    
    # Create the visualization
    fig, ax = plt.subplots(1, 1, figsize=(12, 12))
    
    # Draw the unit square boundary
    square = patches.Rectangle((0, 0), 1, 1, linewidth=3, 
                               edgecolor='black', facecolor='lightgray', alpha=0.2)
    ax.add_patch(square)
    
    # Draw each circle
    colors = plt.cm.tab20(np.linspace(0, 1, n_circles))
    for i in range(n_circles):
        circle = patches.Circle(centers[i], radii[i], 
                               linewidth=1.5, edgecolor='darkblue', 
                               facecolor=colors[i], alpha=0.6)
        ax.add_patch(circle)
        
        # Add circle number at center
        ax.text(centers[i][0], centers[i][1], str(i), 
               ha='center', va='center', fontsize=9, fontweight='bold',
               bbox=dict(boxstyle='round,pad=0.3', facecolor='white', alpha=0.7))
    
    # Set equal aspect ratio and limits
    ax.set_xlim(-0.05, 1.05)
    ax.set_ylim(-0.05, 1.05)
    ax.set_aspect('equal')
    ax.grid(True, alpha=0.3, linestyle='--')
    ax.set_xlabel('X', fontsize=14, fontweight='bold')
    ax.set_ylabel('Y', fontsize=14, fontweight='bold')
    ax.set_title(f'Circle Packing Visualization: {n_circles} circles\nSum of Radii = {radii.sum():.6f}', 
                fontsize=16, fontweight='bold', pad=20)
    
    # Add generation info if available
    if 'generation' in metrics:
        ax.text(0.98, 0.02, f"Generation: {metrics['generation']}", 
               transform=ax.transAxes, ha='right', va='bottom',
               fontsize=10, bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))
    
    # Save the figure to the results directory
    output_path = os.path.join(best_dir, "circle_packing_visualization.png")
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    print(f"\n✓ Visualization saved to: {output_path}")
    
    # Also save to analyze/outputs for convenience
    analyze_output_dir = os.path.join(os.path.dirname(__file__), "outputs")
    os.makedirs(analyze_output_dir, exist_ok=True)
    analyze_output_path = os.path.join(analyze_output_dir, "circle_packing_visualization.png")
    plt.savefig(analyze_output_path, dpi=300, bbox_inches='tight')
    print(f"✓ Also saved to: {analyze_output_path}")
    
    plt.show()


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python visualize_circle_packing.py <results_directory>")
        print("\nExample:")
        print("  python analyze/visualize_circle_packing.py examples/circle_packing/results/results_with_eval_service_gen50_20260203_083049")
        sys.exit(1)
    
    results_dir = sys.argv[1]
    visualize_circle_packing(results_dir)