import matplotlib.pyplot as plt
import numpy as np

# ASCII frames provided (stripping the border '|' characters)
frame_1_ascii = [
    "   +*@%*",
    "   * %**",
    "* **+++*",
    " * + ++*",
    "    ***",
    " *+ ****",
    "  ++****",
    "+++**+**"
]

frame_2_ascii = [
    "   * ***",
    "* * ++**",
    "* *++++*",
    " ++ +++*",
    "  + +*+*",
    "  + ****",
    "  + ****",
    "+++**+**"
]

# Mapping symbols to approximate distance in mm based on ~183mm buckets
# image_3d018b.png legend: @ is nearest, . is farthest, blank is no reading.
depth_mapping = {
    '@': 241,  # ~150 to 330 mm
    '%': 425,  # ~333 to 516 mm
    '#': 608,  # ~516 to 700 mm
    '*': 791,  # ~700 to 883 mm
    '+': 975,  # ~883 to 1066 mm
    '=': 1158, # ~1066 to 1250 mm
    '-': 1341, # ~1250 to 1433 mm
    ':': 1525, # ~1433 to 1616 mm
    '.': 1708, # ~1616 to 1800 mm
    ' ': np.nan # Blank (no reading)
}

def parse_frame(ascii_frame):
    """Converts the ASCII frame into a 2D numpy array of depths."""
    data = np.zeros((8, 8))
    for i, row in enumerate(ascii_frame):
        for j, char in enumerate(row):
            data[i, j] = depth_mapping.get(char, np.nan)
    return data

def plot_3d_bars(ax, data, title):
    """Generates a 3D bar plot for a given 8x8 depth grid."""
    x_data, y_data = np.meshgrid(np.arange(data.shape[1]), np.arange(data.shape[0]))
    
    # Flatten arrays
    x_data = x_data.flatten()
    y_data = y_data.flatten()
    z_data = np.zeros_like(x_data)
    
    # Extract depth values and filter out NaNs (blanks)
    depths = data.flatten()
    mask = ~np.isnan(depths)
    
    x_valid = x_data[mask]
    y_valid = y_data[mask]
    z_valid = z_data[mask]
    
    # We want nearer objects to have taller bars for better visibility
    # Subtracting depth from NEAR_MM (1800) to invert the height representation
    near_mm = 1800
    dz = near_mm - depths[mask] 
    
    dx = dy = 0.8 # Width and depth of the bars
    
    # Create the 3D bars
    bars = ax.bar3d(x_valid, y_valid, z_valid, dx, dy, dz, color='skyblue', shade=True)
    
    ax.set_title(title)
    ax.set_xlabel("X Zone")
    ax.set_ylabel("Y Zone")
    ax.set_zlabel("Proximity (Height = Closeness)")
    
    # Adjust axis limits to match the 8x8 grid
    ax.set_xlim(0, 8)
    ax.set_ylim(0, 8)
    ax.set_zlim(0, near_mm)
    
    # Invert Y axis to match the top-to-bottom reading of the text
    ax.invert_yaxis()

# Parse the ASCII data
data_frame1 = parse_frame(frame_1_ascii)
data_frame2 = parse_frame(frame_2_ascii)

# Setup the side-by-side plot
fig = plt.figure(figsize=(14, 6))

# Plot Frame 1
ax1 = fig.add_subplot(121, projection='3d')
plot_3d_bars(ax1, data_frame1, "Frame 1: Near Cluster Detected (Nearest: 333 mm)")

# Plot Frame 2
ax2 = fig.add_subplot(122, projection='3d')
plot_3d_bars(ax2, data_frame2, "Frame 2: Cluster Gone (Nearest: 780 mm)")

plt.tight_layout()
plt.show()