#!/usr/bin/env python3
"""
Plot circle packing results from metrics.json file
"""
import json
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
import re
from pathlib import Path

def parse_centers_from_string(centers_str):
    """Parse center coordinates from the centers_str field"""
    centers = []
    # Pattern to match: centers[i] = (x, y)
    pattern = r'centers\[\d+\] = \(([0-9.]+), ([0-9.]+)\)'
    matches = re.findall(pattern, centers_str)
    
    for x, y in matches:
        centers.append((float(x), float(y)))
    
    return np.array(centers)

def calculate_radii_from_centers(centers, target_sum):
    """
    Estimate radii assuming roughly equal radii for all circles
    This is just for visualization - actual radii would come from the solution
    """
    n = len(centers)
    # Start with uniform distribution
    uniform_radius = target_sum / n
    
    # Calculate actual minimum distances and adjust
    radii = []
    for i in range(n):
        # Calculate minimum distance to boundaries
        x, y = centers[i]
        min_boundary_dist = min(x, y, 1.0 - x, 1.0 - y)
        
        # Calculate minimum distance to other circles (divided by 2 for radius)
        min_circle_dist = float('inf')
        for j in range(n):
            if i != j:
                dist = np.sqrt(np.sum((centers[i] - centers[j]) ** 2))
                min_circle_dist = min(min_circle_dist, dist)
        
        # Estimate radius as smaller of boundary constraint and half of nearest circle
        estimated_r = min(min_boundary_dist, min_circle_dist / 2.0, uniform_radius * 1.5)
        radii.append(estimated_r)
    
    # Scale to match target sum
    radii = np.array(radii)
    radii = radii * (target_sum / np.sum(radii))
    
    return radii

def plot_circle_packing(metrics_file):
    """Plot circle packing from metrics.json file"""
    # Load metrics
    with open(metrics_file, 'r') as f:
        data = json.load(f)
    
    # Extract information
    combined_score = data['combined_score']
    is_correct = data['correct']
    generation = data.get('generation', 'N/A')
    num_circles = data['primary']['public']['num_circles']
    centers_str = data['primary']['public']['centers_str']
    
    # Try to load actual radii from extra.npz
    extra_file = Path(metrics_file).parent / 'extra.npz'
    if extra_file.exists():
        extra_data = np.load(extra_file)
        centers = extra_data['centers']
        radii = extra_data['radii']
        print(f"Loaded {len(centers)} circles with actual radii from extra.npz")
        print(f"Sum of radii: {np.sum(radii):.6f} (target: {combined_score:.6f})")
    else:
        # Fallback: Parse centers from string and estimate radii
        centers = parse_centers_from_string(centers_str)
        print(f"Parsed {len(centers)} circle centers")
        radii = calculate_radii_from_centers(centers, combined_score)
        print(f"Estimated sum of radii: {np.sum(radii):.6f} (target: {combined_score:.6f})")
        print("WARNING: Using estimated radii. Actual radii not found in extra.npz")
    
    # Create figure with multiple subplots
    fig = plt.figure(figsize=(16, 6))
    
    # Main packing plot
    ax1 = plt.subplot(1, 3, 1)
    ax1.set_xlim(-0.05, 1.05)
    ax1.set_ylim(-0.05, 1.05)
    ax1.set_aspect('equal')
    ax1.set_title(f'Circle Packing (Generation {generation})', fontsize=14, fontweight='bold')
    ax1.set_xlabel('X')
    ax1.set_ylabel('Y')
    
    # Draw unit square
    square = patches.Rectangle((0, 0), 1, 1, linewidth=2, edgecolor='black', facecolor='none')
    ax1.add_patch(square)
    
    # Draw circles with color based on radius
    colors = plt.cm.viridis(radii / np.max(radii))
    for i, (center, radius) in enumerate(zip(centers, radii)):
        circle = patches.Circle(center, radius, linewidth=1, 
                              edgecolor='black', facecolor=colors[i], alpha=0.6)
        ax1.add_patch(circle)
        # Add circle number
        ax1.text(center[0], center[1], str(i), ha='center', va='center', 
                fontsize=6, fontweight='bold')
    
    # Add grid
    ax1.grid(True, alpha=0.3, linestyle='--')
    
    # Radii distribution plot
    ax2 = plt.subplot(1, 3, 2)
    ax2.hist(radii, bins=15, color='steelblue', edgecolor='black', alpha=0.7)
    ax2.set_xlabel('Radius')
    ax2.set_ylabel('Frequency')
    ax2.set_title('Radius Distribution', fontsize=12, fontweight='bold')
    ax2.grid(True, alpha=0.3)
    
    # Statistics text
    ax2.axvline(np.mean(radii), color='red', linestyle='--', linewidth=2, label=f'Mean: {np.mean(radii):.4f}')
    ax2.axvline(np.median(radii), color='green', linestyle='--', linewidth=2, label=f'Median: {np.median(radii):.4f}')
    ax2.legend()
    
    # Metrics summary
    ax3 = plt.subplot(1, 3, 3)
    ax3.axis('off')
    
    # Calculate some additional metrics for display
    total_area = np.sum(np.pi * radii ** 2)
    packing_efficiency = total_area / 1.0  # Unit square area = 1
    
    # Check for overlaps (approximate)
    overlaps = 0
    min_gap = float('inf')
    for i in range(len(centers)):
        for j in range(i+1, len(centers)):
            dist = np.sqrt(np.sum((centers[i] - centers[j]) ** 2))
            gap = dist - (radii[i] + radii[j])
            if gap < 0:
                overlaps += 1
            elif gap < min_gap:
                min_gap = gap
    
    # Check boundary violations
    boundary_violations = 0
    for i in range(len(centers)):
        x, y = centers[i]
        r = radii[i]
        if x - r < 0 or x + r > 1 or y - r < 0 or y + r > 1:
            boundary_violations += 1
    
    metrics_text = f"""
CIRCLE PACKING METRICS
{'='*40}

Primary Metrics:
  • Sum of Radii: {combined_score:.6f}
  • Number of Circles: {num_circles}
  • Valid Solution: {'✓ Yes' if is_correct else '✗ No'}
  • Generation: {generation}

Radius Statistics:
  • Mean Radius: {np.mean(radii):.6f}
  • Std Dev: {np.std(radii):.6f}
  • Min Radius: {np.min(radii):.6f}
  • Max Radius: {np.max(radii):.6f}
  • Median Radius: {np.median(radii):.6f}

Packing Efficiency:
  • Total Circle Area: {total_area:.6f}
  • Area Ratio: {packing_efficiency:.2%}
  • Avg Area per Circle: {total_area/num_circles:.6f}

Validation (Approximate):
  • Overlaps Detected: {overlaps}
  • Boundary Violations: {boundary_violations}
  • Min Gap (non-touching): {min_gap:.6f}

Spatial Distribution:
  • Center of Mass: ({np.mean(centers[:, 0]):.4f}, {np.mean(centers[:, 1]):.4f})
  • X-spread (std): {np.std(centers[:, 0]):.4f}
  • Y-spread (std): {np.std(centers[:, 1]):.4f}
    """
    
    ax3.text(0.1, 0.95, metrics_text, transform=ax3.transAxes, 
            fontsize=10, verticalalignment='top', fontfamily='monospace',
            bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.3))
    
    plt.tight_layout()
    
    # Save figure
    output_dir = Path(metrics_file).parent
    output_file = output_dir / 'circle_packing_visualization.png'
    plt.savefig(output_file, dpi=150, bbox_inches='tight')
    print(f"\nVisualization saved to: {output_file}")
    
    plt.show()

if __name__ == '__main__':
    metrics_file = '/home/tengxiao/pj/ShinkaEvolve/examples/circle_packing/results/results_full_gen200_period10_20260206_062935/best/results/metrics.json'
    plot_circle_packing(metrics_file)