main.py

"""
SONAR 2.0 - Archaeological Site Detection
Beautiful Interactive Interface with Full AOI Visualization
"""

import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')
import plotly.graph_objects as go
from pathlib import Path
import torch
import io
from PIL import Image
import warnings
import zipfile
from scipy.ndimage import sobel, gaussian_filter
from matplotlib import cm
import rasterio
from rasterio.transform import rowcol
import folium
import base64
warnings.filterwarnings('ignore')

from utils import (
    ResUNetAutoencoder, ResUNetEncoder, load_patches,
    load_model, load_kmeans_model, load_unified_probability_matrix
)

# ==============================================================================
# CONFIGURATION
# ==============================================================================

class Config:
    AUTOENCODER_PATH = Path('models/best_model_aoi.pth')
    ENCODER_DIM = 128
    IFOREST_PATH = Path('models/isolation_forest_model_128dim.pkl')
    KMEANS_PATH = Path('models/kmeans_model_128dim.pkl')
    GATE_MODEL_PKL = Path('models/gate_mlp_model.pkl')
    GATE_SCALER_PATH = Path('models/gate_scaler.pkl')
    
    DATA_BASE = Path('Test_dataset')
    PATCHES_DIR = Path('patches_final_file')
    UNIFIED_PROB_DIR = Path('test_unified_probablity_matrices_with_gate')
    
    MODELS_ZIP = Path('models.zip')
    PATCHES_ZIP = Path('patches_final_file.zip')
    UNIFIED_PROB_ZIP = Path('test_unified_probablity_matrices_with_gate.zip')
    DATASET_ZIP = Path('Test_dataset.zip')
    
    BATCH_SIZE = 32
    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    @staticmethod
    def extract_data_files():
        print("\n" + "="*50)
        print("Starting SONAR 2.0...")
        print("="*50)
        
        if Config.MODELS_ZIP.exists() and (not Path('models').exists() or not any(Path('models').iterdir())):
            print("Extracting models...")
            with zipfile.ZipFile(Config.MODELS_ZIP, 'r') as zip_ref:
                zip_ref.extractall('.')
        
        if Config.PATCHES_ZIP.exists() and (not Config.PATCHES_DIR.exists() or not any(Config.PATCHES_DIR.iterdir())):
            print("Extracting patches...")
            with zipfile.ZipFile(Config.PATCHES_ZIP, 'r') as zip_ref:
                zip_ref.extractall('.')
        
        if Config.UNIFIED_PROB_ZIP.exists() and (not Config.UNIFIED_PROB_DIR.exists() or not any(Config.UNIFIED_PROB_DIR.iterdir())):
            print("Extracting probability matrices...")
            with zipfile.ZipFile(Config.UNIFIED_PROB_ZIP, 'r') as zip_ref:
                zip_ref.extractall('.')
        
        if Config.DATASET_ZIP.exists() and not Config.DATA_BASE.exists():
            print("Extracting Test_dataset.zip...")
            with zipfile.ZipFile(Config.DATASET_ZIP, 'r') as zip_ref:
                zip_ref.extractall('.')
            print("✓ Test dataset extracted")
        
        print("Ready!\n")

config = Config()

# ==============================================================================
# DATA MANAGER
# ==============================================================================

class DataManager:
    def __init__(self):
        self.aoi_list = []
        self.current_aoi = None
        self.current_patches = None
        self.current_metadata = None
        self.current_unified_matrix = None
        self.full_dtm = None
        self.reference_transform = None
        self.reference_shape = None
        self.reference_bounds = None
        self.reference_crs = None
        
    def discover_aois(self):
        if not config.PATCHES_DIR.exists():
            return []
        patch_files = list(config.PATCHES_DIR.glob("AOI_*_all_patches.npz"))
        aoi_names = sorted([f.stem.replace('_all_patches', '') for f in patch_files])
        self.aoi_list = aoi_names
        return aoi_names
    
    def load_original_raster(self, aoi_name: str):
        """Load original DTM from Test_dataset"""
        try:
            meta_dir = config.DATA_BASE / aoi_name / 'meta'
            if not meta_dir.exists():
                return None
            
            tif_files = list(meta_dir.glob('*.tif'))
            if not tif_files:
                return None
            
            dtm_path = tif_files[0]
            print(f"✓ Loading DTM: {dtm_path.name}")
            
            with rasterio.open(dtm_path) as src:
                dtm = src.read(1).astype(np.float32)
                if src.nodata is not None:
                    dtm[dtm == src.nodata] = np.nan
                
                self.reference_transform = src.transform
                self.reference_shape = (src.height, src.width)
                self.reference_bounds = src.bounds
                self.reference_crs = src.crs
                
                return dtm
        except Exception as e:
            print(f"Error loading raster: {e}")
            return None
    
    def load_aoi(self, aoi_name: str):
        try:
            patches_file = config.PATCHES_DIR / f"{aoi_name}_all_patches.npz"
            self.current_patches, self.current_metadata = load_patches(patches_file)
            
            matrix_file = config.UNIFIED_PROB_DIR / f"{aoi_name}_unified_prob_matrix.npz"
            if matrix_file.exists():
                self.current_unified_matrix, _, _ = load_unified_probability_matrix(
                    aoi_name, config.UNIFIED_PROB_DIR
                )
            
            self.full_dtm = self.load_original_raster(aoi_name)
            
            if self.full_dtm is not None:
                max_row = max(m['row'] + 64 for m in self.current_metadata)
                max_col = max(m['col'] + 64 for m in self.current_metadata)
                dtm_height, dtm_width = self.full_dtm.shape
                
                if max_row > dtm_height or max_col > dtm_width:
                    print(f"⚠️ Patch bounds exceed raster, falling back to reconstruction")
                    self.full_dtm = None
            
            self.current_aoi = aoi_name
            return f"✅ {aoi_name} loaded ({len(self.current_patches)} patches)"
        except Exception as e:
            return f"❌ Error: {e}"
    
    def get_patch(self, patch_idx: int):
        if self.current_patches is None or patch_idx >= len(self.current_patches):
            return None, None
        return self.current_patches[patch_idx], self.current_metadata[patch_idx]
    
    def find_patch_at_pixel(self, row: int, col: int):
        for idx, meta in enumerate(self.current_metadata):
            patch_row, patch_col = meta['row'], meta['col']
            if (patch_row <= row < patch_row + 64 and patch_col <= col < patch_col + 64):
                return idx
        return None

data_manager = DataManager()

# ==============================================================================
# MAP GENERATION WITH MULTIPLE OVERLAYS (LIKE STREAMLIT VERSION)
# ==============================================================================

def create_interactive_map(aoi_name, threshold=0.5):
    """
    Create beautiful Folium map with multiple terrain overlays
    Matches the Streamlit visualization quality
    """
    
    if data_manager.full_dtm is None or data_manager.reference_bounds is None:
        # Fallback: create a simple map showing patches without terrain overlays
        if data_manager.current_patches is None:
            return "<div style='padding: 20px; text-align: center;'>⚠️ Load an AOI first</div>"
        
        # Calculate approximate center from patches
        if data_manager.current_metadata:
            rows = [m['row'] for m in data_manager.current_metadata]
            cols = [m['col'] for m in data_manager.current_metadata]
            center_row = (min(rows) + max(rows)) / 2
            center_col = (min(cols) + max(cols)) / 2
            
            # Use approximate lat/lon (this is a fallback)
            center_lat, center_lon = 0, 0  # Will be set properly below
            
            m = folium.Map(
                location=[center_lat, center_lon],
                zoom_start=14,
                tiles='OpenStreetMap'
            )
            
            # Add patch markers
            if data_manager.current_unified_matrix is not None:
                gate_channel = data_manager.current_unified_matrix[:, :, :, 4]
                
                for idx, meta in enumerate(data_manager.current_metadata):
                    patch_score = np.mean(gate_channel[idx])
                    
                    if patch_score >= threshold:
                        # For fallback, just use patch indices as approximate locations
                        popup_html = f"""
                        <b>🔴 Anomaly Detected</b><br>
                        <hr>
                        <b>Score:</b> {patch_score:.3f}<br>
                        <b>Patch ID:</b> {idx}<br>
                        <hr>
                        <i>⚠️ DTM not available - using patch coordinates</i>
                        """
                        
                        # This is approximate - will work better with actual DTM
                        lat_approx = meta['row'] / 1000.0
                        lon_approx = meta['col'] / 1000.0
                        
                        folium.CircleMarker(
                            location=[lat_approx, lon_approx],
                            radius=8,
                            popup=folium.Popup(popup_html, max_width=300),
                            tooltip=f"Patch {idx} - Score: {patch_score:.3f}",
                            color='red',
                            fill=True,
                            fillColor='orange',
                            fillOpacity=0.7,
                            weight=2
                        ).add_to(m)
            
            folium.LayerControl().add_to(m)
            
            return m._repr_html_()
        
        return "<div style='padding: 20px; text-align: center; background: #fff3cd; border-radius: 8px;'>⚠️ No DTM data available for this AOI. Map visualization limited.</div>"
    
    bounds = data_manager.reference_bounds
    center_lat = (bounds.bottom + bounds.top) / 2
    center_lon = (bounds.left + bounds.right) / 2
    
    # Create base map
    m = folium.Map(
        location=[center_lat, center_lon],
        zoom_start=14,
        tiles='OpenStreetMap'
    )
    
    # Add center marker
    folium.Marker(
        location=[center_lat, center_lon],
        popup=f'AOI Center<br>Lat: {center_lat:.6f}<br>Lon: {center_lon:.6f}',
        tooltip='AOI Center',
        icon=folium.Icon(color='blue', icon='info-sign')
    ).add_to(m)
    
    # Generate terrain overlays
    dtm = data_manager.full_dtm
    valid_mask = ~np.isnan(dtm)
    
    if valid_mask.any():
        dtm_filled = dtm.copy()
        dtm_filled[~valid_mask] = np.nanmedian(dtm)
        
        # ============================================================
        # LAYER 1: Local Relief Model (ARCHAEOLOGICAL GOLD!)
        # ============================================================
        dtm_smooth = gaussian_filter(dtm_filled, sigma=10)
        local_relief = dtm_filled - dtm_smooth
        
        relief_clipped = np.clip(local_relief, -2, 2)
        relief_norm = (relief_clipped + 2) / 4
        
        rdbu_cmap = cm.get_cmap('RdBu_r')
        relief_rgba = rdbu_cmap(relief_norm)
        relief_rgb = (relief_rgba[:, :, :3] * 255).astype(np.uint8)
        relief_rgb[~valid_mask] = [128, 128, 128]
        
        img_relief = Image.fromarray(relief_rgb, mode='RGB')
        buffered = io.BytesIO()
        img_relief.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()
        
        folium.raster_layers.ImageOverlay(
            image=f'data:image/png;base64,{img_str}',
            bounds=[[bounds.bottom, bounds.left], [bounds.top, bounds.right]],
            opacity=0.75,
            name='🏛️ Local Relief Model (Archaeological)',
            overlay=True,
            control=True,
            show=True  # DEFAULT ON
        ).add_to(m)
        
        # ============================================================
        # LAYER 2: Multi-Directional Hillshade
        # ============================================================
        dx = sobel(dtm_filled, axis=1) / 8.0
        dy = sobel(dtm_filled, axis=0) / 8.0
        slope = np.arctan(np.sqrt(dx**2 + dy**2))
        aspect = np.arctan2(-dy, dx)
        
        azimuths = [315, 45, 225, 135]
        altitude = 45
        hillshades = []
        
        for az_deg in azimuths:
            azimuth = np.radians(az_deg)
            alt_rad = np.radians(altitude)
            hs = (np.sin(alt_rad) * np.sin(slope) + 
                  np.cos(alt_rad) * np.cos(slope) * 
                  np.cos(azimuth - aspect))
            hillshades.append(hs)
        
        hillshade_multi = np.mean(hillshades, axis=0)
        hillshade_multi = np.clip(hillshade_multi, -1, 1)
        hillshade_multi = ((hillshade_multi + 1) / 2 * 255).astype(np.uint8)
        hillshade_multi[~valid_mask] = 128
        
        hillshade_multi_rgb = np.stack([hillshade_multi, hillshade_multi, hillshade_multi], axis=-1)
        
        # Add color tinting
        dtm_norm = (dtm - np.nanpercentile(dtm[valid_mask], 2)) / \
                   (np.nanpercentile(dtm[valid_mask], 98) - 
                    np.nanpercentile(dtm[valid_mask], 2))
        dtm_norm = np.clip(dtm_norm, 0, 1)
        
        terrain_cmap = cm.get_cmap('terrain')
        terrain_rgba = terrain_cmap(dtm_norm)
        terrain_rgb = (terrain_rgba[:, :, :3] * 255).astype(np.uint8)
        
        hillshade_multi_rgb = (hillshade_multi_rgb * 0.75 + terrain_rgb * 0.25).astype(np.uint8)
        hillshade_multi_rgb[~valid_mask] = [128, 128, 128]
        
        img_multi = Image.fromarray(hillshade_multi_rgb, mode='RGB')
        buffered_multi = io.BytesIO()
        img_multi.save(buffered_multi, format="PNG")
        img_str_multi = base64.b64encode(buffered_multi.getvalue()).decode()
        
        folium.raster_layers.ImageOverlay(
            image=f'data:image/png;base64,{img_str_multi}',
            bounds=[[bounds.bottom, bounds.left], [bounds.top, bounds.right]],
            opacity=0.7,
            name='🗻 Multi-Directional Hillshade',
            overlay=True,
            control=True,
            show=False
        ).add_to(m)
        
        # ============================================================
        # LAYER 3: Standard Terrain
        # ============================================================
        dtm_norm_basic = (dtm - np.nanpercentile(dtm[valid_mask], 2)) / \
                         (np.nanpercentile(dtm[valid_mask], 98) - 
                          np.nanpercentile(dtm[valid_mask], 2))
        dtm_norm_basic = np.clip(dtm_norm_basic, 0, 1)
        
        terrain_basic_rgba = terrain_cmap(dtm_norm_basic)
        terrain_basic_rgb = (terrain_basic_rgba[:, :, :3] * 255).astype(np.uint8)
        terrain_basic_rgb[~valid_mask] = [128, 128, 128]
        
        img_basic = Image.fromarray(terrain_basic_rgb, mode='RGB')
        buffered_basic = io.BytesIO()
        img_basic.save(buffered_basic, format="PNG")
        img_str_basic = base64.b64encode(buffered_basic.getvalue()).decode()
        
        folium.raster_layers.ImageOverlay(
            image=f'data:image/png;base64,{img_str_basic}',
            bounds=[[bounds.bottom, bounds.left], [bounds.top, bounds.right]],
            opacity=0.65,
            name='🌍 Standard Terrain',
            overlay=True,
            control=True,
            show=False
        ).add_to(m)
    
    # Add anomaly markers
    if data_manager.current_unified_matrix is not None:
        gate_channel = data_manager.current_unified_matrix[:, :, :, 4]
        
        for idx, meta in enumerate(data_manager.current_metadata):
            patch_score = np.mean(gate_channel[idx])
            
            if patch_score >= threshold:
                row, col = meta['row'], meta['col']
                center_row = row + 32
                center_col = col + 32
                x, y = rasterio.transform.xy(data_manager.reference_transform, center_row, center_col)
                
                if data_manager.reference_crs != 'EPSG:4326':
                    from rasterio.warp import transform as transform_coords
                    lon, lat = transform_coords(data_manager.reference_crs, 'EPSG:4326', [x], [y])
                    lat, lon = lat[0], lon[0]
                else:
                    lat, lon = y, x
                
                popup_html = f"""
                <b>🔴 Anomaly Detected</b><br>
                <hr>
                <b>Score:</b> {patch_score:.3f}<br>
                <b>Patch ID:</b> {idx}<br>
                <b>Location:</b><br>
                Lat: {lat:.6f}<br>
                Lon: {lon:.6f}<br>
                <hr>
                <i>Click map at this location to inspect</i>
                """
                
                folium.CircleMarker(
                    location=[lat, lon],
                    radius=8,
                    popup=folium.Popup(popup_html, max_width=300),
                    tooltip=f"Anomaly Score: {patch_score:.3f}",
                    color='red',
                    fill=True,
                    fillColor='orange',
                    fillOpacity=0.7,
                    weight=2
                ).add_to(m)
    
    # Add boundary
    folium.Rectangle(
        bounds=[[bounds.bottom, bounds.left], [bounds.top, bounds.right]],
        color='red',
        fill=False,
        weight=2,
        popup=f'{aoi_name} boundary'
    ).add_to(m)
    
    folium.LayerControl().add_to(m)
    
    return m._repr_html_()

# ==============================================================================
# PATCH VISUALIZATION
# ==============================================================================

def create_patch_viz(patch: np.ndarray, metadata: dict):
    """Beautiful 2D patch visualization"""
    channel_names = ['DTM', 'Slope', 'Roughness', 'NDVI', 'NDWI', 'Flow Acc', 'Flow Dir']
    
    fig, axes = plt.subplots(2, 4, figsize=(16, 8), facecolor='white')
    axes = axes.flatten()
    
    for i in range(7):
        ax = axes[i]
        data = patch[i]
        cmap = ['terrain', 'YlOrRd', 'viridis', 'RdYlGn', 'Blues', 'cividis', 'twilight'][i]
        
        im = ax.imshow(data, cmap=cmap, interpolation='bilinear')
        ax.set_title(channel_names[i], fontsize=12, fontweight='bold', pad=8)
        ax.axis('off')
        plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    
    axes[7].axis('off')
    
    plt.suptitle(f"Patch {metadata['patch_id']} | Row {metadata['row']} | Col {metadata['col']}", 
                 fontsize=16, fontweight='bold', y=0.98)
    plt.tight_layout()
    
    buf = io.BytesIO()
    plt.savefig(buf, format='png', dpi=150, bbox_inches='tight', facecolor='white')
    plt.close(fig)
    buf.seek(0)
    
    return Image.open(buf)

def create_3d_terrain(patch: np.ndarray, metadata: dict):
    """Enhanced 3D terrain visualization"""
    dtm = patch[0]
    dtm_clean = np.nan_to_num(dtm, nan=np.nanmedian(dtm))
    
    rows, cols = dtm.shape
    x, y = np.arange(cols), np.arange(rows)
    X, Y = np.meshgrid(x, y)
    
    fig = go.Figure(data=[go.Surface(
        z=dtm_clean, x=X, y=Y,
        colorscale='earth',
        showscale=True,
        lighting=dict(
            ambient=0.4,
            diffuse=0.8,
            fresnel=0.2,
            specular=0.3,
            roughness=0.5
        ),
        contours=dict(
            z=dict(
                show=True,
                usecolormap=True,
                highlightcolor="limegreen",
                project=dict(z=True)
            )
        )
    )])
    
    fig.update_layout(
        title=f"3D Terrain · Patch {metadata['patch_id']}",
        scene=dict(
            xaxis_title='X (pixels)',
            yaxis_title='Y (pixels)',
            zaxis_title='Elevation (m)',
            camera=dict(eye=dict(x=1.5, y=1.5, z=1.3)),
            aspectmode='manual',
            aspectratio=dict(x=1, y=1, z=0.5)
        ),
        height=600,
        margin=dict(l=0, r=0, t=40, b=0)
    )
    
    return fig

# ==============================================================================
# UI FUNCTIONS
# ==============================================================================

def load_aoi_and_generate_map(aoi_name, threshold):
    """Load AOI and generate beautiful map"""
    status = data_manager.load_aoi(aoi_name)
    
    if "✅" in status:
        map_html = create_interactive_map(aoi_name, threshold)
        
        # Generate statistics (with safe access)
        stats_html = f"""
        <div style='padding: 15px; background: #f0f7ff; border-radius: 8px; margin: 10px 0;'>
            <h3 style='margin-top: 0; color: #1976d2;'>📊 AOI Statistics</h3>
            <p><b>Total Patches:</b> {len(data_manager.current_patches)}</p>
        """
        
        if data_manager.reference_shape:
            stats_html += f"<p><b>Raster Shape:</b> {data_manager.reference_shape}</p>"
        
        if data_manager.reference_crs:
            stats_html += f"<p><b>CRS:</b> {data_manager.reference_crs}</p>"
        
        if data_manager.reference_bounds:
            center_lat = (data_manager.reference_bounds.bottom + data_manager.reference_bounds.top) / 2
            center_lon = (data_manager.reference_bounds.left + data_manager.reference_bounds.right) / 2
            stats_html += f"<p><b>Center:</b> {center_lat:.6f}, {center_lon:.6f}</p>"
        
        if data_manager.full_dtm is not None:
            stats_html += "<p><b>DTM Status:</b> ✅ Loaded</p>"
        else:
            stats_html += "<p><b>DTM Status:</b> ⚠️ Not available (using reconstructed patches)</p>"
        
        stats_html += "</div>"
        
        return status, map_html, stats_html, "", None, None
    
    return status, "<div style='padding: 20px;'>❌ Failed to load AOI</div>", "", "", None, None

def update_threshold(threshold):
    """Update anomaly detection threshold"""
    if data_manager.current_aoi:
        return create_interactive_map(data_manager.current_aoi, threshold)
    return "<div style='padding: 20px;'>⚠️ Load an AOI first</div>"

def handle_map_click(lat, lon):
    """Handle click on map - extract and visualize patch"""
    if data_manager.full_dtm is None:
        return "⚠️ Load an AOI first", None, None
    
    from rasterio.warp import transform as transform_coords
    
    if data_manager.reference_crs != 'EPSG:4326':
        x, y = transform_coords('EPSG:4326', data_manager.reference_crs, [lon], [lat])
        lon, lat = x[0], y[0]
    
    row, col = rowcol(data_manager.reference_transform, lon, lat)
    patch_idx = data_manager.find_patch_at_pixel(row, col)
    
    if patch_idx is None:
        return f"❌ No patch at ({row}, {col})", None, None
    
    patch, metadata = data_manager.get_patch(patch_idx)
    
    if patch is None:
        return "❌ Error loading patch", None, None
    
    # Generate visualizations
    img_2d = create_patch_viz(patch, metadata)
    fig_3d = create_3d_terrain(patch, metadata)
    
    # Get score
    score_text = ""
    if data_manager.current_unified_matrix is not None:
        gate_channel = data_manager.current_unified_matrix[:, :, :, 4]
        patch_score = np.mean(gate_channel[patch_idx])
        score_text = f" | <span style='color: {'red' if patch_score > 0.5 else 'green'}; font-weight: bold;'>Score: {patch_score:.3f}</span>"
    
    info = f"<div style='padding: 10px; background: #e8f5e9; border-radius: 5px;'><b>✅ Patch {patch_idx}</b> | ID: {metadata['patch_id']}{score_text}</div>"
    
    return info, img_2d, fig_3d

# ==============================================================================
# BUILD INTERFACE
# ==============================================================================

def build_interface():
    
    custom_css = """
    .gradio-container {
        max-width: 1400px !important;
    }
    .map-container {
        border: 2px solid #1976d2;
        border-radius: 8px;
        overflow: hidden;
    }
    """
    
    with gr.Blocks(theme=gr.themes.Soft(), css=custom_css, title="SONAR 2.0") as app:
        
        gr.Markdown("""
        # 🗺️ SONAR 2.0 - Archaeological Anomaly Detection System
        ### Interactive geospatial analysis with multi-layer terrain visualization
        """)
        
        with gr.Row():
            # Left sidebar - Controls
            with gr.Column(scale=1):
                gr.Markdown("### 🎛️ Control Panel")
                
                aoi_dropdown = gr.Dropdown(
                    choices=data_manager.aoi_list,
                    label="📍 Select Area of Interest (AOI)",
                    value=data_manager.aoi_list[0] if data_manager.aoi_list else None,
                    info="Choose an AOI to analyze"
                )
                
                threshold_slider = gr.Slider(
                    minimum=0.0, 
                    maximum=1.0, 
                    value=0.5, 
                    step=0.05,
                    label="🎯 Detection Threshold",
                    info="Higher = fewer, more confident detections"
                )
                
                load_btn = gr.Button("🚀 Load AOI & Generate Map", variant="primary", size="lg")
                
                status_box = gr.HTML(label="Status")
                stats_box = gr.HTML(label="Statistics")
                
                gr.Markdown("""
                ---
                ### 📖 How to Use
                
                1. **Select AOI** from dropdown
                2. **Click "Load AOI"** to generate map
                3. **Explore layers** using map controls:
                   - 🏛️ Local Relief (archaeological features)
                   - 🗻 Multi-directional hillshade
                   - 🌍 Standard terrain
                4. **Click anywhere** on the map to inspect
                5. **View 2D/3D** visualizations below
                
                🔴 **Red markers** = Detected anomalies  
                🔵 **Blue marker** = AOI center
                """)
            
            # Right side - Map
            with gr.Column(scale=3):
                gr.Markdown("### 🗺️ Interactive Map (Click to Inspect)")
                map_display = gr.HTML(
                    value="<div style='padding: 40px; text-align: center; background: #f5f5f5; border-radius: 8px;'>👈 Select an AOI and click 'Load AOI' to view map</div>",
                    elem_classes=["map-container"]
                )
                
                gr.Markdown("### 🔍 Manual Inspection")
                gr.Markdown("*Enter coordinates from map click (Lat/Lon):*")
                
                with gr.Row():
                    lat_input = gr.Number(label="Latitude", precision=6, scale=1)
                    lon_input = gr.Number(label="Longitude", precision=6, scale=1)
                    inspect_btn = gr.Button("🔎 Inspect Patch", variant="primary", scale=1)
                
                patch_info = gr.HTML(label="Patch Information")
        
        gr.Markdown("---")
        gr.Markdown("### 📊 Detailed Patch Analysis")
        
        with gr.Row():
            with gr.Column():
                gr.Markdown("#### 🖼️ Multi-Channel 2D View")
                patch_2d = gr.Image(label="2D Layer Analysis", type="pil")
            
            with gr.Column():
                gr.Markdown("#### 🏔️ 3D Terrain Model")
                terrain_3d = gr.Plot(label="Interactive 3D Visualization")
        
        # Event handlers
        load_btn.click(
            fn=load_aoi_and_generate_map,
            inputs=[aoi_dropdown, threshold_slider],
            outputs=[status_box, map_display, stats_box, patch_info, patch_2d, terrain_3d]
        )
        
        threshold_slider.change(
            fn=update_threshold,
            inputs=[threshold_slider],
            outputs=[map_display]
        )
        
        inspect_btn.click(
            fn=handle_map_click,
            inputs=[lat_input, lon_input],
            outputs=[patch_info, patch_2d, terrain_3d]
        )
        
        gr.Markdown("""
        ---
        ### 💡 Tips
        - Toggle between layers using the map control (top-right)
        - Adjust threshold to see more/fewer anomalies
        - Click directly on red markers for quick inspection
        - Use 3D view to assess terrain relief and features
        
        **Powered by SONAR 2.0** | Archaeological AI Detection System
        """)
    
    return app

# ==============================================================================
# MAIN
# ==============================================================================

if __name__ == "__main__":
    print("\n" + "="*60)
    print("🗺️  SONAR 2.0 - Archaeological Site Detection")
    print("="*60)
    
    config.extract_data_files()
    data_manager.discover_aois()
    
    print(f"\n📊 System Information:")
    print(f"   Device: {config.DEVICE}")
    print(f"   Available AOIs: {len(data_manager.aoi_list)}")
    print(f"   AOI Names: {', '.join(data_manager.aoi_list)}")
    print("="*60 + "\n")
    
    app = build_interface()
    
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,  # Set to True for public URL
        show_error=True,
        show_api=True
    )