Atomizer/tools/adaptive-isogrid/tests/compare_triangle_vs_gmsh.py

"""
Compare Triangle (pure Delaunay) vs Gmsh (Frontal-Delaunay) for isogrid meshing.

Generates side-by-side visualizations showing:
- Triangle count efficiency
- Minimum angle distribution
- Boundary conformance quality
- Density heatmap overlay
- Rib pattern visual comparison
"""

import json
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
import sys
sys.path.insert(0, str(Path(__file__).parent.parent))

from src.brain import triangulation
from src.brain import triangulation_gmsh
from src.brain.geometry_schema import normalize_geometry_schema
from src.brain.density_field import evaluate_density_grid
from src.atomizer_study import DEFAULT_PARAMS


def compute_triangle_angles(vertices, triangles):
    """Compute all angles in all triangles (returns flat array of angles in degrees)."""
    angles = []
    for tri in triangles:
        p0, p1, p2 = vertices[tri]

        # Vectors
        v01 = p1 - p0
        v12 = p2 - p1
        v20 = p0 - p2

        # Angles using law of cosines
        def angle_from_vectors(va, vb):
            cos_angle = np.dot(-va, vb) / (np.linalg.norm(va) * np.linalg.norm(vb))
            return np.degrees(np.arccos(np.clip(cos_angle, -1.0, 1.0)))

        angles.append(angle_from_vectors(v20, v01))
        angles.append(angle_from_vectors(v01, v12))
        angles.append(angle_from_vectors(v12, v20))

    return np.array(angles)


def plot_comparison(geometry, params, output_dir):
    """Generate comparison plots for Triangle vs Gmsh."""
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Generate both triangulations
    print("Generating Triangle mesh...")
    tri_result = triangulation.generate_triangulation(geometry, params)
    tri_verts = tri_result['vertices']
    tri_tris = tri_result['triangles']

    print("Generating Gmsh mesh...")
    gmsh_result = triangulation_gmsh.generate_triangulation_gmsh(geometry, params)
    gmsh_verts = gmsh_result['vertices']
    gmsh_tris = gmsh_result['triangles']

    print(f"\nTriangle: {len(tri_tris)} triangles, {len(tri_verts)} vertices")
    print(f"Gmsh:     {len(gmsh_tris)} triangles, {len(gmsh_verts)} vertices")
    print(f"Reduction: {100 * (1 - len(gmsh_tris) / max(len(tri_tris), 1)):.1f}%")

    # Compute angle statistics
    tri_angles = compute_triangle_angles(tri_verts, tri_tris)
    gmsh_angles = compute_triangle_angles(gmsh_verts, gmsh_tris)

    print(f"\nTriangle angles: min={tri_angles.min():.1f}°, mean={tri_angles.mean():.1f}°, max={tri_angles.max():.1f}°")
    print(f"Gmsh angles:     min={gmsh_angles.min():.1f}°, mean={gmsh_angles.mean():.1f}°, max={gmsh_angles.max():.1f}°")

    # --- Plot 1: Side-by-side mesh comparison ---
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 7), dpi=160)

    # Triangle mesh
    ax1.triplot(tri_verts[:, 0], tri_verts[:, 1], tri_tris, 'b-', linewidth=0.5, alpha=0.8)
    ax1.set_title(f'Triangle (pure Delaunay)\n{len(tri_tris)} triangles, min angle={tri_angles.min():.1f}°', fontsize=11)
    ax1.set_xlabel('x [mm]')
    ax1.set_ylabel('y [mm]')
    ax1.set_aspect('equal')

    # Gmsh mesh
    ax2.triplot(gmsh_verts[:, 0], gmsh_verts[:, 1], gmsh_tris, 'g-', linewidth=0.5, alpha=0.8)
    ax2.set_title(f'Gmsh (Frontal-Delaunay)\n{len(gmsh_tris)} triangles, min angle={gmsh_angles.min():.1f}°', fontsize=11)
    ax2.set_xlabel('x [mm]')
    ax2.set_ylabel('y [mm]')
    ax2.set_aspect('equal')

    # Draw holes on both
    for ax in [ax1, ax2]:
        for hole in geometry.get('holes', []):
            if hole.get('is_circular', False) and 'center' in hole:
                cx, cy = hole['center']
                r = float(hole['diameter']) / 2.0
                circle = plt.Circle((cx, cy), r, color='red', fill=False, linewidth=1.5)
                ax.add_patch(circle)

    fig.tight_layout()
    fig.savefig(output_dir / 'mesh_comparison.png')
    plt.close(fig)
    print(f"\nSaved: {output_dir / 'mesh_comparison.png'}")

    # --- Plot 2: Angle distribution histograms ---
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5), dpi=160)

    ax1.hist(tri_angles, bins=30, color='blue', alpha=0.7, edgecolor='black')
    ax1.axvline(60, color='green', linestyle='--', linewidth=2, label='Equilateral (60°)')
    ax1.axvline(tri_angles.min(), color='red', linestyle='--', linewidth=1.5, label=f'Min={tri_angles.min():.1f}°')
    ax1.set_xlabel('Angle [degrees]')
    ax1.set_ylabel('Count')
    ax1.set_title('Triangle: Angle Distribution')
    ax1.legend()
    ax1.grid(True, alpha=0.3)

    ax2.hist(gmsh_angles, bins=30, color='green', alpha=0.7, edgecolor='black')
    ax2.axvline(60, color='green', linestyle='--', linewidth=2, label='Equilateral (60°)')
    ax2.axvline(gmsh_angles.min(), color='red', linestyle='--', linewidth=1.5, label=f'Min={gmsh_angles.min():.1f}°')
    ax2.set_xlabel('Angle [degrees]')
    ax2.set_ylabel('Count')
    ax2.set_title('Gmsh: Angle Distribution')
    ax2.legend()
    ax2.grid(True, alpha=0.3)

    fig.tight_layout()
    fig.savefig(output_dir / 'angle_distribution.png')
    plt.close(fig)
    print(f"Saved: {output_dir / 'angle_distribution.png'}")

    # --- Plot 3: Density heatmap overlay ---
    X, Y, eta = evaluate_density_grid(geometry, params, resolution=3.0)

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 7), dpi=160)

    # Triangle + heatmap
    ax1.pcolormesh(X, Y, eta, shading='auto', cmap='viridis', alpha=0.3, vmin=0, vmax=1)
    ax1.triplot(tri_verts[:, 0], tri_verts[:, 1], tri_tris, 'k-', linewidth=0.4, alpha=0.6)
    ax1.set_title('Triangle + Density Field')
    ax1.set_xlabel('x [mm]')
    ax1.set_ylabel('y [mm]')
    ax1.set_aspect('equal')

    # Gmsh + heatmap
    im = ax2.pcolormesh(X, Y, eta, shading='auto', cmap='viridis', alpha=0.3, vmin=0, vmax=1)
    ax2.triplot(gmsh_verts[:, 0], gmsh_verts[:, 1], gmsh_tris, 'k-', linewidth=0.4, alpha=0.6)
    ax2.set_title('Gmsh + Density Field')
    ax2.set_xlabel('x [mm]')
    ax2.set_ylabel('y [mm]')
    ax2.set_aspect('equal')

    # Colorbar
    fig.colorbar(im, ax=[ax1, ax2], label='Density η', shrink=0.8)

    fig.tight_layout()
    fig.savefig(output_dir / 'density_overlay.png')
    plt.close(fig)
    print(f"Saved: {output_dir / 'density_overlay.png'}")

    # --- Summary stats ---
    stats = {
        'triangle': {
            'num_triangles': int(len(tri_tris)),
            'num_vertices': int(len(tri_verts)),
            'min_angle': float(tri_angles.min()),
            'mean_angle': float(tri_angles.mean()),
            'max_angle': float(tri_angles.max()),
        },
        'gmsh': {
            'num_triangles': int(len(gmsh_tris)),
            'num_vertices': int(len(gmsh_verts)),
            'min_angle': float(gmsh_angles.min()),
            'mean_angle': float(gmsh_angles.mean()),
            'max_angle': float(gmsh_angles.max()),
        },
        'efficiency_gain_percent': float(100 * (1 - len(gmsh_tris) / max(len(tri_tris), 1))),
        'min_angle_improvement': float(gmsh_angles.min() - tri_angles.min()),
    }

    with open(output_dir / 'comparison_stats.json', 'w') as f:
        json.dump(stats, f, indent=2)
    print(f"Saved: {output_dir / 'comparison_stats.json'}")

    return stats


if __name__ == "__main__":
    # Test on small_plate_200x150
    geometry_file = Path(__file__).parent / 'test_geometries' / 'small_plate_200x150.json'

    with open(geometry_file) as f:
        raw_geometry = json.load(f)

    geometry = normalize_geometry_schema(raw_geometry)

    params = dict(DEFAULT_PARAMS)
    params.update({
        's_min': 12.0,
        's_max': 35.0,
        'w_frame': 8.0,
        'd_keep': 1.5,
        'R_0': 40.0,
        'R_edge': 15.0,
        'lloyd_iterations': 0,  # Disable Lloyd for fair comparison
    })

    output_dir = Path(__file__).parent / 'comparison_output'
    stats = plot_comparison(geometry, params, output_dir)

    print("\n=== Comparison Summary ===")
    print(f"Triangle reduction: {stats['efficiency_gain_percent']:.1f}%")
    print(f"Min angle improvement: {stats['min_angle_improvement']:.1f}°")
    print(f"\nTriangle: {stats['triangle']['num_triangles']} tris, min angle {stats['triangle']['min_angle']:.1f}°")
    print(f"Gmsh:     {stats['gmsh']['num_triangles']} tris, min angle {stats['gmsh']['min_angle']:.1f}°")