# File: pyvista_dem_example_v2.py
# Description: A more robust version of the interactive DEM viewer example.

import numpy as np
import pyvista as pv
import sys # Import sys for checking version

def create_sample_dem_data(width=100, height=100):
    """Creates a sample NumPy array representing DEM data for demonstration."""
    # Create coordinates as float32 from the start to avoid UserWarning
    x = np.arange(-width / 2, width / 2, 1, dtype=np.float32)
    y = np.arange(-height / 2, height / 2, 1, dtype=np.float32)
    xx, yy = np.meshgrid(x, y)
    
    # Create some interesting terrain with hills and valleys
    z = (np.sin(np.sqrt(xx**2 + yy**2) * 0.3) * 20 +
         np.cos(xx * 0.2) * 15 +
         np.sin(yy * 0.3) * 10)
    
    # Add a larger mountain feature
    mountain_x, mountain_y = 20, -15
    mountain_height = 80
    mountain_spread = 25
    distance_from_peak = np.sqrt((xx - mountain_x)**2 + (yy - mountain_y)**2)
    mountain = mountain_height * np.exp(-(distance_from_peak**2 / (2 * mountain_spread**2)))
    z += mountain
    
    return z.astype(np.float32) # Ensure final Z is also float

def run_interactive_dem_viewer(dem_data_array: np.ndarray):
    """
    Creates and displays an interactive 3D scene from a DEM data array using PyVista.
    
    Args:
        dem_data_array (np.ndarray): 2D NumPy array of elevation values.
    """
    if not dem_data_array.any():
        print("Error: DEM data is empty.")
        return

    # 1. Create a 3D mesh from the 2D NumPy array.
    height, width = dem_data_array.shape
    x_coords = np.arange(width, dtype=np.float32)
    y_coords = np.arange(height, dtype=np.float32)
    x_grid, y_grid = np.meshgrid(x_coords, y_coords)
    z_grid = dem_data_array
    
    terrain_mesh = pv.StructuredGrid(x_grid, y_grid, z_grid)
    terrain_mesh.points[:, 2] *= 0.3
    
    # 2. Set up the 3D plotter/scene
    plotter = pv.Plotter(window_size=[1200, 800])
    plotter.add_mesh(terrain_mesh, cmap='terrain', show_edges=False)

    # 3. Define the observer (radar) position and direction
    observer_position = np.array([width / 2, height / 2, 150.0], dtype=np.float32)
    observer_direction = np.array([0.8, 0.8, -0.6], dtype=np.float32)
    observer_direction /= np.linalg.norm(observer_direction)

    # 4. Add visual objects to the scene
    plotter.add_mesh(
        pv.Sphere(radius=2, center=observer_position),
        color='red',
        label='Observer'
    )
    
    ray_length = 200
    ray_end_point = observer_position + observer_direction * ray_length
    plotter.add_mesh(
        pv.Tube(pointa=observer_position, pointb=ray_end_point, radius=0.5),
        color='cyan',
        label='Radar Beam'
    )

    # 5. Perform Ray Tracing - CHECK IF METHOD EXISTS
    if hasattr(plotter, 'ray_trace'):
        intersection_points, _ = plotter.ray_trace(observer_position, ray_end_point)
        
        if len(intersection_points) > 0:
            target_point = intersection_points[0]
            print(f"Radar beam intersects terrain at coordinates: {target_point}")
            plotter.add_mesh(
                pv.Sphere(radius=1.5, center=target_point),
                color='yellow',
                label='Target'
            )
        else:
            print("Radar beam does not intersect the terrain (pointing at the sky).")
    else:
        print("\n--- WARNING ---")
        print("Your version of PyVista is outdated and does not have the 'ray_trace' method.")
        print("Please upgrade with: pip install --upgrade pyvista")
        print("-----------------\n")


    # 6. Customize the scene and show it
    plotter.add_axes()
    plotter.add_legend()
    plotter.camera_position = 'xy'
    plotter.camera.zoom(1.5)
    
    print("\n--- Interactive Controls ---")
    print("Mouse Left: Rotate | Mouse Middle/Shift+Left: Pan | Mouse Right/Scroll: Zoom")
    print("Press 'q' to close the window.")

    try:
        # This call is blocking and opens the interactive window
        plotter.show()
    except Exception as e:
        print("\n---!!! FAILED TO LAUNCH 3D WINDOW !!!---")
        print(f"An error occurred during rendering: {e}")
        print("This is likely an OpenGL/Graphics Driver issue.")
        print("Please check the following:")
        print("1. Are you running this via Remote Desktop (RDP)? Try running it locally.")
        print("2. Are your graphics card drivers up to date? Visit NVIDIA/AMD/Intel's website.")
        print("3. If in a VM, is 3D acceleration enabled?")
        print("------------------------------------------")


if __name__ == '__main__':
    # Generate sample data
    dem_data = create_sample_dem_data(width=200, height=200)
    
    # Run the interactive viewer
    run_interactive_dem_viewer(dem_data)