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

"""
Visualize Gmsh triangulation on real sandbox geometries with adaptive density heatmap.

Creates varying hole weights to generate interesting density fields,
then shows how Gmsh Frontal-Delaunay responds.
"""

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.triangulation_gmsh import generate_triangulation
from src.brain.geometry_schema import normalize_geometry_schema
from src.brain.density_field import evaluate_density_grid
from src.brain.pocket_profiles import generate_pockets
from src.brain.profile_assembly import assemble_profile
from src.shared.arc_utils import typed_segments_to_polyline
from src.atomizer_study import DEFAULT_PARAMS
from shapely.geometry import Polygon


def add_varying_hole_weights(geometry, strategy='mixed'):
    """Add varying hole weights to create interesting density field."""
    num_holes = len(geometry.get('inner_boundaries', []))

    if num_holes == 0:
        return geometry  # No holes

    if strategy == 'mixed':
        # Create interesting pattern: critical holes near edges, low in center
        weights = []
        for i, hole in enumerate(geometry['inner_boundaries']):
            # Extract center from arc segments
            seg = hole['segments'][0]
            if seg['type'] == 'arc':
                cx, cy = seg['center']

                # Distance to nearest boundary determines importance
                # (simplified - just use y-coordinate for variation)
                # High weight near top/bottom, low in middle
                y_normalized = abs(cy + 200) / 400  # Normalize to 0-1
                weight = 0.15 + 0.7 * (1 - abs(0.5 - y_normalized) * 2)  # U-shaped
                weights.append(weight)

        # Assign weights
        for i, hole in enumerate(geometry['inner_boundaries']):
            hole['weight'] = weights[i]

    elif strategy == 'gradient':
        # Linear gradient from low to high
        for i, hole in enumerate(geometry['inner_boundaries']):
            hole['weight'] = 0.1 + 0.8 * (i / max(num_holes - 1, 1))

    elif strategy == 'random':
        # Random weights
        for hole in geometry['inner_boundaries']:
            hole['weight'] = np.random.uniform(0.1, 0.9)

    return geometry


def geometry_to_legacy_format(geometry):
    """Convert v2.0 sandbox geometry to legacy format for density_field module."""
    # Extract outer boundary coordinates
    outer_segments = geometry.get('outer_boundary', [])
    outer_coords = typed_segments_to_polyline(outer_segments, arc_pts=64)

    # Extract holes
    holes = []
    for i, inner in enumerate(geometry.get('inner_boundaries', [])):
        hole_segments = inner.get('segments', [])
        if not hole_segments:
            continue

        # For circular holes (single arc segment)
        if len(hole_segments) == 1 and hole_segments[0]['type'] == 'arc':
            seg = hole_segments[0]
            cx, cy = seg['center']
            radius = seg['radius']
            diameter = radius * 2.0

            # Sample circle
            theta = np.linspace(0, 2 * np.pi, 32, endpoint=False)
            boundary = [[cx + radius * np.cos(t), cy + radius * np.sin(t)] for t in theta]

            holes.append({
                'index': i,
                'center': [cx, cy],
                'diameter': diameter,
                'is_circular': True,
                'boundary': boundary,
                'weight': inner.get('weight', 0.5)
            })
        else:
            # Non-circular hole - use polyline
            hole_coords = typed_segments_to_polyline(hole_segments, arc_pts=32)
            poly = Polygon(hole_coords)
            centroid = poly.centroid
            area = poly.area
            diameter = 2 * np.sqrt(area / np.pi)  # Equivalent circle diameter

            holes.append({
                'index': i,
                'center': [centroid.x, centroid.y],
                'diameter': diameter,
                'is_circular': False,
                'boundary': hole_coords,
                'weight': inner.get('weight', 0.5)
            })

    return {
        'outer_boundary': outer_coords,
        'outer_boundary_typed': outer_segments,  # Keep typed for inset
        'holes': holes
    }


def visualize_sandbox(geometry_file, output_dir, params):
    """Generate comprehensive visualization for sandbox geometry."""
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Load and normalize geometry
    with open(geometry_file) as f:
        raw_geom = json.load(f)

    sandbox_id = raw_geom.get('sandbox_id', 'unknown')
    print(f"\n{'='*60}")
    print(f"Processing {sandbox_id}")
    print(f"{'='*60}")

    # Add varying hole weights
    geometry_v2 = add_varying_hole_weights(raw_geom, strategy='mixed')

    # Use normalize_geometry_schema which handles v2.0 conversion properly
    geometry = normalize_geometry_schema(geometry_v2)

    num_holes = len(geometry.get('holes', []))
    print(f"Holes: {num_holes}")
    if num_holes > 0:
        weights = [h.get('weight', 0.5) for h in geometry['holes']]
        print(f"Weight range: {min(weights):.2f} - {max(weights):.2f}")

    # Generate triangulation with Gmsh
    print("Generating Gmsh Frontal-Delaunay mesh...")
    try:
        tri_result = generate_triangulation(geometry, params)
        verts = tri_result['vertices']
        tris = tri_result['triangles']

        print(f"  > {len(tris)} triangles, {len(verts)} vertices")

        if len(tris) == 0:
            print("  ERROR: Meshing produced 0 triangles!")
            print(f"  Vertices: {len(verts)}")
            print("  Possible causes: Invalid geometry, size field too restrictive, or polygon errors")
            return None
    except Exception as e:
        print(f"  ERROR during triangulation: {e}")
        import traceback
        traceback.print_exc()
        return None

    # Compute triangle quality
    angles = []
    for tri in tris:
        p0, p1, p2 = verts[tri]
        v01 = p1 - p0
        v12 = p2 - p1
        v20 = p0 - p2

        def angle(va, vb):
            cos_a = np.dot(-va, vb) / (np.linalg.norm(va) * np.linalg.norm(vb) + 1e-12)
            return np.degrees(np.arccos(np.clip(cos_a, -1, 1)))

        angles.extend([angle(v20, v01), angle(v01, v12), angle(v12, v20)])

    angles = np.array(angles)
    print(f"  > Min angle: {angles.min():.1f} deg, Mean: {angles.mean():.1f} deg")

    # Generate pockets
    pockets = generate_pockets(tri_result, geometry, params)
    print(f"  > {len(pockets)} pockets generated")

    # Assemble final profile
    ribbed_plate = assemble_profile(geometry, pockets, params)

    # --- VISUALIZATION 1: Density heatmap overlay ---
    print("Creating density heatmap...")
    X, Y, eta = evaluate_density_grid(geometry, params, resolution=2.5)

    fig, ax = plt.subplots(figsize=(14, 10), dpi=160)

    # Heatmap background
    im = ax.pcolormesh(X, Y, eta, shading='auto', cmap='viridis', alpha=0.4, vmin=0, vmax=1)

    # Triangle mesh overlay
    ax.triplot(verts[:, 0], verts[:, 1], tris, 'k-', linewidth=0.4, alpha=0.7)

    # Draw holes
    for hole in geometry.get('holes', []):
        if hole.get('is_circular'):
            cx, cy = hole['center']
            r = hole['diameter'] / 2.0
            weight = hole.get('weight', 0.5)
            circle = plt.Circle((cx, cy), r, color='red', fill=False, linewidth=1.5)
            ax.add_patch(circle)
            # Weight label
            ax.text(cx, cy, f"{weight:.2f}", ha='center', va='center',
                   fontsize=8, color='white', weight='bold',
                   bbox=dict(boxstyle='round,pad=0.3', facecolor='black', alpha=0.7))

    ax.set_aspect('equal')
    ax.set_title(f'{sandbox_id}: Gmsh Frontal-Delaunay + Density Field\n'
                f'{len(tris)} triangles | Min angle {angles.min():.1f}° | {len(pockets)} pockets',
                fontsize=12, weight='bold')
    ax.set_xlabel('x [mm]')
    ax.set_ylabel('y [mm]')

    cbar = fig.colorbar(im, ax=ax, label='Density η (rib reinforcement)', shrink=0.7)

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

    # --- VISUALIZATION 2: Rib profile (pockets) ---
    print("Creating rib profile...")
    fig, ax = plt.subplots(figsize=(14, 10), dpi=160)

    # Plot outer boundary
    outer = np.array(geometry['outer_boundary'])
    ax.plot(np.r_[outer[:, 0], outer[0, 0]], np.r_[outer[:, 1], outer[0, 1]],
           'g-', linewidth=2.5, label='Sandbox boundary', zorder=5)

    # Plot pockets (material removed)
    for pocket in pockets:
        polyline = pocket.get('polyline', pocket.get('vertices', []))
        if len(polyline) < 3:
            continue
        pv = np.array(polyline)
        ax.fill(pv[:, 0], pv[:, 1], color='#ffcccc', alpha=0.4, edgecolor='#cc6677', linewidth=0.8)

    # Plot holes
    for hole in geometry.get('holes', []):
        hb = np.array(hole['boundary'])
        ax.plot(np.r_[hb[:, 0], hb[0, 0]], np.r_[hb[:, 1], hb[0, 1]],
               'b-', linewidth=1.2, label='_nolegend_')

    ax.set_aspect('equal')
    ax.set_title(f'{sandbox_id}: Final Rib Profile\n'
                f'{len(pockets)} pockets | Material: Ribs (white) + Frame (green)',
                fontsize=12, weight='bold')
    ax.set_xlabel('x [mm]')
    ax.set_ylabel('y [mm]')
    ax.legend(loc='upper right')

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

    # --- VISUALIZATION 3: Angle distribution ---
    print("Creating angle histogram...")
    fig, ax = plt.subplots(figsize=(10, 6), dpi=160)

    ax.hist(angles, bins=40, color='#1f77b4', alpha=0.75, edgecolor='black')
    ax.axvline(60, color='green', linestyle='--', linewidth=2, label='Equilateral (60°)')
    ax.axvline(angles.min(), color='red', linestyle='--', linewidth=1.5,
              label=f'Min={angles.min():.1f}°')
    ax.axvline(angles.mean(), color='orange', linestyle='--', linewidth=1.5,
              label=f'Mean={angles.mean():.1f}°')

    ax.set_xlabel('Angle [degrees]')
    ax.set_ylabel('Count')
    ax.set_title(f'{sandbox_id}: Triangle Angle Distribution (Gmsh Quality)', fontsize=12, weight='bold')
    ax.legend()
    ax.grid(True, alpha=0.3)

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

    return {
        'sandbox_id': sandbox_id,
        'num_triangles': len(tris),
        'num_vertices': len(verts),
        'num_pockets': len(pockets),
        'min_angle': float(angles.min()),
        'mean_angle': float(angles.mean()),
        'max_angle': float(angles.max()),
    }


if __name__ == "__main__":
    # Parameters optimized for isogrid
    params = dict(DEFAULT_PARAMS)
    params.update({
        's_min': 15.0,        # Min triangle spacing (dense areas)
        's_max': 45.0,        # Max triangle spacing (sparse areas)
        'w_frame': 3.0,       # Frame width (smaller for complex geometries)
        'd_keep': 0.8,        # Hole keepout multiplier
        'R_0': 40.0,          # Hole influence radius
        'R_edge': 15.0,       # Edge influence radius
        'alpha': 1.0,         # Hole influence weight
        'beta': 0.3,          # Edge influence weight
        'eta_0': 0.1,         # Baseline density
        't_min': 2.5,         # Min rib thickness
        't_0': 3.5,           # Nominal rib thickness
        'gamma': 1.0,         # Density-thickness coupling
        'r_f': 1.5,           # Pocket fillet radius
    })

    # Process both sandboxes
    sandbox_files = [
        Path(__file__).parent / 'geometry_sandbox_1.json',
        Path(__file__).parent / 'geometry_sandbox_2.json',
    ]

    output_dir = Path(__file__).parent / 'sandbox_results'

    results = []
    for geom_file in sandbox_files:
        if geom_file.exists():
            stats = visualize_sandbox(geom_file, output_dir, params)
            results.append(stats)

    # Summary
    print(f"\n{'='*60}")
    print("SUMMARY")
    print(f"{'='*60}")
    for stats in results:
        if stats:  # Skip None results
            print(f"\n{stats['sandbox_id']}:")
            print(f"  Triangles: {stats['num_triangles']}")
            print(f"  Pockets: {stats['num_pockets']}")
            print(f"  Min angle: {stats['min_angle']:.1f} deg")
            print(f"  Mean angle: {stats['mean_angle']:.1f} deg")

    # Save summary JSON
    with open(output_dir / 'summary.json', 'w') as f:
        json.dump(results, f, indent=2)

    print(f"\nAll results saved to: {output_dir}")