#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Molecular Viewer - Modular Coordinator (Complete Feature Restoration)
======================================================================
"""
from __future__ import annotations
import math
import sys
import os
import time
from dataclasses import dataclass
from typing import Optional, Sequence, Tuple, List, Dict, Any

import numpy as np
try:
    import glfw
    from OpenGL.GL import *
    import imgui
    from imgui.integrations.glfw import GlfwRenderer
    HAS_VIZ = True
except ImportError:
    HAS_VIZ = False

# Subsystems
from .camera import OrbitCamera, rodrigues
from .scene import Scene
from .constants import WORLD_UP, CPK_COLORS, ATOMIC_RADII, CELL_EDGES
from .io.providers import StructureProvider
from .io.exporters import export_structure_xyz, export_structure
from .renderer.utils import look_at, perspective, ortho

# Analysis Modules
from .analysis.rdf import compute_rdf
from .analysis.coordination import compute_cn, compute_coordination_sphere, compute_bad
from .analysis.hbonds import compute_hbonds
from .analysis.geometry import compute_bounds, cell_corners, compute_rmsd
try:
    from .analysis.voronoi import compute_voronoi, VoronoiResult
    HAS_VORONOI = True
except ImportError:
    HAS_VORONOI = False

# Renderer Subsystem
from .renderer.sphere_renderer import SphereRenderer
from .renderer.bond_renderer import BondRenderer
from .renderer.line_renderer import LineRenderer
from .renderer.post_process import PostProcessor
from .renderer.picking_renderer import PickingRenderer

# UI Subsystem
from .ui.plots import draw_rdf_plot, draw_bad_plot, draw_energy_plots, draw_dos_plot, draw_eigenvalue_ladder

@dataclass
class RenderQualityState:
    target_fps: float = 60.0
    current_fps: float = 60.0
    skip_px: float = 0.5
    point_px: float = 3.0
    impostor_px: float = 12.0
    bond_mode: str = "cylinder" # cylinder, line, off
    max_visible_replicas: int = 125
    playback_low_quality: bool = True
    
    def adjust(self, fps: float):
        self.current_fps = fps
        if fps < self.target_fps * 0.7:
             # Reduce quality
             self.skip_px = min(self.skip_px + 0.1, 2.0)
             self.point_px = min(self.point_px + 0.5, 12.0)
             if self.bond_mode == "cylinder": self.bond_mode = "line"
        elif fps > self.target_fps * 1.2:
             # Increase quality
             self.skip_px = max(self.skip_px - 0.05, 0.4)
             self.point_px = max(self.point_px - 0.2, 3.0)
             if self.bond_mode == "line": self.bond_mode = "cylinder"

class Viewer:
    def __init__(self, provider: StructureProvider, 
                 start_index: int = 0, 
                 replicas: Tuple[int,int,int] = (0,0,0), 
                 win_size=(1280, 800)):
        
        if not HAS_VIZ:
            raise ImportError("Missing required visualization dependencies (GLFW, OpenGL, ImGui).")
            
        self.provider = provider
        self.index = start_index
        self.win_size = win_size
        
        # Core Subsystems
        self.scene = Scene()
        self.scene.set_replicas(*replicas)
        
        # Adaptive Quality
        self.quality = RenderQualityState()
        self.frame_times: List[float] = []
        
        # Camera & View State
        self.cam = OrbitCamera(np.zeros(3, dtype=np.float32), np.array([10, 10, 10], dtype=np.float32))
        self.is_ortho = True
        self.show_gui = True
        self.show_cell = True
        self.pretty_render = False
        self.bg_color = [0.05, 0.07, 0.09]
        self.atom_scale = 1.0
        self.bond_scale = 0.2
        self.show_bonds = False
        
        # Lights & Materials
        self.lights = [
            {'pos': [1.0, 1.0, 1.0], 'col': [1.0, 1.0, 1.0]},
            {'pos': [-1.0, -1.0, 0.5], 'col': [0.4, 0.5, 0.7]}
        ]
        self.selected_light_idx = 0
        self.materials = {
            'exposure': 1.0, 'headlight': 0.6, 
            'shadow_intensity': 0.6, 'shadow_softness': 2.762, 'shadow_bias': 0.002,
            'reflection_intensity': 1.0, 'reflection_gloss': 1.0,
            'fog_density': 0.002, 'gamma': 2.2,
            'ssao_strength': 1.0, 'ssao_radius': 0.1, 'ssao_bias': 0.0,
            'bloom_intensity': 0.489, 'bloom_threshold': 0.8, 'selection_glow': 0.688,
            'sheen': 0.993, 'clearcoat': 1.0,
            'pulse_speed': 4.0, 'pulse_amp': 0.0,
            'show_ssao': True, 'show_bloom': True, 'show_fxaa': False,
            'show_outlines': True, 'show_shadows': True
        }
        self.preset_options = ["Realistic", "Plastic", "Metallic", "Velvet", "Glassy", "Neon", "Clay", "Chrome", "Candy", "Frosted"]
        self.material_presets = {
            "Realistic": {'ambient': 0.4, 'specular': 0.6, 'shininess': 48.0, 'rim': 0.8, 'sheen': 0.0, 'clearcoat': 0.0, 'env': 0.3, 'gloss': 0.5},
            "Plastic":   {'ambient': 0.3, 'specular': 1.2, 'shininess': 64.0, 'rim': 1.0, 'sheen': 0.0, 'clearcoat': 0.5, 'env': 0.2, 'gloss': 0.3},
            "Metallic":  {'ambient': 0.1, 'specular': 4.0, 'shininess': 128.0, 'rim': 2.0, 'sheen': 0.0, 'clearcoat': 0.0, 'env': 1.0, 'gloss': 1.0},
            "Velvet":    {'ambient': 0.5, 'specular': 0.1, 'shininess': 16.0, 'rim': 1.5, 'sheen': 1.0, 'clearcoat': 0.0, 'env': 0.1, 'gloss': 0.1},
            "Glassy":    {'ambient': 0.1, 'specular': 2.0, 'shininess': 256.0, 'rim': 4.0, 'sheen': 0.0, 'clearcoat': 1.0, 'env': 0.8, 'gloss': 1.0},
            "Neon":      {'ambient': 5.0, 'specular': 0.1, 'shininess': 1.0, 'rim': 2.0, 'sheen': 1.0, 'clearcoat': 1.0, 'env': 0.0, 'gloss': 0.0},
            "Clay":      {'ambient': 0.6, 'specular': 0.1, 'shininess': 4.0, 'rim': 0.2, 'sheen': 0.0, 'clearcoat': 0.0, 'env': 0.0, 'gloss': 0.0},
            "Chrome":    {'ambient': 0.0, 'specular': 8.0, 'shininess': 512.0, 'rim': 3.0, 'sheen': 0.0, 'clearcoat': 0.0, 'env': 2.0, 'gloss': 1.0},
            "Candy":     {'ambient': 0.3, 'specular': 2.0, 'shininess': 128.0, 'rim': 1.0, 'sheen': 0.2, 'clearcoat': 1.5, 'env': 0.5, 'gloss': 0.8},
            "Frosted":   {'ambient': 0.4, 'specular': 0.4, 'shininess': 16.0, 'rim': 2.0, 'sheen': 1.0, 'clearcoat': 0.2, 'env': 0.2, 'gloss': 0.1},
        }
        self.element_materials = {} # Dict[Element, MaterialDict]
        self._mat_tbo_dirty = True
        self.light_rotating = False
        self.shadow_res_idx = 0 # 512
        self.global_preset_idx = 0
        
        # Playback State
        self.playback_playing = False
        self.playback_speed = 30.0
        self.last_time = glfw.get_time()
        self.time_acc = 0.0
        
        # Analysis State
        self.show_rdf = False
        self.rdf_mode_idx = 0
        self.rdf_mode = "bulk"
        self.show_cn = False
        self.show_hbonds = False
        self.show_voronoi = False
        self.show_z_profile = False
        self.show_selection_analytics = True
        self.force_full_pairwise = False
        self.rdf_rmax = 5.0
        self.rdf_bins = 200
        self.z_profile_x = np.zeros(1)
        self.z_profile_y = np.zeros(1)
        self.voro_vol_range = [0.0, 20.0]
        self.rdf_results = {}
        self.rdf_visibility = {} # Dict[str, bool]
        self.rdf_normalize = False # Peak normalization
        
        # Centralized Coloring Architecture
        self.color_mode = "Element" # Element, Z Position, Displacement, Neighbors, Coordination, Voronoi Volume, Selection, Index
        self.color_palette = "BlueRed" # BlueRed, Viridis, Magma, Hot, Grayscale
        self.color_invert = False
        self.color_range_auto = True
        self.color_range = [0.0, 10.0]
        self.color_z_fractional = False # Z in fractional coords?
        self.color_nb_cutoff = 3.0 # Neighbor count cutoff
        self.color_highlight_selection = True # Show selection on top
        self.displacement_ref_index = 0
        self.displacement_ref_pos = None # Reference positions
        
        # New Interactive Analysis Parameters
        self.bond_factor = 1.2 # Tunable covalent factor
        self.rdf_species_a = "All"
        self.rdf_species_b = "All"
        self.rdf_smooth = False
        self.rdf_manual_mode = True
        
        self.displacement_ref_index = 0
        self.analysis_scope = "Global" # Global, Selection, Region
        
        # UI state
        self.show_gui = True
        self.show_cell = True
        self.show_help = False
        
        # Clipping State
        self.clipping_enabled = False
        self.clipping_axis = 2 # 0:X, 1:Y, 2:Z
        self.clipping_range = [0.0, 1.0]
        self.clipping_fractional = False
        self.clipping_invert = False
        self.clipping_alpha = 0.0 # 0.0 means hidden, >0 means transparent
        
        # Volumetric State
        self.vol_data = None
        self.vol_enabled = False
        
        # Editing State
        self.edit_translate_delta = [0.0, 0.0, 0.0]
        self.edit_rotate_angle = 90.0
        self.edit_rotate_axis_idx = 2 # Z
        self.edit_new_element = "C"
        
        # Analysis Export State
        self.cn_cutoff = 3.0
        self.cn_results = None
        self.export_selected_only = False

        # Tight-Binding State
        from .analysis.tight_binding import TBParameters, TBResults
        self.tb_params = TBParameters()
        self.tb_results: Optional[Any] = None
        self.tb_enabled = False
        self.tb_status = ""
        self.tb_show_dos = False
        self.tb_show_ladder = False
        self.tb_color_mode = "None"
        self.tb_viz_mode = "None"
        self.tb_show_volume = False
        self.tb_vol_iso = 0.05
        self.vol_iso_level = 0.5
        self.vol_opacity = 0.5
        self.vol_show_slices = False
        self.vol_color = [0.3, 0.6, 1.0, 0.5]
        self.vol_mode_idx = 0 # 0: Isosurface, 1: MIP
        
        self.export_status = ""
        self.export_fmt_idx = 0
        self.export_formats = ["xyz", "extxyz", "cif", "pdb", "vasp", "aims"]
        self.export_prefix = "sage_export"
        self._export_path = "sage_export" # Added to fix AttributeError
        self._cube_path = "density.cube"  # Standard default
        self._comp_path = ""              # Simulation/Comparison path
        self.export_all_frames = False
        self.gif_frame_count = 10
        self.gif_random_selection = False
        
        # Navigation dual state
        self.cam1 = OrbitCamera(center=np.zeros(3, dtype=np.float32), 
                                offset=np.array([10.0, 10.0, 10.0], dtype=np.float32))
        self.cam2 = OrbitCamera(center=np.zeros(3, dtype=np.float32), 
                                offset=np.array([10.0, 10.0, 10.0], dtype=np.float32))
        self.cam = self.cam1 # Pointer to active camera for callbacks
        
        self.index1 = 0
        self.index2 = 0
        self.index = 0 # Pointer to active master index
        
        self.sync_cameras = True
        self.sync_frames = True
        self.split_view = False
        self.active_viewport = 1 # 1 for Left/A, 2 for Right/B
        
        # Interactions
        self.last_x, self.last_y = None, None
        self.orbiting = False
        self.panning = False
        self.box_selecting = False
        self.light_rotating = False
        self.show_axes = True
        self.show_origin_axes = False
        self.show_com = False
        self.panning = False
        self.rolling = False
        self.keys_down = set()
        self.shadow_res_idx = 3
        self.win_size = (1200, 800)
        self._trigger_selection_clear = False
        
        # Caches & Persistence
        self._bond_cache: Dict[Tuple[int, float], np.ndarray] = {}
        self._cached_element_colors = None
        self._gpu_status = {"pos": True, "sel": True, "mat": True}
        
        self.selection_dirty = True
        self._selection_cache = {
            "key": None,
            "dists_seq": [],
            "angles_seq": [],
            "dihedrals_seq": [],
            "dists_pw": None,
        }
        
        # Initialize OpenGL
        self._init_gl()
        self.shadow_res_idx = 0 # 512 by default
        self._init_renderers()
        
        # Initial Load
        self.provider2: Optional[Any] = None
        self.scene2: Optional[Scene] = None
        
        # Energy Trace Cache
        try:
            self.all_energies = np.array(self.provider.get_all_E(), dtype=np.float32)
        except Exception:
            self.all_energies = np.array([], dtype=np.float32)
            
        self._load_frame(self.index1)
        self._center_camera()

    def _init_gl(self):
        """Initializes GLFW window and ImGui context with proper callback chaining."""
        if not glfw.init(): raise RuntimeError("Failed to initialize GLFW")
        glfw.window_hint(glfw.CONTEXT_VERSION_MAJOR, 3)
        glfw.window_hint(glfw.CONTEXT_VERSION_MINOR, 3)
        glfw.window_hint(glfw.OPENGL_PROFILE, glfw.OPENGL_CORE_PROFILE)
        glfw.window_hint(glfw.OPENGL_FORWARD_COMPAT, GL_TRUE)
        glfw.window_hint(glfw.SAMPLES, 4)
        
        self.win = glfw.create_window(self.win_size[0], self.win_size[1], "Molecular Visualizer", None, None)
        if not self.win:
            glfw.terminate()
            raise RuntimeError("Failed to create GLFW window")
            
        glfw.make_context_current(self.win)
        glfw.swap_interval(1)
        
        # 1. Initialize ImGui first
        imgui.create_context()
        self.gui_impl = GlfwRenderer(self.win)
        
        # 2. Set custom callbacks AFTER ImGui to override/complement
        glfw.set_scroll_callback(self.win, self._on_scroll)
        glfw.set_cursor_pos_callback(self.win, self._on_mouse_move)
        glfw.set_mouse_button_callback(self.win, self._on_mouse_button)
        glfw.set_key_callback(self.win, self._on_key)
        
        glEnable(GL_DEPTH_TEST)
        glEnable(GL_MULTISAMPLE)

    def _init_renderers(self):
        self.rnd_sphere = SphereRenderer()
        self.rnd_bond = BondRenderer()
        self.rnd_line = LineRenderer()
        self.post_process = PostProcessor()
        self.rnd_pick = PickingRenderer()
        from .renderer.volume_renderer import VolumeRenderer
        self.rnd_volume = VolumeRenderer()

    def load_comparison(self, provider: Any = None):
        """Loads a second trajectory or clones the current one."""
        self.provider2 = provider if provider else self.provider
        self.scene2 = Scene()
        self.split_view = True
        self.index2 = self.index1
        # Copy camera state
        self.cam2.center = np.copy(self.cam1.center)
        self.cam2.offset = np.copy(self.cam1.offset)
        self.cam2.pan = np.copy(self.cam1.pan)
        self.cam2.ortho_scale = self.cam1.ortho_scale
        self._load_frame2(self.index2)

    def _get_selection_cache_key(self):
        """Generates a stable key for selection analytics caching."""
        return (
            self.index,
            tuple(sorted(self.scene.selected_atoms)),
            self.scene.N,
            self.scene.is_periodic,
            # If lattice is dynamic, we'd add it here too
        )

    def _compute_selection_analytics(self):
        """Lazy-computes distances, angles, and dihedrals for the current selection."""
        key = self._get_selection_cache_key()
        if self._selection_cache["key"] == key and not self.selection_dirty:
            return self._selection_cache

        selected = self.scene.selected_atoms
        pos = self.scene.pos
        lat = self.scene.lat
        is_periodic = self.scene.is_periodic

        def get_diff(pa, pb):
            d = pb - pa
            if is_periodic:
                d = d - np.round(d @ np.linalg.inv(lat)) @ lat
            return d

        dists_seq = []
        angles_seq = []
        dihedrals_seq = []

        # 1. Sequential Chain Metrics (O(N))
        for i in range(len(selected) - 1):
            p1, p2 = pos[selected[i]], pos[selected[i + 1]]
            dists_seq.append(np.linalg.norm(get_diff(p1, p2)))

        if len(selected) >= 3:
            for i in range(len(selected) - 2):
                p1, p2, p3 = pos[selected[i]], pos[selected[i + 1]], pos[selected[i + 2]]
                v1, v2 = get_diff(p2, p1), get_diff(p2, p3)
                cos_a = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2) + 1e-9)
                angles_seq.append(np.degrees(np.arccos(np.clip(cos_a, -1.0, 1.0))))

        if len(selected) >= 4:
            for i in range(len(selected) - 3):
                p1, p2, p3, p4 = [pos[selected[j]] for j in range(i, i + 4)]
                b1, b2, b3 = get_diff(p1, p2), get_diff(p2, p3), get_diff(p3, p4)
                n1 = np.cross(b1, b2); n1 /= (np.linalg.norm(n1) + 1e-9)
                n2 = np.cross(b2, b3); n2 /= (np.linalg.norm(n2) + 1e-9)
                m1 = np.cross(n1, b2 / (np.linalg.norm(b2) + 1e-9))
                x, y = np.dot(n1, n2), np.dot(m1, n2)
                dihedrals_seq.append(np.degrees(np.arctan2(y, x)))

        self._selection_cache.update({
            "key": key,
            "dists_seq": dists_seq,
            "angles_seq": angles_seq,
            "dihedrals_seq": dihedrals_seq,
            # Do NOT reset dists_pw here; it has its own demand flag
        })
        self.selection_dirty = False
        return self._selection_cache

    def _mark_dirty(self, pos=False, sel=False, mat=False):
        """Marks specific GPU data categories as needing synchronization."""
        if pos: self._gpu_status["pos"] = True
        if sel: self._gpu_status["sel"] = True
        if mat: self._gpu_status["mat"] = True

    def _load_frame1(self, index: int):
        if not self.provider: return
        self.index1 = int(index) % len(self.provider)
        try:
            p, l, e, els, c, r = self.provider.get(self.index1)
            self.scene.ingest(p, l, e, els, c, r)
            self._validate_selection(self.scene)
            self.scene.set_replicas(*self.scene.replicas)
            self.selection_dirty = True
            self._mark_dirty(pos=True, sel=True)
            self._sync_gpu(self.scene)
            # Live analysis & Coloring triggers
            self._update_coloring(self.scene)
            if self.show_rdf: self._calc_rdf()
            if self.show_z_profile: self._calc_z_profile()
        except Exception as err:
            print(f"Error loading frame1 {index}: {err}")

    def _load_frame2(self, index: int):
        if not self.provider2: return
        self.index2 = int(index) % len(self.provider2)
        try:
            p, l, e, els, c, r = self.provider2.get(self.index2)
            if self.scene2 is None: self.scene2 = Scene()
            self.scene2.ingest(p, l, e, els, c, r)
            self._validate_selection(self.scene2)
            self.scene2.set_replicas(*self.scene.replicas)
            self.selection_dirty = True
            self._mark_dirty(pos=True, sel=True)
            self._sync_gpu(self.scene2)
            self._mat_tbo_dirty = True # Force material sync on frame load
            # Live analysis & Coloring for secondary view
            self._update_coloring(self.scene2)
            if self.show_rdf: self._calc_rdf() 
            if self.show_z_profile: self._calc_z_profile()
        except Exception as err:
            print(f"Error loading frame2 {index}: {err}")

    def _load_frame(self, index: int):
        self._load_frame1(index)
        if self.sync_frames:
            self._load_frame2(index)

    def _validate_selection(self, scene: Scene):
        """Removes out-of-bounds selection indices when frame atom counts change."""
        if not scene.selected_atoms: return
        valid = [idx for idx in scene.selected_atoms if idx < scene.N]
        if len(valid) != len(scene.selected_atoms):
            scene.selected_atoms = valid
            scene.selected_set = set(valid)
            self.selection_dirty = True
            self._mark_dirty(sel=True)
            self._mat_tbo_dirty = True # Ensure GPU highlights update

    def _sync_gpu(self, scene: Optional[Scene] = None):
        """Updates GPU buffers lazily using dirty flags."""
        if scene is None: scene = self.scene
        
        # 1. Structural Sync (Positions, Colors, Replicas)
        if self._gpu_status.get("pos", True):
            colors = self._get_element_colors(scene, force_refresh=True)
            self.rnd_sphere.upload_data(scene.pos, scene.radii, colors)
            self.rnd_sphere.upload_replicas(scene.replica_offsets)
            self.rnd_bond.upload_replicas(scene.replica_offsets)
            
            if self.show_bonds:
                if scene.bond_count == 0:
                    self._calc_bonds_simple(scene)
                self.rnd_bond.upload_data(scene.pos, np.array(scene.bonds), 
                                          colors, lattice=scene.lat)
            self._gpu_status["pos"] = False

        # 2. Selection Sync
        if self._gpu_status.get("sel", True):
            self.rnd_sphere.upload_selection(scene.selected_atoms, scene.N)
            self._gpu_status["sel"] = False

    def _get_element_colors(self, scene: Scene, force_refresh: bool = False):
        """Builds a color array, with caching for Element mode."""
        # If in an analysis mode (Coordination, Z-pos, etc.), use the scene's dynamic colors
        if self.color_mode != "Element":
            return scene.colors
        
        if self._cached_element_colors is not None and not force_refresh:
            return self._cached_element_colors
            
        from .constants import CPK_COLORS
        colors = np.zeros((scene.N, 3), dtype=np.float32)
        for i, el in enumerate(scene.elements):
            if el not in self.element_materials:
                # Fallback / Initializer
                self.element_materials[el] = {
                    'col': list(CPK_COLORS.get(el, [0.5, 0.5, 0.5])),
                    'ambient': 0.4, 'specular': 0.6, 'shininess': 48.0, 'rim': 0.8,
                    'sheen': 0.0, 'clearcoat': 0.0, 'env': 0.3, 'gloss': 0.5
                }
            colors[i] = self.element_materials[el]['col']
            
        self._cached_element_colors = colors
        return colors

    def _center_camera(self):
        mn, mx = self.scene.get_full_bounds()
        center = (mn + mx) * 0.5
        diag = np.linalg.norm(mx - mn)
        self.cam.center = center.astype(np.float32)
        self.cam.offset = np.array([1.5*diag, 1.5*diag, diag], dtype=np.float32)
        self.cam.ortho_scale = diag * 1.5
        self.cam.snap_isometric()

    def loop(self):
        fps_hist = []
        while not glfw.window_should_close(self.win):
            t0 = glfw.get_time()
            glfw.poll_events()
            
            # --- Camera Physics ---
            tnow = glfw.get_time()
            if not hasattr(self, "_last_time"): self._last_time = tnow
            dt_step = tnow - self._last_time
            self._last_time = tnow
            
            self.cam1.update(dt_step)
            if self.split_view: self.cam2.update(dt_step)
            self._update_camera_keyboard(dt_step)
            
            self._update_playback()
            
            # Adaptive Quality Logic
            avg_fps = self.quality.current_fps
            if fps_hist: avg_fps = sum(fps_hist) / len(fps_hist)
            self.quality.adjust(avg_fps)
            
            # Rendering
            w, h = glfw.get_framebuffer_size(self.win)
            glViewport(0, 0, w, h)
            glClearColor(*self.bg_color, 1.0)
            glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
            
            self._render_scene(w, h)
            if self.show_gui: self._render_gui()
            
            glfw.swap_buffers(self.win)
            
            # FPS Tracking
            t1 = glfw.get_time(); dt = t1 - t0
            if dt > 0:
                fps_hist.append(1.0 / dt)
                if len(fps_hist) > 20: fps_hist.pop(0)
            
        self.gui_impl.shutdown()
        glfw.terminate()

    def _render_scene(self, w, h):
        """Orchestrates split-view or single-view rendering."""
        if self.split_view:
            # Clear once for both views
            bg = [0.03, 0.04, 0.06] if self.pretty_render else self.bg_color
            glClearColor(*bg, 1.0)
            glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
            
            # 1. Left Viewport (Scene 1, Cam 1)
            self._render_pass(self.scene, self.cam1, 0, 0, w//2, h)
            # 2. Right Viewport (Scene 2, Cam 2)
            s2 = self.scene2 if self.scene2 else self.scene
            self._render_pass(s2, self.cam2, w//2, 0, w//2, h)
        else:
            # Single Viewport
            bg = [0.03, 0.04, 0.06] if self.pretty_render else self.bg_color
            glClearColor(*bg, 1.0)
            glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
            self._render_pass(self.scene, self.cam1, 0, 0, w, h)

    def _update_clipping_planes(self, scene: Scene, cam: OrbitCamera):
        """Calculates dynamic near/far planes based on scene bounds for optimal depth precision."""
        mn, mx = scene.get_full_bounds()
        scene_center = 0.5 * (mn + mx)
        scene_diag = float(np.linalg.norm(mx - mn))
        eye = cam.eye()
        
        dist_to_center = float(np.linalg.norm(scene_center - eye))
        radius = max(1.0, 0.5 * scene_diag)
        
        cam.near = max(0.01, dist_to_center - 2.5 * radius)
        cam.far = dist_to_center + 2.5 * radius + 10.0 # Ensure far plane covers the scene

    def _render_pass(self, scene: Scene, cam: OrbitCamera, x: int, y: int, w: int, h: int):
        """Drives a complete render pass for a specific scene, camera, and viewport."""
        if scene.N == 0: return
        
        # Current Camera Matrices
        aspect = w / h if h > 0 else 1.0
        view = cam.get_view_matrix()
        # proj will be calculated later after near/far update
            
        # --- Performance Optimization: Visible Pair System ---
        # 1. Get visible supercell replicas (coarse culling)
        # Use a dummy proj for initial coarse culling, it's not critical for accuracy here
        # as it's just for broad-phase replica selection.
        dummy_proj = cam.get_projection_matrix(aspect, is_ortho=self.is_ortho)
        visible_offsets = self.rnd_sphere.get_visible_offsets(view, dummy_proj, scene.replica_offsets, scene.get_full_bounds())
        if not visible_offsets: return
        
        num_atoms = scene.N
        num_vis_reps = len(visible_offsets)
        mn, mx = scene.get_full_bounds()
        scene_diag = float(np.linalg.norm(mx - mn))
        
        # 2. Vectorized Culling for all atoms in visible replicas
        atom_indices = np.tile(np.arange(num_atoms, dtype=np.uint32), num_vis_reps)
        rep_indices = np.repeat(np.arange(num_vis_reps, dtype=np.uint32), num_atoms)
        offsets_arr = np.array(visible_offsets, dtype=np.float32)
        all_pos = scene.pos[atom_indices] + offsets_arr[rep_indices]
        
        # --- Interactive Clipping ---
        mask_clip = np.ones(all_pos.shape[0], dtype=bool)
        if self.clipping_enabled:
            if self.clipping_fractional:
                ext_lat = scene.lat
                inv_lat = np.linalg.inv(ext_lat)
                pos_frac = all_pos @ inv_lat
                vals = pos_frac[:, self.clipping_axis]
            else:
                vals = all_pos[:, self.clipping_axis]
            
            mask_clip = (vals >= self.clipping_range[0]) & (vals <= self.clipping_range[1])
            if self.clipping_invert: mask_clip = ~mask_clip
        
        # Dynamic Near/Far calculation for optimal depth precision
        self._update_clipping_planes(scene, cam)
        # Recalculate Projection Matrix AFTER updating near/far for coherence
        proj = cam.get_projection_matrix(aspect, is_ortho=self.is_ortho)
        
        # Frustum Culling (robust)
        vp = proj @ view
        p4 = np.ones((all_pos.shape[0], 4), dtype=np.float32)
        p4[:, :3] = all_pos
        p_clip = p4 @ vp.T
        w_vals = p_clip[:, 3]
        
        max_dist = cam.far
        
        mask_wide = (w_vals > 0.1) & (w_vals < max_dist)
        
        # Apply hard clip (alpha=0 mode) after mask_wide is defined
        if self.clipping_enabled and self.clipping_alpha <= 0.001:
            mask_wide &= mask_clip
        # mask_wide &= mask_clip # Apply clipping mask here - now handled above
        mask_wide &= (np.abs(p_clip[:, 0]) < (w_vals * 1.8))
        mask_wide &= (np.abs(p_clip[:, 1]) < (w_vals * 1.8))
        
        vis_indices_wide = np.where(mask_wide)[0]
        # Proceed even if nothing visible to allow ground/background to render
        
        # Adjust shadow resolution based on scene complexity to maintain FPS
        if vis_indices_wide.size > 200000:
             self.rnd_sphere.change_shadow_res(2048)
        else:
             self.rnd_sphere.change_shadow_res(4096)
        
        # Atom pairs for shadow pass (wide view)
        if vis_indices_wide.size > 0:
            v_pairs_atoms_wide = np.column_stack([
                atom_indices[vis_indices_wide], 
                rep_indices[vis_indices_wide]
            ]).astype(np.uint32)
        else:
            v_pairs_atoms_wide = np.zeros((0, 2), dtype=np.uint32)
        
        # Filter down for main pass (tight view + sorting)
        mask_main = (np.abs(p_clip[vis_indices_wide, 0]) < (w_vals[vis_indices_wide] * 1.5))
        mask_main &= (np.abs(p_clip[vis_indices_wide, 1]) < (w_vals[vis_indices_wide] * 1.5))
        vis_indices_main = np.where(mask_main)[0]
        
        # Depth Sorting (Front-to-Back for opaque, Back-to-Front for transparent)
        active_pos = all_pos[vis_indices_wide][vis_indices_main]
        cam_pos = cam.center + cam.offset
        dists_sq = np.sum((active_pos - cam_pos)**2, axis=1)
        
        # Back-to-front for transparency, Front-to-back for early-Z
        if self.clipping_enabled and self.clipping_alpha > 0.001:
            sort_idx = np.argsort(-dists_sq) 
        else:
            sort_idx = np.argsort(dists_sq) 
        
        if vis_indices_main.size > 0:
            v_pairs_atoms_main = np.column_stack([
                atom_indices[vis_indices_wide][vis_indices_main][sort_idx], 
                rep_indices[vis_indices_wide][vis_indices_main][sort_idx]
            ]).astype(np.uint32)
        else:
            v_pairs_atoms_main = np.zeros((0, 2), dtype=np.uint32)
        
        # Shadow Matrix calculation
        shadow_matrix = None
        if self.pretty_render and self.lights and self.materials.get('show_shadows', True):
            from .renderer.utils import look_at, ortho
            
            l_dir = np.array(self.lights[0]['pos'], dtype=np.float32)
            l_dir /= np.linalg.norm(l_dir) + 1e-8
            
            center = 0.5 * (mn + mx)
            # Position light far enough to cover the whole scene reliably
            l_eye = center + l_dir * (4.0 * scene_diag + 1.0)
            l_up = np.array([0,1,0], dtype=np.float32) if abs(l_dir[1]) < 0.9 else np.array([0,0,1], dtype=np.float32)
            
            l_view = look_at(l_eye, center, l_up)
            
            # Project 8 corners of the bounding box to light-space
            corners = np.array([
                [mn[0], mn[1], mn[2]], [mn[0], mn[1], mx[2]],
                [mn[0], mx[1], mn[2]], [mn[0], mx[1], mx[2]],
                [mx[0], mn[1], mn[2]], [mx[0], mn[1], mx[2]],
                [mx[0], mx[1], mn[2]], [mx[0], mx[1], mx[2]]
            ], dtype=np.float32)
            corners_h = np.c_[corners, np.ones(8, dtype=np.float32)]
            pts_ls = (l_view @ corners_h.T).T[:, :3]
            
            l_mn, l_mx = pts_ls.min(axis=0), pts_ls.max(axis=0)
            pad = 0.1 * scene_diag + 1.0
            
            l_proj = ortho(l_mn[0]-pad, l_mx[0]+pad, l_mn[1]-pad, l_mx[1]+pad, 
                           max(0.1, -l_mx[2]-pad), max(10.0, -l_mn[2]+pad))
            
            # Shadow Pass: Atoms
            glEnable(GL_POLYGON_OFFSET_FILL)
            glPolygonOffset(2.0, 8.0)
            self.rnd_sphere.draw_shadows(l_view, l_proj, v_pairs_atoms_wide, self.atom_scale)
            glDisable(GL_POLYGON_OFFSET_FILL)
            
            # Shadow Pass: Bonds
            if self.show_bonds and scene.bond_count > 0:
                b_s_idx, b_e_idx = scene.bonds[:, 0].astype(int), scene.bonds[:, 1].astype(int)
                b_shifts = scene.bonds[:, 2:5]
                p0_u = scene.pos[b_s_idx]
                p1_u = scene.pos[b_e_idx] + b_shifts @ scene.lat
                bc_u = (p0_u + p1_u) * 0.5
                
                b_idx_all = np.tile(np.arange(scene.bond_count, dtype=np.uint32), num_vis_reps)
                b_rep_all = np.repeat(np.arange(num_vis_reps, dtype=np.uint32), scene.bond_count)
                bc_w = bc_u[b_idx_all] + offsets_arr[b_rep_all]
                
                p4b = np.ones((bc_w.shape[0], 4), dtype=np.float32)
                p4b[:, :3] = bc_w
                pcb = p4b @ (proj @ view).T
                wb = pcb[:, 3]
                
                # Apply clipping to bonds (based on bond center)
                mask_b_clip = np.ones(bc_w.shape[0], dtype=bool)
                if self.clipping_enabled:
                    if self.clipping_fractional:
                        bc_frac = bc_w @ inv_lat
                        vals_b = bc_frac[:, self.clipping_axis]
                    else:
                        vals_b = bc_w[:, self.clipping_axis]
                    mask_b_clip = (vals_b >= self.clipping_range[0]) & (vals_b <= self.clipping_range[1])
                    if self.clipping_invert: mask_b_clip = ~mask_b_clip
                
                # Use a tighter wide margin (1.8x) for bonds in shadow pass to save FPS
                mask_b_wide = (wb > 0.1) & (np.abs(pcb[:, 0]) < (wb * 1.8)) & (np.abs(pcb[:, 1]) < (wb * 1.8))
                mask_b_wide &= mask_b_clip
                vis_b_wide = np.where(mask_b_wide)[0]
                if vis_b_wide.size > 0:
                    v_pairs_bonds_wide = np.column_stack([b_idx_all[vis_b_wide], b_rep_all[vis_b_wide]]).astype(np.uint32)
                    glEnable(GL_POLYGON_OFFSET_FILL)
                    glPolygonOffset(2.0, 8.0)
                    self.rnd_bond.draw_shadows(l_view, l_proj, v_pairs_bonds_wide, self.bond_scale,
                                               self.rnd_sphere.shadow_fbo, self.rnd_sphere.shadow_size)
                    glDisable(GL_POLYGON_OFFSET_FILL)
                    
                    mask_b_main = (np.abs(pcb[vis_b_wide, 0]) < (wb[vis_b_wide] * 1.5)) & (np.abs(pcb[vis_b_wide, 1]) < (wb[vis_b_wide] * 1.5))
                    self._v_pairs_bonds_main = v_pairs_bonds_wide[np.where(mask_b_main)[0]]
                else:
                    self._v_pairs_bonds_main = np.zeros((0, 2), dtype=np.uint32)
            
            bias = np.array([[0.5, 0.0, 0.0, 0.5], [0.0, 0.5, 0.0, 0.5], [0.0, 0.0, 0.5, 0.5], [0.0, 0.0, 0.0, 1.0]], dtype=np.float32)
            shadow_matrix = bias @ l_proj @ l_view
            
            # Apply shadow debug state to the renderer
            self.rnd_sphere.set_shadow_debug(self.materials.get('shadow_debug', False))

        # --- Rendering Pass ---
        if self.pretty_render:
            self.post_process.resize(w, h)
            self.post_process.begin() # Binds fbo_main
            glViewport(0, 0, w, h)    # Local FBO coordinates
        else:
            glBindFramebuffer(GL_FRAMEBUFFER, 0)
            glViewport(x, y, w, h)    # Global screen coordinates
        self._sync_gpu(scene)
        self._update_material_tbo()
        
        # Main Pass (Atoms)
        t_now = glfw.get_time()
        self.rnd_sphere.draw(view, proj, v_pairs_atoms_main, 
                             ortho_scale=cam.ortho_scale, viewport_height=h, 
                             lights=self.lights, materials=self.materials,
                             atom_scale=self.atom_scale, fov=cam.fov, is_ortho=self.is_ortho,
                             skip_px=self.quality.skip_px, point_px=self.quality.point_px,
                             pretty_render=self.pretty_render, shadow_matrix=shadow_matrix,
                             reflection_intensity=self.materials['reflection_intensity'],
                             reflection_gloss=self.materials['reflection_gloss'],
                             fog_density=self.materials['fog_density'], time=t_now)
        
        # Main Pass (Bonds)
        if self.show_bonds and scene.bond_count > 0:
            if not hasattr(self, '_v_pairs_bonds_main'):
                # Handle case where pretty_render is False or shadows are Off
                b_s_idx, b_e_idx = scene.bonds[:, 0].astype(int), scene.bonds[:, 1].astype(int)
                b_shifts = scene.bonds[:, 2:5]; p0_u = scene.pos[b_s_idx]
                p1_u = scene.pos[b_e_idx] + b_shifts @ scene.lat; bc_u = (p0_u + p1_u) * 0.5
                b_idx_all = np.tile(np.arange(scene.bond_count, dtype=np.uint32), num_vis_reps)
                b_rep_all = np.repeat(np.arange(num_vis_reps, dtype=np.uint32), scene.bond_count)
                bc_w = bc_u[b_idx_all] + offsets_arr[b_rep_all]
                p4b = np.ones((bc_w.shape[0], 4), dtype=np.float32); pcb = p4b @ (proj @ view).T; wb = pcb[:, 3]
                # Clipping for main bond pass
                mask_b_clip = np.ones(bc_w.shape[0], dtype=bool)
                if self.clipping_enabled:
                    if self.clipping_fractional:
                        bc_frac = bc_w @ inv_lat
                        vals_b = bc_frac[:, self.clipping_axis]
                    else:
                        vals_b = bc_w[:, self.clipping_axis]
                    mask_b_clip = (vals_b >= self.clipping_range[0]) & (vals_b <= self.clipping_range[1])
                    if self.clipping_invert: mask_b_clip = ~mask_b_clip
                
                m_b = (wb > 0.1) & (np.abs(pcb[:, 0]) < (wb * 1.5)) & (np.abs(pcb[:, 1]) < (wb * 1.5))
                m_b &= mask_b_clip
                vi_b = np.where(m_b)[0]
                self._v_pairs_bonds_main = np.column_stack([b_idx_all[vi_b], b_rep_all[vi_b]]).astype(np.uint32)

            if self._v_pairs_bonds_main.size > 0:
                self.rnd_bond.draw(view, proj, self._v_pairs_bonds_main, 
                                   lights=self.lights, materials=self.materials,
                                   bond_scale=self.bond_scale, pretty_render=self.pretty_render,
                                   shadow_matrix=shadow_matrix, shadow_map=self.rnd_sphere.shadow_tex,
                                   shadow_size=self.rnd_sphere.shadow_size,
                                   reflection_intensity=self.materials['reflection_intensity'],
                                   reflection_gloss=self.materials['reflection_gloss'],
                                    fog_density=self.materials['fog_density'], time=t_now,
                                    mat_tbo=self.rnd_sphere.mat_tbo,
                                    mat_tbo2=self.rnd_sphere.mat_tbo2)

        # Volumetric Pass (Standard or TB)
        if (self.vol_enabled and self.vol_data) or self.tb_show_volume:
            cam_pos = cam.center + cam.offset
            iso = self.tb_vol_iso if self.tb_show_volume else self.vol_iso_level
            self.rnd_volume.draw(view, proj, cam_pos, 
                                 iso_level=iso, 
                                 mode=self.vol_mode_idx, 
                                 col=tuple(self.vol_color))

        # Silhouette Mask Pass
        if self.pretty_render and self.materials.get('show_outlines', True):
            glBindFramebuffer(GL_FRAMEBUFFER, self.post_process.fbo_selection)
            glViewport(0, 0, w, h)
            glClearColor(0, 0, 0, 0)
            glClear(GL_COLOR_BUFFER_BIT)
            # Re-render ONLY selected objects as solid mask
            self.rnd_sphere.draw(view, proj, v_pairs_atoms_main, 
                                 ortho_scale=cam.ortho_scale, viewport_height=h, 
                                 lights=self.lights, materials=self.materials,
                                 atom_scale=self.atom_scale, fov=cam.fov, is_ortho=self.is_ortho,
                                 pretty_render=True, selection_only=True)
            if self.show_bonds and scene.bond_count > 0:
                self.rnd_bond.draw(view, proj, self._v_pairs_bonds_main,
                                   lights=self.lights, materials=self.materials,
                                   bond_scale=self.bond_scale, pretty_render=True,
                                   sel_full_tbo=self.rnd_sphere.sel_full_tbo, selection_only=True)
            
            if hasattr(self, '_v_pairs_bonds_main'):
                del self._v_pairs_bonds_main # Clean for next pass
        if self.show_cell and scene.is_periodic:
            from .analysis.geometry import cell_corners, CELL_EDGES
            corners = cell_corners(scene.lat)
            edges = corners[CELL_EDGES.flatten()]
            self.rnd_line.draw(view, proj, edges, len(scene.replica_offsets), 
                               self.rnd_sphere.rep_tbo, color=[1, 1, 1, 0.4])
        
        if self.pretty_render:
            # Restore screen viewport before final resolve
            glBindFramebuffer(GL_FRAMEBUFFER, 0)
            glViewport(x, y, w, h)
            # Map individual debug checkboxes to the unified uDebugMode
            debug_mode = 0
            if self.materials.get('debug_ao', False): debug_mode = 1
            elif self.materials.get('debug_depth', False): debug_mode = 2
            
            self.materials['debug_mode'] = debug_mode
            self.materials['proj'] = proj
            self.materials['projInv'] = np.linalg.inv(proj)
            self.post_process.end(materials=self.materials, 
                                  show_ssao=self.materials.get('show_ssao', True), 
                                  show_bloom=self.materials.get('show_bloom', True), 
                                  show_fxaa=self.materials.get('show_fxaa', True),
                                  near=cam.near, far=cam.far,
                                   viewport=(x, y, w, h),
                                   is_ortho=self.is_ortho)
            
            pass
        
        # 4. World Axes (at origin)
        if self.show_origin_axes:
            glEnable(GL_DEPTH_TEST)
            glDepthFunc(GL_LEQUAL)
            tripod = np.array([
                [0,0,0], [5,0,0], # X (Red)
                [0,0,0], [0,5,0], # Y (Green)
                [0,0,0], [0,0,5], # Z (Blue)
            ], dtype=np.float32)
            for i, color in enumerate([[1,0,0,1], [0,1,0,1], [0,0,1,1]]):
                self.rnd_line.draw(view, proj, tripod[2*i:2*i+2], 1, 
                                   self.rnd_sphere.rep_tbo, color=color)
        
        # 5. Center of Mass Indicator
        if self.show_com:
            self._draw_com(view, proj)
        
        # orientation HUD
        if self.show_axes:
            self._draw_axes_hud(cam, x, y, w, h)

    def _draw_axes_hud(self, cam: OrbitCamera, x, y, w, h):
        """Renders small RGB coordinate axes in the corner of the viewport."""
        glEnable(GL_DEPTH_TEST)
        glDepthFunc(GL_ALWAYS) # HUD on top
        
        # 1. Setup small viewport in corner
        hud_size = 80
        glViewport(x + 10, y + 10, hud_size, hud_size)
        
        # 2. View: Rotation only
        # We want the axes to rotate with the camera but stay centered
        view = cam.get_view_matrix()
        view[3, :3] = 0.0 # Remove translation
        view[3][3] = 1.0
        
        # 3. Projection: Small ortho box
        proj = ortho(-1.5, 1.5, -1.5, 1.5, -5.0, 5.0)
        
        # 4. Data
        tripod = np.array([
            [0,0,0], [1,0,0], # X (Red)
            [0,0,0], [0,1,0], # Y (Green)
            [0,0,0], [0,0,1], # Z (Blue)
        ], dtype=np.float32)
        
        # Render each axis with its color
        for i, color in enumerate([[1,0,0,1], [0,1,0,1], [0,0,1,1]]):
            self.rnd_line.draw(view, proj, tripod[2*i:2*i+2], 1, 
                               self.rnd_sphere.rep_tbo, color=color)
        
        # Restore viewport
        glViewport(x, y, w, h)

    def _draw_com(self, view, proj):
        """Renders a high-visibility crosshair at the center of mass."""
        com = self.scene.get_com()
        glDisable(GL_DEPTH_TEST) # Crosshair always on top
        
        s = 1.0 # crosshair size
        cross = np.array([
            [com[0]-s, com[1], com[2]], [com[0]+s, com[1], com[2]], # X
            [com[0], com[1]-s, com[2]], [com[0], com[1]+s, com[2]], # Y
            [com[0], com[1], com[2]-s], [com[0], com[1], com[2]+s], # Z
        ], dtype=np.float32)
        
        # Gold/Yellow color for COM
        color = [1.0, 0.8, 0.0, 1.0]
        self.rnd_line.draw(view, proj, cross, 1, self.rnd_sphere.rep_tbo, color=color)
        glEnable(GL_DEPTH_TEST)

    def reset_view(self):
        """Re-centers camera on the structural center or COM."""
        if self.scene.N == 0: return
        
        mn, mx = self.scene.get_full_bounds()
        center = self.scene.get_com() # Center on COM specifically
        dist = np.linalg.norm(mx - mn) * 1.2
        if dist < 5.0: dist = 15.0
        
        self.cam1.center = center
        self.cam1.offset = np.array([dist, dist, dist], dtype=np.float32)
        self.cam1.pan = np.zeros(2, dtype=np.float32)
        
        if self.split_view:
            self.cam2.center = np.copy(center)
            self.cam2.offset = np.copy(self.cam1.offset)
            self.cam2.pan = np.zeros(2, dtype=np.float32)

    def _update_material_tbo(self):
        """Syncs material and color TBOs with per-element settings."""
        if not self._mat_tbo_dirty: return
        
        # Sync Materials TBOs
        mat_arr1 = np.zeros((self.scene.N, 4), dtype=np.float32)
        mat_arr2 = np.zeros((self.scene.N, 4), dtype=np.float32)
        
        for i, el in enumerate(self.scene.elements):
            m = self.element_materials[el]
            mat_arr1[i] = [m.get('ambient',0.4), m.get('specular',0.6), m.get('shininess',48.0), m.get('rim',0.8)]
            mat_arr2[i] = [m.get('sheen',0.0), m.get('clearcoat',0.0), m.get('env',0.3), m.get('gloss',0.5)]
            
        self.rnd_sphere.upload_materials(mat_arr1, mat_arr2)
        
        # Color synchronization is now handled within _sync_gpu 
        # which is called every frame, so we just need to ensure 
        # it has the latest data by re-loading frames if needed 
        # or just letting the next frame's _sync_gpu pick it up.
        # However, to be immediate:
        colors = self._get_element_colors(self.scene)
        self.rnd_sphere.upload_data(self.scene.pos, self.scene.radii, colors)
        
        self.rnd_sphere.upload_selection(self.scene.selected_atoms, self.scene.N)
        self._mat_tbo_dirty = False

    def _update_clipping_alpha(self):
        """Updates the alpha channel of atom colors based on clipping range and alpha."""
        if not self.clipping_enabled or self.clipping_alpha <= 0.001:
            # If clipping is off or hard clipping, ensure full alpha
            if self.scene.colors is not None and self.scene.colors.shape[1] == 4:
                self.scene.colors[:, 3] = 1.0
                self._mark_dirty(pos=True)
            return

        # Get current colors (RGB)
        current_colors = self._get_element_colors(self.scene, force_refresh=True)
        if current_colors.shape[1] == 3:
            # Add alpha channel if not present
            rgba_colors = np.ones((current_colors.shape[0], 4), dtype=np.float32)
            rgba_colors[:, :3] = current_colors
        else:
            rgba_colors = np.copy(current_colors)

        # Calculate clipping mask
        if self.clipping_fractional:
            ext_lat = self.scene.lat
            inv_lat = np.linalg.inv(ext_lat)
            pos_frac = self.scene.pos @ inv_lat
            vals = pos_frac[:, self.clipping_axis]
        else:
            vals = self.scene.pos[:, self.clipping_axis]
        
        mask_in_range = (vals >= self.clipping_range[0]) & (vals <= self.clipping_range[1])
        if self.clipping_invert:
            mask_clipped = mask_in_range # Atoms *inside* the box are now the ones to be ghosted
        else:
            mask_clipped = ~mask_in_range # Atoms *outside* the box are the ones to be ghosted

        # Apply alpha to clipped atoms
        rgba_colors[mask_clipped, 3] = self.clipping_alpha
        rgba_colors[~mask_clipped, 3] = 1.0 # Ensure non-clipped atoms are fully opaque

        self.scene.colors = rgba_colors
        self._mark_dirty(pos=True) # Mark positions dirty to re-upload colors

    def _render_pretty_window(self):
        """Advanced Visualizer Controls (Final Version)."""
        imgui.set_next_window_size(400, 600, imgui.FIRST_USE_EVER)
        imgui.begin("Pretty Render Settings", True)

        if imgui.collapsing_header("Global HQ Options", imgui.TREE_NODE_DEFAULT_OPEN)[0]:
            _, self.materials['show_shadows'] = imgui.checkbox("Shadow Mapping (High Quality)", self.materials.get('show_shadows', True))
            
            # Resolution Dropdown
            res_options = ["512", "1024", "2048", "4096 (High)", "8192 (Ultra)"]
            vals = [512, 1024, 2048, 4096, 8192]
            changed_res, self.shadow_res_idx = imgui.combo("Shadow Quality", self.shadow_res_idx, res_options)
            if changed_res:
                self.rnd_sphere.change_shadow_res(vals[self.shadow_res_idx])
            
            _, self.materials['exposure'] = imgui.slider_float("Exposure", self.materials['exposure'], 0.1, 20.0)
            _, self.materials['gamma'] = imgui.slider_float("Gamma", self.materials['gamma'], 0.1, 5.0)
            _, self.materials['headlight'] = imgui.slider_float("Headlight", self.materials['headlight'], 0.0, 5.0)
            _, self.materials['shadow_intensity'] = imgui.slider_float("Shadow Intensity", self.materials['shadow_intensity'], 0.0, 1.0)
            _, self.materials['shadow_softness'] = imgui.slider_float("Shadow Soft (Blur)", self.materials['shadow_softness'], 0.1, 8.0)
            _, self.materials['shadow_bias'] = imgui.slider_float("Shadow Bias", self.materials.get('shadow_bias', 0.002), 0.0, 0.02)
            
            # Cleanup VFX controls
            imgui.separator()
            imgui.text("Post-Processing Effects:")
            if self.materials.get('show_ssao', True):
                _, self.materials['ssao_strength'] = imgui.slider_float("SSAO Strength", self.materials['ssao_strength'], 0.0, 5.0)
                _, self.materials['ssao_radius'] = imgui.slider_float("SSAO Radius (UV)", self.materials['ssao_radius'], 0.001, 0.1)
                _, self.materials['ssao_bias'] = imgui.slider_float("SSAO Bias (Depth)", self.materials['ssao_bias'], 0.0, 1.0)
                _, self.materials['debug_ao'] = imgui.checkbox("Debug: Show AO Only", self.materials.get('debug_ao', False))
                _, self.materials['debug_depth'] = imgui.checkbox("Debug: Show Depth Map", self.materials.get('debug_depth', False))
            
            _, self.materials['show_bloom'] = imgui.checkbox("Bloom (Glow)", self.materials['show_bloom'])
            if self.materials['show_bloom']:
                _, self.materials['bloom_intensity'] = imgui.slider_float("Bloom Intensity", self.materials['bloom_intensity'], 0.0, 5.0)
                _, self.materials['selection_glow'] = imgui.slider_float("Selection Glow", self.materials['selection_glow'], 0.0, 5.0)
            
            _, self.materials['show_fxaa'] = imgui.checkbox("FXAA (Smooth Edges)", self.materials['show_fxaa'])
            _, self.materials['show_outlines'] = imgui.checkbox("Selection Outlines (Silhouettes)", self.materials.get('show_outlines', True))
            _, self.materials['gamma'] = imgui.slider_float("Gamma Correction", self.materials['gamma'], 0.5, 3.0)
            
            imgui.separator()
            imgui.text("Atmosphere & Motion:")
            _, self.materials['fog_density'] = imgui.slider_float("Fog Density", self.materials['fog_density'], 0.0, 0.01)
            _, self.materials['pulse_amp'] = imgui.slider_float("Pulse Amplitude", self.materials['pulse_amp'], 0.0, 0.2)
            if self.materials['pulse_amp'] > 0:
                _, self.materials['pulse_speed'] = imgui.slider_float("Pulse Speed", self.materials['pulse_speed'], 0.0, 10.0)
            
            imgui.separator()
            imgui.text("Material Aesthetics:")
            _, self.materials['reflection_intensity'] = imgui.slider_float("Env Reflection", self.materials['reflection_intensity'], 0.0, 1.0)
            _, self.materials['reflection_gloss'] = imgui.slider_float("Glossiness/Sharpness", self.materials['reflection_gloss'], 0.0, 1.0)
            _, self.materials['sheen'] = imgui.slider_float("Sheen (Velvet)", self.materials['sheen'], 0.0, 1.0)
            _, self.materials['clearcoat'] = imgui.slider_float("Clearcoat (Gloss)", self.materials['clearcoat'], 0.0, 1.0)
            
            self._render_gui_lighting()

        if imgui.collapsing_header("Material Library (Per Element)")[0]:
            imgui.text("Bulk Styling (Apply to ALL):")
            changed_gp, self.global_preset_idx = imgui.combo("Global Preset", self.global_preset_idx, self.preset_options)
            if changed_gp:
                p_name = self.preset_options[self.global_preset_idx]
                preset_vals = self.material_presets[p_name]
                for e in self.element_materials:
                    self.element_materials[e].update(preset_vals)
                self._mat_tbo_dirty = True
            
            imgui.separator()
            for el, m in self.element_materials.items():
                if imgui.tree_node(f"[{el}] Properties"):
                    # Preset Selector
                    changed_p, p_idx = imgui.combo(f"Preset##{el}", 0, self.preset_options)
                    if changed_p:
                        m.update(self.material_presets[self.preset_options[p_idx]])
                        self._mat_tbo_dirty = True

                    ch1, m['col'] = imgui.color_edit3(f"Color##{el}", *m['col'])
                    ch2, m['ambient'] = imgui.slider_float(f"Ambient##{el}", m.get('ambient', 0.4), 0.0, 5.0)
                    ch3, m['specular'] = imgui.slider_float(f"Specular##{el}", m.get('specular', 0.6), 0.0, 10.0)
                    ch4, m['shininess'] = imgui.slider_float(f"Shininess##{el}", m.get('shininess', 48.0), 1.0, 512.0)
                    ch5, m['rim'] = imgui.slider_float(f"Rim##{el}", m.get('rim', 0.8), 0.0, 10.0)
                    
                    imgui.text("Advanced Aesthetic:")
                    ch6, m['sheen'] = imgui.slider_float(f"Sheen##{el}", m.get('sheen', 0.0), 0.0, 2.0)
                    ch7, m['clearcoat'] = imgui.slider_float(f"Clearcoat##{el}", m.get('clearcoat', 0.0), 0.0, 2.0)
                    ch8, m['env'] = imgui.slider_float(f"Env Intensity##{el}", m.get('env', 0.3), 0.0, 2.0)
                    ch9, m['gloss'] = imgui.slider_float(f"Glossiness##{el}", m.get('gloss', 0.5), 0.0, 1.0)

                    if any([ch1, ch2, ch3, ch4, ch5, ch6, ch7, ch8, ch9]): self._mat_tbo_dirty = True
                    imgui.tree_pop()

        imgui.text_disabled("(Tip: Hold 'L' + [1/2/3/4] + Drag to rotate different lights)")
        imgui.end()

    def _render_gui_lighting(self):
        """Interactive Multi-Light Switcher."""
        imgui.separator()
        imgui.text("Multi-Light Config:")
        
        # New selection for active light
        light_names = [f"Light {i+1}" for i in range(len(self.lights))]
        if light_names:
            _, self.selected_light_idx = imgui.combo("Interaction Focus", self.selected_light_idx, light_names)
            imgui.text_colored("Focus this light to use 'L' + Mouse rotation.", 0.6, 0.8, 1.0)

        if len(self.lights) < 4:
            if imgui.button("Add New Light"):
                self.lights.append({'pos': [1.0, 0.0, 0.0], 'col': [1.0, 1.0, 1.0]})

        for i in range(len(self.lights)-1, -1, -1):
            light = self.lights[i]
            opened, visible = imgui.collapsing_header(f"Light {i+1}##{i}", True)
            if not visible: self.lights.pop(i); continue
            if opened:
                changed_d, d = imgui.input_float3(f"Vector##{i}", *light['pos'])
                if changed_d: light['pos'] = list(d)
                changed_c, c = imgui.color_edit3(f"Color##{i}", *light['col'])
                if changed_c: light['col'] = list(c)
                if imgui.button(f"Remove Light##btn{i}"): self.lights.pop(i)

    def _render_gui(self):
        """Builds all ImGui UI windows."""
        self.gui_impl.process_inputs()
        imgui.new_frame()
        
        self._draw_fps_overlay()
        if self.show_help:
            self._draw_help_overlay()
            
        self._draw_annotations()

        if self.pretty_render:
            self._render_pretty_window()

        # Regular Control Panel
        imgui.set_next_window_bg_alpha(0.65)
        imgui.set_next_window_size(300, 600, imgui.FIRST_USE_EVER)
        imgui.begin("Control Panel", True)

        if imgui.collapsing_header("Structural Info", imgui.TREE_NODE_DEFAULT_OPEN)[0]:
            imgui.text(f"Composition: {self.scene.composition_str}")
            imgui.text(f"Atoms: {self.scene.N}")
            e_val = self.scene.energy
            if e_val is None or (isinstance(e_val, float) and (e_val != e_val)):  # NaN check
                e_str = "N/A"
            elif abs(e_val) > 1e12:
                e_str = f"{e_val:.4e} eV"
            else:
                e_str = f"{e_val:.6f} eV"
            imgui.text(f"Energy: {e_str}")

        if imgui.collapsing_header("Navigation & Playback", imgui.TREE_NODE_DEFAULT_OPEN)[0]:
            # Sync Toggles (Moved here for accessibility)
            if self.split_view:
                _, self.sync_frames = imgui.checkbox("Sync Time Control", self.sync_frames)
                imgui.same_line()
                _, self.sync_cameras = imgui.checkbox("Sync Camera Orbit", self.sync_cameras)
                imgui.separator()

            if not self.split_view or self.sync_frames:
                imgui.text(f"Global Index (0 to {len(self.provider)-1}):")
                c, self.index1 = imgui.drag_int("##f_global", self.index1, 1, 0, len(self.provider)-1)
                if c: 
                    self._load_frame1(self.index1)
                    if self.split_view: 
                         self.index2 = self.index1
                         self._load_frame2(self.index2)
            else:
                imgui.text("Index A (Left):")
                c1, self.index1 = imgui.drag_int("##f1", self.index1, 1, 0, len(self.provider)-1)
                if c1: self._load_frame1(self.index1)
                
                imgui.text("Index B (Right):")
                max2 = (len(self.provider2)-1) if self.provider2 else (len(self.provider)-1)
                c2, self.index2 = imgui.drag_int("##f2", self.index2, 1, 0, max2)
                if c2: self._load_frame2(self.index2)
            
            # Control Buttons
            if imgui.button(" |< "): 
                self.index1 = 0; self._load_frame1(0)
                if self.split_view and self.sync_frames: self.index2 = 0; self._load_frame2(0)
            imgui.same_line()
            
            if imgui.button("  <  "): 
                new1 = max(0, self.index1 - 1); self.index1 = new1; self._load_frame1(new1)
                if self.split_view and self.sync_frames: self.index2 = new1; self._load_frame2(new1)
            imgui.same_line()
            
            btn_play = "Pause" if self.playback_playing else " Play "
            if imgui.button(btn_play):
                self.playback_playing = not self.playback_playing
            imgui.same_line()
            
            if imgui.button("  >  "): 
                new1 = min(len(self.provider)-1, self.index1 + 1); self.index1 = new1; self._load_frame1(new1)
                if self.split_view and self.sync_frames: self.index2 = new1; self._load_frame2(new1)
            imgui.same_line()
            
            if imgui.button(" >| "): 
                last1 = len(self.provider)-1; self.index1 = last1; self._load_frame1(last1)
                if self.split_view and self.sync_frames: self.index2 = last1; self._load_frame2(last1)
            imgui.same_line()
            
            if imgui.button("Reset View"): self._center_camera()
            imgui.same_line()
            if imgui.button(" Help (F1) "):
                self.show_help = not self.show_help

            # FPS
            imgui.separator()
            imgui.text("Playback Speed (FPS):")
            changed_f, self.playback_speed = imgui.drag_float("##fps_drag", self.playback_speed, 1.0, 1.0, 240.0, "%.1f FPS")
            if changed_f: self.playback_speed = max(0.1, self.playback_speed)

            # Energy Trace & Histogram
            imgui.separator()
            if hasattr(self, "all_energies") and self.all_energies.size > 0:
                new_idx = draw_energy_plots(self.all_energies, self.index1)
                if new_idx != -1:
                    self.index1 = new_idx
                    self._load_frame1(self.index1)
                    if self.split_view and self.sync_frames:
                        self.index2 = self.index1
                        self._load_frame2(self.index2)

        if imgui.collapsing_header("Trajectory Comparison")[0]:
            changed_split, self.split_view = imgui.checkbox("Split View Mode", self.split_view)
            if changed_split and self.split_view and self.provider2 is None:
                # Auto-clone if nothing loaded yet
                self.load_comparison(None)

            if imgui.button("Clone Current (Double View)"):
                self.load_comparison(None)
            
            imgui.separator()
            imgui.text("Load Second Trajectory (File Path):")
            changed_pth, self._comp_path = imgui.input_text("##comp_path", self._comp_path, 255)
            if imgui.button("Load & Sync"):
                from .io.providers import StructureProvider
                try:
                    p2 = StructureProvider(self._comp_path)
                    self.load_comparison(p2)
                except Exception as e:
                    print(f"Failed to load comparison: {e}")
            if self.provider2 and self.provider2 != self.provider:
                imgui.text_colored(f"Comparing with: {os.path.basename(self._comp_path)}", 0.3, 1.0, 0.3)
        
        if imgui.collapsing_header("Display Settings")[0]:
            _, self.is_ortho = imgui.checkbox("Orthographic Mode", self.is_ortho)
            
            p_changed, self.pretty_render = imgui.checkbox("Pretty Render (HQ)", self.pretty_render)
            if p_changed and self.pretty_render:
                self.materials['show_shadows'] = True
                self.materials['show_bloom'] = True
                
            _, self.show_axes = imgui.checkbox("Show Coordinate Axes", self.show_axes)
            
            if imgui.tree_node("Camera Settings"):
                # Navigation Mode Dropdown
                modes = ["ORBIT", "FLY"]
                try: 
                    m_idx = modes.index(self.cam1.mode)
                except ValueError: 
                    m_idx = 0
                
                changed, m_idx = imgui.combo("Nav Mode", m_idx, modes)
                if changed:
                    self.cam1.mode = modes[m_idx]
                    if self.split_view: self.cam2.mode = modes[m_idx]
                
                _, self.cam1.invert_y = imgui.checkbox("Invert Mouse Y", self.cam1.invert_y)
                _, self.cam1.clamped_pitch = imgui.checkbox("Clamp Pitch (Poles)", self.cam1.clamped_pitch)
                _, self.show_origin_axes = imgui.checkbox("Show Origin Axes (0,0,0)", self.show_origin_axes)
                _, self.show_com = imgui.checkbox("Show Center of Mass", self.show_com)
                
                _, self.cam1.momentum = imgui.checkbox("Enable Momentum", self.cam1.momentum)
                _, self.cam1.sensitivity = imgui.slider_float("Sensitivity", self.cam1.sensitivity, 0.001, 0.05)
                _, self.cam1.damping = imgui.slider_float("Damping", self.cam1.damping, 0.0, 0.99)
                _, self.cam1.fly_speed = imgui.slider_float("Flight Speed", self.cam1.fly_speed, 1.0, 100.0)
                
                if imgui.button("Reset View (COM)"):
                    self.reset_view()
                imgui.tree_pop()
            
            _, self.atom_scale = imgui.slider_float("Atom Scale", self.atom_scale, 0.1, 3.0)
            
            
            # Supercell Controls
            imgui.text("Supercell (nx, ny, nz):")
            rx, ry, rz = self.scene.replicas
            changed_nx, rx = imgui.slider_int("nx", rx, 0, 20)
            changed_ny, ry = imgui.slider_int("ny", ry, 0, 20)
            changed_nz, rz = imgui.slider_int("nz", rz, 0, 20)
            if changed_nx or changed_ny or changed_nz:
                self.scene.set_replicas(rx, ry, rz)
                self._mark_dirty(pos=True)
                self._sync_gpu()
            
            imgui.separator()
            
            changed_b, self.show_bonds = imgui.checkbox("Show Bonds", self.show_bonds)
            if changed_b: 
                self._mark_dirty(pos=True)
            _, self.show_cell = imgui.checkbox("Show Unit Cell", self.show_cell)
            _, self.bg_color = imgui.color_edit3("Background", *self.bg_color)

        if self.scene.selected_atoms and imgui.collapsing_header("Selection & Atom Editing", imgui.TREE_NODE_DEFAULT_OPEN)[0]:
            imgui.text_colored(f"Selected: {len(self.scene.selected_atoms)} atoms", 0.4, 0.8, 1.0)
            
            if imgui.button("  Delete Selected  "):
                self.scene.remove_atoms(self.scene.selected_atoms)
                self._mark_dirty(pos=True)
                self._sync_gpu()
                self._update_coloring()
            
            imgui.separator()
            # Transformation Tools
            if imgui.tree_node("Transformations"):
                _, self.edit_translate_delta = imgui.drag_float3("Translate Delta", *self.edit_translate_delta, 0.1)
                if imgui.button(" Apply Translation "):
                    self.scene.translate_atoms(self.scene.selected_atoms, np.array(self.edit_translate_delta))
                    self._mark_dirty(pos=True)
                    self._sync_gpu()
                
                imgui.spacing()
                axes_list = ["X", "Y", "Z"]
                _, self.edit_rotate_axis_idx = imgui.combo("Rotation Axis", self.edit_rotate_axis_idx, axes_list)
                _, self.edit_rotate_angle = imgui.drag_float("Angle (deg)", self.edit_rotate_angle, 1.0)
                if imgui.button(" Apply Rotation "):
                    axis = np.zeros(3); axis[self.edit_rotate_axis_idx] = 1.0
                    self.scene.rotate_atoms(self.scene.selected_atoms, axis, self.edit_rotate_angle)
                    self._mark_dirty(pos=True)
                    self._sync_gpu()
                imgui.tree_pop()
            
            # Element Editing
            if imgui.tree_node("Chemical Editing"):
                _, self.edit_new_element = imgui.input_text("New Element", self.edit_new_element, 3)
                if imgui.button(" Change Element "):
                    self.scene.set_elements(self.scene.selected_atoms, self.edit_new_element.strip())
                    self._mark_dirty(pos=True)
                    self._sync_gpu()
                    self._update_coloring()
                imgui.tree_pop()

            if imgui.button("Clear Selection"):
                self.scene.clear_selection()
                self._sync_gpu()

        if imgui.collapsing_header("Atom Coloring", imgui.TREE_NODE_DEFAULT_OPEN)[0]:
            modes = ["Element", "Z Position", "Displacement", "Neighbors", "Coordination", "Voronoi Volume", "Selection", "Index"]
            try: cur_m_idx = modes.index(self.color_mode)
            except: cur_m_idx = 0
            changed_m, m_idx = imgui.combo("Mode", cur_m_idx, modes)
            if changed_m:
                self.color_mode = modes[m_idx]
                self._mark_dirty(pos=True)
                self._update_coloring()
            
            pals = ["BlueRed", "Viridis", "Magma", "Hot", "Grayscale"]
            try: cur_p_idx = pals.index(self.color_palette)
            except: cur_p_idx = 0
            changed_p, p_idx = imgui.combo("Palette", cur_p_idx, pals)
            if changed_p:
                self.color_palette = pals[p_idx]
                self._mark_dirty(pos=True)
                self._update_coloring()

            changed_ia, self.color_invert = imgui.checkbox("Invert Palette", self.color_invert)
            changed_ha, self.color_highlight_selection = imgui.checkbox("Highlight selected atoms", self.color_highlight_selection)
            
            changed_ar, self.color_range_auto = imgui.checkbox("Auto-Range Mapping", self.color_range_auto)
            if not self.color_range_auto:
                changed_rv, self.color_range = imgui.drag_float2("Value Range", *self.color_range, 0.1)
                if changed_rv: self._update_coloring()
            
            imgui.separator()
            if self.color_mode == "Z Position":
                changed_z, self.color_z_fractional = imgui.checkbox("Fractional (0-1)", self.color_z_fractional)
                if changed_z: self._update_coloring()
            elif self.color_mode == "Displacement":
                imgui.text(f"Ref: Frame {self.displacement_ref_index}")
                if imgui.button("Set current as Reference"):
                    self.displacement_ref_index = self.index1
                    self.displacement_ref_pos = np.copy(self.scene.pos)
                    self._update_coloring()
            elif self.color_mode == "Neighbors":
                changed_rc, self.color_nb_cutoff = imgui.drag_float("Cutoff (A)", self.color_nb_cutoff, 0.05, 1.0, 15.0)
                if changed_rc: self._update_coloring()
            
            if changed_ia or changed_ha or changed_ar: self._update_coloring()
            
            if imgui.button("  Manual Refresh Colors  "): self._update_coloring()
            imgui.same_line()
            if imgui.button("Reset (CPK)"):
                self.color_mode = "Element"
                self._mark_dirty(pos=True)
                self._update_coloring()

        if imgui.collapsing_header("Analysis Tools", imgui.TREE_NODE_DEFAULT_OPEN)[0]:
            # --- Global Scope & Computation ---
            imgui.text("Selection Scope:")
            scopes = ["Global", "Selected Atoms", "Region"]
            try: cur_scope_idx = scopes.index(self.analysis_scope)
            except: cur_scope_idx = 0
            changed_scope, cur_scope_idx = imgui.combo("##scope", cur_scope_idx, scopes)
            if changed_scope: self.analysis_scope = scopes[cur_scope_idx]
            
            _, self.show_selection_analytics = imgui.checkbox("Selection Geometry & Analytics", self.show_selection_analytics)
            imgui.separator()

            # --- 1. Graph Build ---
            if imgui.tree_node("Bonding & Graph Setup"):
                imgui.text("Radii Factor (Graph Threshold):")
                # Slider for real-time preview (cheap if cached)
                changed_bf, self.bond_factor = imgui.slider_float("##bf", self.bond_factor, 0.5, 2.0)
                
                # Visibility Toggle
                _, self.show_bonds = imgui.checkbox("Show Network / Bonds", self.show_bonds)
                
                if imgui.button("Build Graph") or changed_bf:
                    self._bond_cache.clear()
                    self._calc_bonds_simple(self.scene)
                    if self.scene2: self._calc_bonds_simple(self.scene2)
                    self._sync_gpu(self.scene)
                    if self.scene2: self._sync_gpu(self.scene2)
                imgui.tree_pop()

            if imgui.tree_node("Volumetric & Density Analysis"):
                _, self.vol_enabled = imgui.checkbox("Enable Overlay", self.vol_enabled)
                
                # Parameters for generation
                if not hasattr(self, "vol_res"): self.vol_res = 0.25
                if not hasattr(self, "vol_sigma"): self.vol_sigma = 0.5
                _, self.vol_res = imgui.slider_float("Grid Res (A)", self.vol_res, 0.1, 1.0)
                _, self.vol_sigma = imgui.slider_float("Gaussian Sigma", self.vol_sigma, 0.1, 2.0)

                if imgui.button("  Generate Density  "):
                    self._generate_gaussian_density()
                imgui.same_line()
                if imgui.button("  Load .cube  "):
                    self._load_volumetric_cube()
                
                if self.vol_enabled and self.vol_data:
                    imgui.separator()
                    
                    # Data Info
                    v_min = getattr(self, "vol_min", 0.0)
                    v_max = getattr(self, "vol_max", 1.0)
                    imgui.text_disabled(f"Grid Range: [{v_min:.3f}, {v_max:.3f}]")
                    
                    modes = ["Isosurface", "MIP (Maximum Intensity)"]
                    _, self.vol_mode_idx = imgui.combo("Render Mode", self.vol_mode_idx, modes)
                    
                    # Use actual range for slider
                    _, self.vol_iso_level = imgui.slider_float("Iso Level", self.vol_iso_level, v_min, v_max)
                    _, self.vol_opacity = imgui.slider_float("Opacity", self.vol_opacity, 0.0, 1.0)
                    
                    changed_c, c = imgui.color_edit3("Volume Color", *self.vol_color[:3])
                    if changed_c: self.vol_color[:3] = list(c)
                    self.vol_color[3] = self.vol_opacity
                    
                    if imgui.button("Clear Data"):
                        self.vol_data = None
                        self.vol_enabled = False
                imgui.tree_pop()
            
            if imgui.tree_node("Structure Clipping & Slicing"):
                changed_ce, self.clipping_enabled = imgui.checkbox("Enable Clipping Plane", self.clipping_enabled)
                if changed_ce: self._update_clipping_alpha()

                if self.clipping_enabled:
                    # Transparency slider for clipped atoms
                    changed_ca, self.clipping_alpha = imgui.slider_float("Ghost Alpha", self.clipping_alpha, 0.0, 1.0)
                    
                    imgui.same_line()
                    changed_ci, self.clipping_invert = imgui.checkbox("Invert Box", self.clipping_invert)
                    
                    # Axis Selection
                    axes = ["X-Axis", "Y-Axis", "Z-Axis"]
                    changed_ca, self.clipping_axis = imgui.combo("Clip Axis", self.clipping_axis, axes)
                    
                    # Coordinate Mode
                    changed_mode, self.clipping_fractional = imgui.checkbox("Fractional Mode", self.clipping_fractional)
                    
                    # Bounds Calculation for Slider
                    if self.clipping_fractional:
                        c_min, c_max = 0.0, 1.0
                        step = 0.01
                    else:
                        mn, mx = self.scene.get_full_bounds()
                        c_min, c_max = mn[self.clipping_axis], mx[self.clipping_axis]
                        padding = 0.1 * (c_max - c_min + 1.0)
                        c_min -= padding; c_max += padding
                        step = 0.1
                    
                    if changed_mode:
                        self.clipping_range = [c_min, c_max]
                    
                    changed_cr, self.clipping_range = imgui.drag_float2("Clip Bounds", *self.clipping_range, step, c_min, c_max)
                    
                    if imgui.button("  Reset Bounds  "):
                        self.clipping_range = [c_min, c_max]
                    
                    if any([changed_ce, changed_ca, changed_ci, changed_ca, changed_mode, changed_cr]):
                        self._update_clipping_alpha()

                imgui.tree_pop()

            # --- 2. RDF (Radial Distribution) ---
            if imgui.tree_node("Radial Distribution (RDF)"):
                imgui.text("Analysis Parameters (Live):")
                m_list = ["Bulk g(r)", "Partial (A-B)", "2D (In-plane)"]
                changed_m, self.rdf_mode_idx = imgui.combo("##rdf_m", self.rdf_mode_idx, m_list)
                self.rdf_mode = ["bulk", "cluster", "2d"][self.rdf_mode_idx]
                
                changed_r, self.rdf_rmax = imgui.slider_float("r_max", self.rdf_rmax, 1.0, 20.0)
                changed_b, self.rdf_bins = imgui.slider_int("bins", self.rdf_bins, 10, 1000)
                
                _, self.rdf_normalize = imgui.checkbox("Normalize Peak (Max=1.0)", self.rdf_normalize)
                
                if changed_m or changed_r or changed_b:
                    self._calc_rdf()

                imgui.text("Visualization:")
                if imgui.button("  Calculate  "): self._calc_rdf()
                
                if self.show_rdf:
                    from .ui.plots import draw_rdf_plot
                    draw_rdf_plot(self.rdf_results, self.rdf_visibility, self.rdf_normalize)
                imgui.tree_pop()

            # --- 3. Density Profiles ---
            if imgui.tree_node("Z-Density Profiles"):
                _, self.show_z_profile = imgui.checkbox("Enable Plot", self.show_z_profile)
                if imgui.button("Compute Z-Profile"): self._calc_z_profile()
                if self.show_z_profile:
                    from .ui.plots import draw_z_profile
                    draw_z_profile(self.z_profile_x, self.z_profile_y)
                imgui.tree_pop()

            # --- 4. Structural Dynamics ---
            if imgui.tree_node("Structural Dynamics"):
                imgui.text("Reference Frame:")
                _, self.displacement_ref_index = imgui.slider_int("##ref_idx", self.displacement_ref_index, 0, len(self.provider)-1)
                
                if imgui.button("Compute Displacement"):
                    p0_any, _, _, _, _, _ = self.provider.get(self.displacement_ref_index)
                    p0 = np.asarray(p0_any, dtype=np.float32).reshape(-1, 3)
                    p1 = self.scene.pos
                    
                    if p0.shape != p1.shape:
                        self._rmsd_err = f"Error: Atom count mismatch ({p0.shape[0]} vs {p1.shape[0]})"
                        if hasattr(self, "_rmsd"): del self._rmsd
                    else:
                        dist_sq = np.sum((p1 - p0)**2, axis=1)
                        self._rmsd = np.sqrt(np.mean(dist_sq))
                        if hasattr(self, "_rmsd_err"): del self._rmsd_err
                        print(f"RMSD vs Frame {self.displacement_ref_index}: {self._rmsd:.4f} A")
                
                if hasattr(self, "_rmsd"):
                    imgui.text_colored(f"RMSD vs Ref: {self._rmsd:.4f} A", 0.3, 1.0, 0.3)
                if hasattr(self, "_rmsd_err"):
                    imgui.text_colored(self._rmsd_err, 1.0, 0.3, 0.3)
                imgui.tree_pop()

            # --- 5. Voronoi & Volumes ---
            if imgui.tree_node("Voronoi Analysis"):
                if imgui.button("Compute Voronoi"): self._calc_voronoi()
                if hasattr(self, "voro_results"):
                    vols = self.voro_results.volumes
                    imgui.text(f"Mean Volume: {np.mean(vols):.2f} A^3")
                    imgui.text(f"Std Dev: {np.std(vols):.2f} A^3")
                imgui.tree_pop()

            # --- 6. Coordination Analysis ---
            if imgui.tree_node("Coordination Numbers"):
                _, self.cn_cutoff = imgui.slider_float("Cutoff (A)", self.cn_cutoff, 1.0, 10.0)
                if imgui.button("  Calculate CN  "): self._calc_cn_distribution()
                
                if self.cn_results is not None:
                    imgui.text(f"Mean CN: {np.mean(self.cn_results):.2f}")
                    if imgui.button(" Export CN (CSV) "):
                        ts = int(time.time())
                        hist, bins = np.histogram(self.cn_results, bins=np.arange(np.max(self.cn_results)+2))
                        np.savetxt(f"cn_dist_{ts}.csv", np.column_stack((bins[:-1], hist)), header="CN, count", delimiter=",")
                        self.export_status = f"CN exported to cn_dist_{ts}.csv"
                imgui.tree_pop()

            # --- 7. Measurement Tracker ---
            if imgui.tree_node("Measurements & Tools"):
                if len(self.scene.selected_atoms) == 2:
                    i1, i2 = self.scene.selected_atoms
                    d = np.linalg.norm(self.scene.pos[i1] - self.scene.pos[i2])
                    imgui.text_colored(f"Distance ({i1}-{i2}): {d:.4f} A", 1.0, 0.8, 0.4)
                    if imgui.button(" Export Distance (CSV) "):
                        ts = int(time.time())
                        with open(f"measurement_{ts}.csv", "w") as f:
                            f.write(f"Atom1,Atom2,Distance(A)\n{i1},{i2},{d:.6f}")
                        self.export_status = "Distance exported."
                elif len(self.scene.selected_atoms) == 1:
                    imgui.text_disabled("Select 2 atoms to measure distance.")
                imgui.tree_pop()

            # --- 8. Electronic Structure (TB) ---
            if imgui.tree_node("Electronic Structure (TB)"):
                imgui.text_disabled("Orthogonal tight-binding, 1 orbital/atom")
                imgui.separator()

                # Model settings
                _, self.tb_params.cutoff_radius   = imgui.slider_float("Cutoff (A)",      self.tb_params.cutoff_radius,   1.0, 8.0)
                _, self.tb_params.hopping_scale    = imgui.slider_float("Hopping Scale",   self.tb_params.hopping_scale,   0.1, 4.0)
                _, self.tb_params.beta_scale       = imgui.slider_float("Decay Beta",      self.tb_params.beta_scale,      0.1, 3.0)
                _, self.tb_params.smearing_kT      = imgui.slider_float("Smearing kT(eV)", self.tb_params.smearing_kT,     0.0, 0.5)
                _, self.tb_params.dos_sigma        = imgui.slider_float("DOS Sigma (eV)",  self.tb_params.dos_sigma,       0.01, 1.0)

                ref_modes = ["fermi", "homo", "none"]
                try: ref_idx = ref_modes.index(self.tb_params.energy_ref)
                except: ref_idx = 0
                changed_ref, ref_idx = imgui.combo("E Reference", ref_idx, ref_modes)
                if changed_ref: self.tb_params.energy_ref = ref_modes[ref_idx]

                imgui.separator()
                if imgui.button("  Run TB Calculation  "):
                    self._run_tb_calculation()
                
                if self.tb_status:
                    col = (0.4, 1.0, 0.4) if "Done" in self.tb_status else (1.0, 0.6, 0.2)
                    imgui.text_colored(self.tb_status, *col)

                if self.tb_results is not None:
                    res = self.tb_results
                    imgui.separator()
                    imgui.text_colored(f"HOMO: {res.homo:+.3f} eV", 0.4, 1.0, 0.6)
                    imgui.same_line()
                    imgui.text_colored(f"  LUMO: {res.lumo:+.3f} eV", 1.0, 0.6, 0.4)
                    imgui.text_colored(f"Gap: {res.gap:.3f} eV  |  N_e: {res.n_electrons}", 1.0, 1.0, 0.4)

                    imgui.separator()
                    _, self.tb_show_dos    = imgui.checkbox("Show DOS/PDOS Window", self.tb_show_dos)
                    _, self.tb_show_ladder = imgui.checkbox("Show Energy Levels",   self.tb_show_ladder)

                    # --- Sub-tab: 3D Visualization ---
                    if imgui.tree_node("3D Field Visualization"):
                        imgui.text_disabled("Reconstruct real-space orbitals/density")
                        _, self.tb_params.grid_spacing = imgui.slider_float("Grid Spacing (A)", self.tb_params.grid_spacing, 0.05, 0.4)
                        _, self.tb_params.grid_padding = imgui.slider_float("Padding (A)",      self.tb_params.grid_padding, 1.0, 6.0)
                        _, self.tb_vol_iso             = imgui.slider_float("Isosurface",       self.tb_vol_iso,            0.001, 0.4)
                        
                        imgui.separator()
                        viz_modes = ["None", "HOMO", "LUMO", "Total Density"]
                        try: vm_idx = viz_modes.index(self.tb_viz_mode)
                        except: vm_idx = 0
                        changed_vm, vm_idx = imgui.combo("Target Field", vm_idx, viz_modes)
                        if changed_vm: 
                            self.tb_viz_mode = viz_modes[vm_idx]
                            
                        if imgui.button("  Update/Reconstruct Field  "):
                            self._reconstruct_tb_field()
                            
                        _, self.tb_show_volume = imgui.checkbox("Display 3D Volume", self.tb_show_volume)
                        if self.tb_show_volume:
                            imgui.text_colored("Rendering using Single-Pass Raymarching", 0.4, 0.8, 1.0)
                        
                        imgui.tree_pop()

                    # Color mode
                    tb_color_modes = ["None", "TB Charge", "TB Bond Order"]
                    try: tc_idx = tb_color_modes.index(self.tb_color_mode)
                    except: tc_idx = 0
                    changed_tc, tc_idx = imgui.combo("Atom Color Mode", tc_idx, tb_color_modes)
                    if changed_tc:
                        self.tb_color_mode = tb_color_modes[tc_idx]
                        self._apply_tb_coloring()

                    # Export
                    if imgui.button(" Export DOS (CSV) "):
                        ts = int(time.time())
                        header = "E(eV),DOS"
                        data = np.column_stack([res.dos_x, res.dos_y])
                        for sp, py in res.pdos_y.items():
                            header += f",PDOS_{sp}"
                            data = np.column_stack([data, py])
                        np.savetxt(f"tb_dos_{ts}.csv", data, header=header, delimiter=",")
                        self.export_status = f"DOS exported to tb_dos_{ts}.csv"

                imgui.tree_pop()

        if self.clipping_enabled:
            self._update_clipping_alpha()
            
        # Draw Passes
        # Draw focus frame for Split View
        if self.split_view:
            w, h = glfw.get_window_size(self.win)
            dl = imgui.get_foreground_draw_list()
            # Dynamic colors
            col_active = imgui.get_color_u32_rgba(1.0, 0.6, 0.1, 0.8) # Orange glow
            col_inactive = imgui.get_color_u32_rgba(0.5, 0.5, 0.5, 0.2) # Dim grey
            
            # Left Half (A)
            is_L = self.active_viewport == 1
            dl.add_rect(4, 4, w//2 - 4, h - 4, col_active if is_L else col_inactive, thickness=4.0 if is_L else 1.0)
            
            # Right Half (B)
            is_R = self.active_viewport == 2
            dl.add_rect(w//2 + 4, 4, w - 4, h - 4, col_active if is_R else col_inactive, thickness=4.0 if is_R else 1.0)

        if imgui.collapsing_header("Output & Export", imgui.TREE_NODE_DEFAULT_OPEN)[0]:
            imgui.text("Capture System:")
            if imgui.button("  Take Screenshot (PNG)  "):
                self.save_screenshot()
            
            imgui.separator(); imgui.text("Structure Export:")
            _, self.export_prefix = imgui.input_text("Filename Prefix", self.export_prefix, 255)
            changed_fmt, self.export_fmt_idx = imgui.combo("Format", self.export_fmt_idx, self.export_formats)
            _, self.export_all_frames = imgui.checkbox("Export All Frames", self.export_all_frames)
            
            btn_txt = "  Export All Frames  " if self.export_all_frames else "  Export Current Structure  "
            imgui.separator(); imgui.text("GIF Trajectory Export:")
            _, self.gif_frame_count = imgui.slider_int("Frames", self.gif_frame_count, 2, 100)
            _, self.gif_random_selection = imgui.checkbox("Random Frame Selection", self.gif_random_selection)
            
            if imgui.button("  Export GIF Animation  "):
                self.export_gif()
            
            imgui.separator(); imgui.text("Structure Export:")
            imgui.text("Export Path:")
            _, self._export_path = imgui.input_text("##fout", self._export_path, 256)
            _, self.export_selected_only = imgui.checkbox("Selected Atoms Only", self.export_selected_only)
            
            if imgui.button("  Export Structure  "):
                self.export_current_structure()
            
            if self.export_status:
                imgui.text_colored(self.export_status, 0.5, 1.0, 0.5)

        imgui.end() # End Main Control
        
        # Draw Additional Analytics Windows
        self._draw_selection_analytics_window()
        
        # TB Electronic Structure Windows
        if self.tb_results is not None:
            if self.tb_show_dos:
                res = self.tb_results
                draw_dos_plot(res.dos_x, res.dos_y, res.fermi_level,
                              pdos_y=res.pdos_y, homo=res.homo, 
                              lumo=res.lumo, gap=res.gap)
            if self.tb_show_ladder:
                res = self.tb_results
                draw_eigenvalue_ladder(res.eigenvalues, res.n_occ_states,
                                       ipr=res.ipr, fermi=res.fermi_level)

        imgui.render()
        self.gui_impl.render(imgui.get_draw_data())

    def _draw_fps_overlay(self):
        """Renders a minimal FPS counter in the top-right corner."""
        imgui.set_next_window_position(self.win_size[0] - 10, 10, imgui.ALWAYS, 1.0, 0.0)
        imgui.set_next_window_bg_alpha(0.35)
        flags = (imgui.WINDOW_NO_TITLE_BAR | imgui.WINDOW_NO_RESIZE | 
                 imgui.WINDOW_ALWAYS_AUTO_RESIZE | imgui.WINDOW_NO_MOVE | 
                 imgui.WINDOW_NO_SAVED_SETTINGS | imgui.WINDOW_NO_FOCUS_ON_APPEARING | 
                 imgui.WINDOW_NO_NAV)
        
        imgui.begin("FPS Overlay", False, flags)
        fps = self.quality.current_fps
        col = (0.2, 1.0, 0.2) if fps > 45 else (1.0, 0.8, 0.0) if fps > 25 else (1.0, 0.2, 0.2)
        imgui.text_colored(f"{fps:.1f} FPS", *col)
        imgui.end()

    def _draw_help_overlay(self):
        """Renders a comprehensive help window listing all controls."""
        imgui.set_next_window_size(450, 500, imgui.FIRST_USE_EVER)
        imgui.begin("Controls & Help", True)
        if not imgui.is_window_focused() and imgui.is_window_appearing():
            imgui.set_window_focus()

        imgui.text_colored("Mouse Controls:", 0.6, 0.8, 1.0)
        imgui.bullet_text("Left Click: Orbit / Rotate Scene")
        imgui.bullet_text("Ctrl + Left Click: Pan View")
        imgui.bullet_text("Shift + Left Click: Box Selection (Drag)")
        imgui.bullet_text("Middle Click: Roll (Camera Rotation)")
        imgui.bullet_text("Scroll: Zoom In/Out")
        imgui.separator()

        imgui.text_colored("Keyboard Shortcuts:", 0.6, 0.8, 1.0)
        imgui.bullet_text("F1: Toggle this Help Overlay")
        imgui.bullet_text("H: Toggle GUI (Control Panels)")
        imgui.bullet_text("ESC: Close Application")
        imgui.bullet_text("L + Mouse: Rotate Active Light")
        imgui.bullet_text("L + [1-4] + Mouse: Select and Rotate Light")
        imgui.separator()

        imgui.text_colored("Navigation Modes (FLY Mode):", 0.6, 0.8, 1.0)
        imgui.bullet_text("W / S: Move Forward / Backward")
        imgui.bullet_text("A / D: Move Left / Right")
        imgui.bullet_text("Q / E: Move Up / Down")
        imgui.bullet_text("Shift: Boost Movement Speed")
        
        imgui.separator()
        if imgui.button("Close (F1)"):
            self.show_help = False
        imgui.end()

    def _update_playback(self):
        now = glfw.get_time()
        dt = now - self.last_time
        self.last_time = now
        
        if self.playback_playing:
            self.time_acc += dt
            interval = 1.0 / max(0.1, self.playback_speed)
            if self.time_acc >= interval:
                self.time_acc = 0
                # Move to next frame
                self.index1 = (self.index1 + 1) % len(self.provider)
                self._load_frame1(self.index1)
                
                # Sync if needed
                if self.split_view and self.sync_frames:
                    self.index2 = self.index1
                    self._load_frame2(self.index2)

    # --- Callbacks ---
    def _on_scroll(self, win, dx, dy):
        if imgui.get_io().want_capture_mouse: return
        zoom_speed = 0.1
        factor = 1.0 - dy * zoom_speed
        
        # Zoom camera under cursor
        x, _ = glfw.get_cursor_pos(win)
        w, _ = glfw.get_window_size(win)
        target_cam = self.cam1
        if self.split_view and x > w/2: target_cam = self.cam2

        if self.is_ortho:
            target_cam.ortho_scale = np.clip(target_cam.ortho_scale * factor, 0.1, 2000.0)
        else:
            dist = np.linalg.norm(target_cam.offset)
            new_dist = np.clip(dist * factor, 0.1, 2000.0)
            target_cam.offset = (target_cam.offset / dist) * new_dist
        
        if self.split_view and self.sync_cameras:
            other = self.cam2 if target_cam is self.cam1 else self.cam1
            other.offset = np.copy(target_cam.offset)
            other.ortho_scale = target_cam.ortho_scale

    def _on_mouse_button(self, win, btn, act, mods):
        if imgui.get_io().want_capture_mouse: return
        
        x, y = glfw.get_cursor_pos(win)
        w, _ = glfw.get_window_size(win)
        
        # Set Focus
        if self.split_view:
             self.active_viewport = 2 if x > w/2 else 1

        if btn == glfw.MOUSE_BUTTON_LEFT:
            if act == glfw.PRESS:
                # Hotkeys for light rotation: L + (1-4) or L alone
                if glfw.get_key(win, glfw.KEY_L) == glfw.PRESS:
                    self.light_rotating = True
                    for i, key in enumerate([glfw.KEY_1, glfw.KEY_2, glfw.KEY_3, glfw.KEY_4]):
                        if glfw.get_key(win, key) == glfw.PRESS:
                            if i < len(self.lights): self.selected_light_idx = i
                elif mods & glfw.MOD_CONTROL: self.panning = True
                elif mods & glfw.MOD_SHIFT:
                    self.box_selecting = True
                    self.scene.box_start = (x, y)
                    self.scene.box_end = (x, y)
                else: self.orbiting = True
            elif act == glfw.RELEASE:
                if self.box_selecting:
                    self.scene.box_end = (x, y)
                    dist = np.linalg.norm(np.array(self.scene.box_end) - np.array(self.scene.box_start))
                    if dist < 5.0:
                         idx = self._pick_atom(x, y)
                         if idx is not None:
                             self.scene.toggle_atom_selection(idx)
                         else:
                             self.scene.clear_selection()
                    else:
                         self._select_atoms_in_box()
                    
                    self.box_selecting = False
                    self.selection_dirty = True
                    self._mark_dirty(sel=True)
                    self._sync_gpu()
                self.orbiting = self.panning = self.light_rotating = False
        elif btn == glfw.MOUSE_BUTTON_MIDDLE:
            self.rolling = (act == glfw.PRESS)

    def _on_mouse_move(self, win, x, y):
        if imgui.get_io().want_capture_mouse: return
        w, _ = glfw.get_window_size(win)
        
        # Camera selection logic
        target_cam = self.cam1
        if self.split_view and x > w/2: target_cam = self.cam2
        
        if self.last_x is None: self.last_x, self.last_y = x, y; return
        dx, dy = x - self.last_x, y - self.last_y
        self.last_x, self.last_y = x, y
        
        if self.light_rotating and self.lights:
            # Rotate Selected Light around Scene Center
            from .camera import rodrigues, WORLD_UP
            idx = min(self.selected_light_idx, len(self.lights)-1)
            l = self.lights[idx]
            # Horizontal (Yaw-like)
            l['pos'] = list(rodrigues(np.array(l['pos']), WORLD_UP, -0.01 * dx))
            # Vertical (Pitch-like)
            right = np.cross(np.array(l['pos']), WORLD_UP); right = right / (np.linalg.norm(right) + 1e-9)
            l['pos'] = list(rodrigues(np.array(l['pos']), right, 0.01 * dy))
        
        if self.orbiting:
            # Set target velocity based on mouse delta
            target_cam.yaw_v = dx * target_cam.sensitivity
            target_cam.pitch_v = -dy * target_cam.sensitivity
            
            if self.split_view and self.sync_cameras:
                other = self.cam2 if target_cam is self.cam1 else self.cam1
                other.yaw_v = target_cam.yaw_v
                other.pitch_v = target_cam.pitch_v
                other.offset = np.copy(target_cam.offset)
                other.center = np.copy(target_cam.center)
        
        elif self.rolling:
                target_cam.rotate_roll(-dx * target_cam.sensitivity * 5.0)
        
        elif self.panning:
            scale = target_cam.ortho_scale / 1000.0 if self.is_ortho else np.linalg.norm(target_cam.offset) * 0.001
            target_cam.pan[0] -= dx * scale
            target_cam.pan[1] += dy * scale
            if self.split_view and self.sync_cameras:
                other = self.cam2 if target_cam is self.cam1 else self.cam1
                other.pan = np.copy(target_cam.pan)
                other.ortho_scale = target_cam.ortho_scale

    def _on_key(self, win, key, sc, act, mods):
        if imgui.get_io().want_capture_keyboard: return
        
        if act == glfw.PRESS:
            self.keys_down.add(key)
            if key == glfw.KEY_ESCAPE: glfw.set_window_should_close(win, True)
            if key == glfw.KEY_H: self.show_gui = not self.show_gui
            if key == glfw.KEY_F1: self.show_help = not self.show_help
        elif act == glfw.RELEASE:
            if key in self.keys_down: self.keys_down.remove(key)

    def _update_camera_keyboard(self, dt):
        """Processes WASD/QE movement for cameras in FLY mode."""
        speed = self.cam1.fly_speed * dt
        if glfw.KEY_LEFT_SHIFT in self.keys_down: speed *= 3.0
        
        if self.cam1.mode == "FLY":
            dx, dy, dz = 0, 0, 0
            if glfw.KEY_W in self.keys_down: dz += speed
            if glfw.KEY_S in self.keys_down: dz -= speed
            if glfw.KEY_A in self.keys_down: dx -= speed
            if glfw.KEY_D in self.keys_down: dx += speed
            if glfw.KEY_Q in self.keys_down: dy += speed
            if glfw.KEY_E in self.keys_down: dy -= speed
            if dx or dy or dz: self.cam1.move_local(dx, dy, dz)
            
        if self.split_view and self.cam2.mode == "FLY":
            # (Similar logic for cam2 if needed, but usually we sync them)
            if self.sync_cameras:
                self.cam2.center = np.copy(self.cam1.center)
                self.cam2.offset = np.copy(self.cam1.offset)

    def _pick_atom(self, x, y):
        w, h = glfw.get_framebuffer_size(self.win)
        win_w, win_h = glfw.get_window_size(self.win)
        ratio = w / win_w if win_w > 0 else 1.0
        
        # Sync matrices exactly with main render pass
        aspect = w / h if h > 0 else 1.0
        view = self.cam1.get_view_matrix()
        self._update_clipping_planes(self.scene, self.cam1)
        proj = self.cam1.get_projection_matrix(aspect, is_ortho=self.is_ortho)
            
        return self.rnd_pick.pick(x*ratio, y*ratio, w, h, self.scene.N, len(self.scene.replica_offsets),
                                  self.rnd_sphere.pos_tbo, self.rnd_sphere.rep_tbo, view, proj, self.is_ortho)

    def _select_atoms_in_box(self):
        """Identifies and selects atoms whose screen projection falls within the box."""
        w, h = glfw.get_framebuffer_size(self.win)
        win_w, win_h = glfw.get_window_size(self.win)
        
        # Sync matrices exactly with rendering pipeline
        aspect = w / h if h > 0 else 1.0
        view = self.cam1.get_view_matrix()
        self._update_clipping_planes(self.scene, self.cam1)
        proj = self.cam1.get_projection_matrix(aspect, is_ortho=self.is_ortho)
        vp = proj @ view
        
        # Selection Box Bounds (Screen space)
        x1, y1 = self.scene.box_start
        x2, y2 = self.scene.box_end
        sx1, sx2 = min(x1, x2), max(x1, x2)
        sy1, sy2 = min(y1, y2), max(y1, y2)
        
        new_indices = []
        for i in range(self.scene.N):
            pos_w = np.append(self.scene.pos[i], 1.0)
            pos_v = vp @ pos_w
            if pos_v[3] <= 0: continue
            
            ndc_x, ndc_y = pos_v[0] / pos_v[3], pos_v[1] / pos_v[3]
            cur_x = (ndc_x * 0.5 + 0.5) * win_w
            cur_y = (1.0 - (ndc_y * 0.5 + 0.5)) * win_h
            
            if sx1 <= cur_x <= sx2 and sy1 <= cur_y <= sy2:
                new_indices.append(i)
        
        if new_indices:
            self.scene.select_atoms(new_indices, append=True)
            self.selection_dirty = True
            self._mark_dirty(sel=True)
            self._sync_gpu()
        else:
            # If a large box was drawn and nothing was found, clear selection
            # (matches standard GUI behavior for 'empty click/drag')
            dist = np.linalg.norm(np.array(self.scene.box_end) - np.array(self.scene.box_start))
            if dist > 50.0:
                self.scene.clear_selection()
                self._mark_dirty(sel=True)
                self._sync_gpu()
            # User earlier requested clear if clicking empty space, so let's clear if dragging empty space too.

    def _draw_selection_box(self):
        """Draws the selection rectangle overlay using ImGui's draw list."""
        if not self.box_selecting: return
        draw_list = imgui.get_background_draw_list()
        x1, y1 = self.scene.box_start
        x2, y2 = self.scene.box_end
        draw_list.add_rect(x1, y1, x2, y2, imgui.get_color_u32_rgba(1.0, 1.0, 0.3, 0.8), thickness=2.0)
        draw_list.add_rect_filled(x1, y1, x2, y2, imgui.get_color_u32_rgba(1.0, 1.0, 0.3, 0.1))

    def _draw_annotations(self):
        """Overlays real-time geometric info (IDs, distances, angles) on selected atoms."""
        if not self.scene.selected_atoms: return
        
        # 1. Setup projection matrices
        w, h = glfw.get_framebuffer_size(self.win)
        win_w, win_h = glfw.get_window_size(self.win)
        view = self.cam1.get_view_matrix()
        aspect = w / h if h > 0 else 1.0
        self._update_clipping_planes(self.scene, self.cam1)
        proj = self.cam1.get_projection_matrix(aspect, is_ortho=self.is_ortho)
        vp = proj @ view
        
        draw_list = imgui.get_foreground_draw_list()
        selected = self.scene.selected_atoms
        
        def project(pos):
            p_v = vp @ np.append(pos, 1.0)
            if p_v[3] <= 0: return None
            ndc_x, ndc_y = p_v[0] / p_v[3], p_v[1] / p_v[3]
            return (ndc_x * 0.5 + 0.5) * win_w, (1.0 - (ndc_y * 0.5 + 0.5)) * win_h
            
        def draw_text_res(x, y, color, text):
            # Draw shadow for contrast
            draw_list.add_text(x + 1, y + 1, imgui.get_color_u32_rgba(0,0,0,1), text)
            draw_list.add_text(x, y, color, text)

        # Adaptive performance threshold: If > 20 selected, show only 3D markers
        if len(selected) > 20:
            draw_text_res(20, win_h - 40, imgui.get_color_u32_rgba(1, 0.5, 0.5, 1.0), 
                          "Performance Mode: Measurement details & circles hidden (>50 atoms selected).")
            return

        # 0. Selection Circle Outlines (for all selected - only if < 50)
        for idx in selected:
            if idx < self.scene.N:
                p2d = project(self.scene.pos[idx])
                if p2d:
                    # Estimate screen radius (rough but effective)
                    dist_cam = np.linalg.norm(self.cam.eye() - self.scene.pos[idx])
                    r_px = (self.atom_scale * 500.0) / (dist_cam + 1e-6)
                    if self.is_ortho: r_px = self.atom_scale * 100.0 / self.cam.ortho_scale
                    draw_list.add_circle(p2d[0], p2d[1], max(r_px, 5.0), imgui.get_color_u32_rgba(1, 1, 0, 0.8), thickness=2.0)
            
        # 1. Measurement Chains (Consecutive Segments)
            
        count = 0
        limit = 1000
        
        # Distances
        for i in range(len(selected)-1):
            if count >= limit: break
            idx1, idx2 = selected[i], selected[i+1]
            if idx1 < self.scene.N and idx2 < self.scene.N:
                p1, p2 = self.scene.pos[idx1], self.scene.pos[idx2]
                diff = p2 - p1
                if self.scene.is_periodic:
                    diff = diff - np.round(diff @ np.linalg.inv(self.scene.lat)) @ self.scene.lat
                dist = np.linalg.norm(diff)
                p1_2d, p2_2d = project(p1), project(p1 + diff)
                if p1_2d and p2_2d:
                    draw_list.add_line(p1_2d[0], p1_2d[1], p2_2d[0], p2_2d[1], imgui.get_color_u32_rgba(1,1,0.5,0.7), 2.0)
                    mid_2d = project(p1 + diff * 0.5)
                    if mid_2d:
                        draw_text_res(mid_2d[0], mid_2d[1], imgui.get_color_u32_rgba(1,1,0.5,1), f"{dist:.3f} A")
                    count += 1

        # Angles
        for i in range(len(selected)-2):
            if count >= limit: break
            idx1, idx2, idx3 = selected[i], selected[i+1], selected[i+2]
            if all(idx < self.scene.N for idx in [idx1, idx2, idx3]):
                p1, p2, p3 = self.scene.pos[idx1], self.scene.pos[idx2], self.scene.pos[idx3]
                def get_diff(pa, pb):
                    d = pb - pa
                    if self.scene.is_periodic:
                        d = d - np.round(d @ np.linalg.inv(self.scene.lat)) @ self.scene.lat
                    return d
                d21, d23 = get_diff(p2, p1), get_diff(p2, p3)
                cos_a = np.dot(d21, d23) / (np.linalg.norm(d21) * np.linalg.norm(d23) + 1e-9)
                angle = np.degrees(np.arccos(np.clip(cos_a, -1.0, 1.0)))
                p2_2d = project(p2)
                if p2_2d:
                    draw_text_res(p2_2d[0]+15, p2_2d[1], imgui.get_color_u32_rgba(0.5,1,1,1), f"{angle:.1f} deg")
                    count += 1
                    
        # Individual Labels (if only 1 selected)
        if len(selected) == 1:
            idx = selected[0]
            if idx < self.scene.N:
                p2d = project(self.scene.pos[idx])
                if p2d:
                    txt = f"ID:{idx} {self.scene.elements[idx]}"
                    draw_text_res(p2d[0]+12, p2d[1]-12, imgui.get_color_u32_rgba(1,1,1,1), txt)

    def _draw_selection_analytics_window(self):
        """Displays a dedicated window with pairwise distances and angles for the current selection."""
        if not self.show_selection_analytics or not self.scene.selected_atoms: return
        selected = self.scene.selected_atoms
        if len(selected) < 2: return
        
        imgui.set_next_window_size(350, 400, imgui.FIRST_USE_EVER)
        expanded, self.show_selection_analytics = imgui.begin("Selection Analytics", True)
        if expanded:
            # 1. Fetch from Cache (Lazy O(N) or O(1) if valid)
            cache = self._compute_selection_analytics()
            dists_seq = cache["dists_seq"]
            angles_seq = cache["angles_seq"]
            dihedrals_seq = cache["dihedrals_seq"]
            
            # 2. Pairwise Calculation (Strictly On-Demand)
            dists_pw = cache.get("dists_pw")
            if self.force_full_pairwise and dists_pw is None:
                dists_pw = []
                pos = self.scene.pos
                lat = self.scene.lat
                is_periodic = self.scene.is_periodic
                def get_diff(pa, pb):
                    d = pb - pa
                    if is_periodic:
                        d = d - np.round(d @ np.linalg.inv(lat)) @ lat
                    return d
                for i in range(len(selected)):
                    for j in range(i+1, len(selected)):
                        dists_pw.append(np.linalg.norm(get_diff(pos[selected[i]], pos[selected[j]])))
                cache["dists_pw"] = dists_pw

            # UI Rendering
            imgui.text_colored(f"Selected Atoms: {len(selected)}", 0.4, 0.8, 1.0)
            if imgui.button("  Clear Selection  "): 
                self.scene.clear_selection()
                self.selection_dirty = True
                self._mark_dirty(sel=True)
                self._sync_gpu()
            
            # --- Pairwise Distances ---
            if dists_pw:
                header_txt = f"Pairwise Distances ({len(dists_pw)} pairs)"
                if imgui.collapsing_header(header_txt)[0]:
                    imgui.begin_child("dist_scroll", 0, 150, border=True)
                    # UI Cap: Max 100 rows to prevent ImGui lag
                    max_rows = 100
                    row_count = min(len(dists_pw), max_rows)
                    k = 0
                    for i in range(len(selected)):
                        if k >= row_count: break
                        for j in range(i+1, len(selected)):
                            if k >= row_count: break
                            imgui.text(f"Atoms {selected[i]}-{selected[j]}:")
                            imgui.same_line(150); imgui.text_colored(f"{dists_pw[k]:.4f} A", 1, 1, 0.6)
                            k += 1
                    imgui.end_child()
                    if len(dists_pw) > max_rows:
                        imgui.text_colored(f"Showing {max_rows} / {len(dists_pw)} pairs (UI Cap)", 1, 0.8, 0.4)
                    imgui.text(f"Min: {np.min(dists_pw):.2f} | Max: {np.max(dists_pw):.2f} | Mean: {np.mean(dists_pw):.2f}")
            else:
                if imgui.collapsing_header("Pairwise Distances [HIDDEN]"):
                    imgui.text_colored("Pairwise O(N²) calculation is manual for performance.", 1, 0.5, 0.5)
                    if imgui.button("Compute Full Pairwise Matrix"):
                        self.force_full_pairwise = True
            
            # --- Sequential Analytics (Chain) ---
            if len(selected) >= 2:
                if imgui.collapsing_header("Consecutive Chains", imgui.TREE_NODE_DEFAULT_OPEN)[0]:
                    imgui.begin_child("chain_scroll", 0, 120, border=True)
                    # UI Cap: Max 100 rows
                    max_chain = 100
                    for i in range(min(len(dists_seq), max_chain)):
                        imgui.text(f"Dist {selected[i]}-{selected[i+1]}:")
                        imgui.same_line(150); imgui.text_colored(f"{dists_seq[i]:.4f} A", 1, 1, 0.6)
                        if i < len(angles_seq):
                            imgui.text(f"Angle {selected[i]}-{selected[i+1]}-{selected[i+2]}:")
                            imgui.same_line(150); imgui.text_colored(f"{angles_seq[i]:.2f} deg", 0.6, 1, 1)
                    imgui.end_child()
                    if len(dists_seq) > max_chain:
                        imgui.text_colored(f"Showing {max_chain} / {len(dists_seq)} steps (UI Cap)", 1, 0.8, 0.4)

            # --- Z-Matrix View ---
            if imgui.collapsing_header("Z-Matrix Representation"):
                imgui.begin_child("zmat_scroll", 0, 200, border=True)
                imgui.columns(5, "zmat_cols")
                imgui.text("Atom"); imgui.next_column()
                imgui.text("Dist"); imgui.next_column()
                imgui.text("Angle"); imgui.next_column()
                imgui.text("Dihedral"); imgui.next_column()
                imgui.text("Refs"); imgui.next_column()
                imgui.separator()
                
                max_rows = 100
                for i in range(min(len(selected), max_rows)):
                    idx = selected[i]
                    elem = self.scene.elements[idx] if idx < len(self.scene.elements) else "?"
                    imgui.text(f"{idx}({elem})"); imgui.next_column()
                    
                    if i >= 1: imgui.text(f"{dists_seq[i-1]:.3f}"); imgui.next_column()
                    else: imgui.text("-"); imgui.next_column()
                    
                    if i >= 2: imgui.text(f"{angles_seq[i-2]:.1f}"); imgui.next_column()
                    else: imgui.text("-"); imgui.next_column()
                    
                    if i >= 3: imgui.text(f"{dihedrals_seq[i-3]:.1f}"); imgui.next_column()
                    else: imgui.text("-"); imgui.next_column()
                    
                    refs = []
                    if i >= 1: refs.append(str(selected[i-1]))
                    if i >= 2: refs.append(str(selected[i-2]))
                    if i >= 3: refs.append(str(selected[i-3]))
                    imgui.text(",".join(refs)); imgui.next_column()
                
                imgui.columns(1)
                imgui.end_child()
                if imgui.button("Copy Z-Matrix to Clipboard"):
                    lines = ["# SAGE Z-Matrix Export"]
                    for i in range(len(selected)):
                        idx = selected[i]
                        elem = self.scene.elements[idx] if idx < len(self.scene.elements) else "?"
                        line = f"{elem:<2}"
                        if i >= 1: line += f" {selected[i-1]:>4} {dists_seq[i-1]:>10.5f}"
                        if i >= 2: line += f" {selected[i-2]:>4} {angles_seq[i-2]:>10.3f}"
                        if i >= 3: line += f" {selected[i-3]:>4} {dihedrals_seq[i-3]:>10.3f}"
                        lines.append(line)
                    glfw.set_clipboard_string(self.win, "\n".join(lines))

            # --- Distribution Shortcut ---
            if len(selected) > 5:
                imgui.separator()
                imgui.text("Distance Distribution:")
                plot_data = dists_pw if dists_pw else dists_seq
                imgui.plot_histogram("", np.array(plot_data, dtype=np.float32), graph_size=(0, 60))
        
        imgui.end()

    def _calc_bonds_simple(self, scene: Optional[Scene] = None):
        if scene is None: scene = self.scene
        # Use appropriate index for scene
        s_idx = self.index1 if scene is self.scene else self.index2
        
        from .analysis.coordination import compute_coordination_sphere
        key = (id(scene), s_idx, self.bond_factor)
        if key in self._bond_cache:
            scene.bonds = self._bond_cache[key]
            scene.bond_count = len(scene.bonds)
            return

        # Use the tunable covalent factor
        _, _, _, bonds = compute_coordination_sphere(scene.pos, scene.elements, scene.lat, factor=self.bond_factor)
        
        scene.bonds = np.array(bonds, dtype=np.int32)
        scene.bond_count = len(scene.bonds)
        self._bond_cache[key] = scene.bonds

    def _get_clipped_indices(self, scene: Scene):
        """Returns indices of atoms that pass the current clipping filter."""
        if not self.clipping_enabled:
            return np.arange(scene.N)
            
        if self.clipping_fractional:
            inv_lat = np.linalg.inv(scene.lat)
            frac = scene.pos @ inv_lat
            vals = frac[:, self.clipping_axis]
        else:
            vals = scene.pos[:, self.clipping_axis]
            
        mask = (vals >= self.clipping_range[0]) & (vals <= self.clipping_range[1])
        if self.clipping_invert: mask = ~mask
        return np.where(mask)[0]

    def _calc_rdf(self):
        """Computes all-species-pairs RDF and automatically shows the plot."""
        if not self.scene.N: 
            self.rdf_results.clear()
            return
            
        self.show_rdf = True # Auto-open window on calculation
        from .analysis.rdf import compute_rdf
        self.rdf_results.clear()
        
        # 1. Global (All-All)
        indices = self._get_clipped_indices(self.scene)
        if not len(indices): return
        
        pos_clipped = self.scene.pos[indices]
        r, g = compute_rdf(pos_clipped, pos_clipped, self.scene.lat, 
                           r_max=self.rdf_rmax, bins=self.rdf_bins, mode=self.rdf_mode)
        self.rdf_results["Global"] = (r, g)
        
        # 2. All Unique Pairs (Species-A - Species-B)
        species = self.scene.unique_elements
        if len(species) > 1:
            for i, sa in enumerate(species):
                idx_a = [k for k in indices if self.scene.elements[k] == sa]
                if not idx_a: continue
                # Diagonal: A-A
                # Cross: A-B
                for j in range(i, len(species)):
                    sb = species[j]
                    idx_b = [k for k in indices if self.scene.elements[k] == sb]
                    if not idx_a or not idx_b: continue
                    
                    r_p, g_p = compute_rdf(self.scene.pos[idx_a], self.scene.pos[idx_b], self.scene.lat,
                                           r_max=self.rdf_rmax, bins=self.rdf_bins, mode=self.rdf_mode)
                    lbl = f"{sa}-{sb}"
                    self.rdf_results[lbl] = (r_p, g_p)
                    if lbl not in self.rdf_visibility: self.rdf_visibility[lbl] = True
            
        # Ensure Global is also visible
        if "Global" not in self.rdf_visibility: self.rdf_visibility["Global"] = True

    def _calc_z_profile(self):
        """Computes Z-density profile for the current system."""
        if not self.scene.N: return
        from .analysis.rdf import compute_z_profile
        
        # We could also filter here by species in the future
        indices = self._get_clipped_indices(self.scene)
        if not len(indices): return
        x, y = compute_z_profile(self.scene.pos[indices])
        self.z_profile_x, self.z_profile_y = x, y

    def _calc_voronoi(self):
        """Computes Voronoi tessellation and updates scene state."""
        from .analysis.voronoi import compute_voronoi
        indices = self._get_clipped_indices(self.scene)
        if len(indices) > 0:
            res = compute_voronoi(self.scene.pos[indices], self.scene.lat)
            self.voro_results = res
            self.scene.voro_volumes = res["volumes"]
            print(f"Voronoi: Mean Volume = {np.mean(res['volumes']):.2f} A^3")

    def _update_coloring(self, scene: Optional[Scene] = None):
        """Centralized controller for atomic coloring modes."""
        if scene is None: 
            self._update_coloring(self.scene)
            if self.scene2: self._update_coloring(self.scene2)
            return
            
        if not scene.N: return
        
        # 1. Calculate scalars based on active mode
        scalars = np.zeros(scene.N, dtype=np.float32)
        
        if self.color_mode == "Element":
            from .constants import CPK_COLORS
            scene.colors = np.array([CPK_COLORS.get(e, [0.5, 0.5, 0.5]) for e in scene.elements], dtype=np.float32)
            self._apply_selection_highlight(scene)
            self._sync_gpu(scene)
            return

        elif self.color_mode == "Z Position":
            if self.color_z_fractional and scene.lat is not None:
                inv_lat = np.linalg.inv(scene.lat)
                frac = scene.pos @ inv_lat
                scalars = frac[:, 2]
            else:
                scalars = scene.pos[:, 2]

        elif self.color_mode == "Displacement":
            if self.displacement_ref_pos is not None and len(self.displacement_ref_pos) == scene.N:
                scalars = np.linalg.norm(scene.pos - self.displacement_ref_pos, axis=1)

        elif self.color_mode == "Neighbors":
            from .analysis.coordination import compute_cn
            _, nb_counts = compute_cn(scene.pos, scene.lat, r_cut=self.color_nb_cutoff)
            scalars = nb_counts.astype(np.float32)

        elif self.color_mode == "Coordination":
            if hasattr(scene, "cn") and scene.cn is not None: 
                scalars = scene.cn
            else:
                from .analysis.coordination import compute_coordination_sphere
                # compute_coordination_sphere returns (avg, counts, avg_len, bonds)
                _, cnts, _, _ = compute_coordination_sphere(scene.pos, scene.elements, scene.lat, factor=self.bond_factor)
                scalars = cnts.astype(np.float32)

        elif self.color_mode == "Voronoi Volume":
            if hasattr(scene, "voro_volumes") and scene.voro_volumes is not None:
                scalars = scene.voro_volumes
            else:
                from .analysis.voronoi import compute_voronoi
                v_res = compute_voronoi(scene.pos, scene.lat)
                scene.voro_volumes = v_res.volumes
                scalars = scene.voro_volumes

        elif self.color_mode == "Selection":
            for idx in scene.selected_atoms:
                if idx < scene.N: scalars[idx] = 1.0

        elif self.color_mode == "Index":
            scalars = np.arange(scene.N, dtype=np.float32)

        # 2. Map to Palette
        scene.colors = self._map_to_palette(scalars)
        
        # 3. Enhance Selection & Sync
        self._apply_selection_highlight(scene)
        self._sync_gpu(scene)

    def _map_to_palette(self, scalars: np.ndarray) -> np.ndarray:
        """Turns a vector of scalars into RGB colors using CPK or linear maps."""
        if scalars.size == 0: return np.zeros((0, 3), dtype=np.float32)
        
        valid_mask = np.isfinite(scalars)
        if not np.any(valid_mask):
            return np.ones((len(scalars), 3), dtype=np.float32) * 0.5 # Gray for all-invalid
            
        v_min = np.nanmin(scalars[valid_mask]) if self.color_range_auto else self.color_range[0]
        v_max = np.nanmax(scalars[valid_mask]) if self.color_range_auto else self.color_range[1]
        
        if v_max == v_min: v_max += 1e-7
        
        t = np.zeros_like(scalars, dtype=np.float32)
        t[valid_mask] = np.clip((scalars[valid_mask] - v_min) / (v_max - v_min), 0.0, 1.0)
        # Handle Inf/Nan by mapping them to zero or middle? Let's use 0.0 for invalid.
        t[~valid_mask] = 0.0
        if self.color_invert: t = 1.0 - t
        
        colors = np.zeros((len(t), 3), dtype=np.float32)
        
        if self.color_palette == "BlueRed":
            # Blue -> White -> Red
            mask_lo = t < 0.5
            t_lo = t[mask_lo] * 2.0
            colors[mask_lo] = np.column_stack((t_lo, t_lo, np.ones_like(t_lo)))
            mask_hi = t >= 0.5
            t_hi = 1.0 - (t[mask_hi] - 0.5) * 2.0
            colors[mask_hi] = np.column_stack((np.ones_like(t_hi), t_hi, t_hi))
            
        elif self.color_palette == "Viridis":
            # Approx Viridis
            colors[:, 0] = t**2
            colors[:, 1] = t
            colors[:, 2] = 0.5 * (1.0 - t)
            
        elif self.color_palette == "Magma":
            colors[:, 0] = np.clip(t * 1.2, 0, 1)
            colors[:, 1] = t**3
            colors[:, 2] = 0.2 + 0.3 * t
            
        elif self.color_palette == "Hot":
            colors[:, 0] = np.clip(t * 3.0, 0, 1)
            colors[:, 1] = np.clip(t * 3.0 - 1.0, 0, 1)
            colors[:, 2] = np.clip(t * 3.0 - 2.0, 0, 1)
            
        elif self.color_palette == "Grayscale":
            colors[:, 0] = t; colors[:, 1] = t; colors[:, 2] = t
            
        return colors

    def _apply_selection_highlight(self, scene: Scene):
        """Force yellow highlight on selected atoms regardless of mode."""
        if not self.color_highlight_selection or not scene.selected_atoms: return
        for idx in scene.selected_atoms:
            if idx < scene.N:
                scene.colors[idx] = np.array([1.0, 0.9, 0.1], dtype=np.float32)

    def save_screenshot(self):
        """Captures the current viewport and saves to PNG."""
        from PIL import Image
        import time
        # Get actual framebuffer size (handles Retina/High-DPI)
        w, h = glfw.get_framebuffer_size(self.win)
        pixels = glReadPixels(0, 0, w, h, GL_RGB, GL_UNSIGNED_BYTE)
        img = Image.frombytes("RGB", (w, h), pixels)
        img = img.transpose(Image.FLIP_TOP_BOTTOM)
        fname = f"sage_snapshot_{int(time.time())}.png"
        img.save(fname)
        self.export_status = f"Saved: {fname}"

    def export_gif(self):
        """Generates a GIF animation from the trajectory with smart sampling."""
        import imageio
        from PIL import Image
        import random
        import time
        
        N_total = len(self.provider)
        count = min(self.gif_frame_count, N_total)
        
        if self.gif_random_selection:
            indices = sorted(random.sample(range(N_total), count))
        else:
            indices = np.linspace(0, N_total-1, count, dtype=int)
            
        frames = []
        original_idx = self.index1
        self.export_status = f"Rendering {count} frames..."
        
        # We need to render each frame to capture it
        for idx in indices:
            self.index1 = idx
            self._load_frame1(idx)
            
            # 1. Clear & Setup Viewport
            glClearColor(*self.bg_color, 1.0)
            glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
            w, h = glfw.get_framebuffer_size(self.win)
            glViewport(0, 0, w, h)
            
            # 2. Draw scene (Clean, no GUI)
            self._render_pass(self.scene, self.cam1, 0, 0, w, h)
            
            # 3. Capture
            pixels = glReadPixels(0, 0, w, h, GL_RGB, GL_UNSIGNED_BYTE)
            img = Image.frombytes("RGB", (w, h), pixels)
            img = img.transpose(Image.FLIP_TOP_BOTTOM)
            # Scale down if too large to save space/time
            if w > 800:
                img.thumbnail((800, 800), Image.Resampling.LANCZOS)
            frames.append(np.array(img))
            
            glfw.swap_buffers(self.win) # Visual feedback during export
            
        fname = f"sage_movie_{int(time.time())}.gif"
        self.export_status = "Encoding GIF (please wait)..."
        imageio.mimsave(fname, frames, fps=8, loop=0)
        self.export_status = f"Exported: {fname}"
        
        # Restore original state
        self.index1 = original_idx
        self._load_frame1(original_idx)

    def export_current_structure(self):
        """Exports the current scene structure using the selected format, respecting clipping and selection."""
        import time
        from .io.exporters import export_structure
        fmt = self.export_formats[self.export_fmt_idx]
        fname = f"{self.export_prefix}_{int(time.time())}.{fmt}"
        
        # Determine subset of indices to export
        if self.export_selected_only and self.scene.selected_atoms:
            indices = np.array(self.scene.selected_atoms, dtype=int)
        else:
            indices = self._get_clipped_indices(self.scene)

        if not len(indices):
            self.export_status = "Error: No atoms in visible/selected region to export."
            return

        if self.export_all_frames:
            n_frames = len(self.provider)
            self.export_status = f"Exporting {n_frames} frames..."
            succ_count = 0
            
            for i in range(n_frames):
                p, l, e, els, c, r = self.provider.get(i)
                pos_all = np.asarray(p, dtype=np.float32).reshape(-1, 3)
                
                # Apply current mask (clipping or selection) to this frame
                if self.export_selected_only and self.scene.selected_atoms:
                    frame_indices = np.array(self.scene.selected_atoms, dtype=int)
                elif self.clipping_enabled:
                    if self.clipping_fractional:
                        inv_l = np.linalg.inv(l)
                        vals = pos_all @ inv_l[:, self.clipping_axis]
                    else:
                        vals = pos_all[:, self.clipping_axis]
                    mask = (vals >= self.clipping_range[0]) & (vals <= self.clipping_range[1])
                    if self.clipping_invert: mask = ~mask
                    frame_indices = np.where(mask)[0]
                else:
                    frame_indices = np.arange(len(pos_all))
                
                if not len(frame_indices): continue
                
                ok, _ = export_structure(
                    fname, fmt, pos_all[frame_indices], l, [els[k] for k in frame_indices], 
                    energy=e, append=(succ_count > 0)
                )
                if ok: succ_count += 1
            
            self.export_status = f"Exported {succ_count} frames to {fname}"
        else:
            success, msg = export_structure(
                fname, fmt, 
                self.scene.pos[indices], 
                self.scene.lat, 
                [self.scene.elements[k] for k in indices], 
                energy=self.scene.energy,
                append=False
            )
            self.export_status = f"Exported to {fname}" if success else msg

    def _get_clipped_indices(self, scene: Scene) -> np.ndarray:
        """Returns indices of atoms that are currently within the clipping box."""
        if not self.clipping_enabled:
            return np.arange(scene.N)
        
        if self.clipping_fractional:
            inv_lat = np.linalg.inv(scene.lat)
            vals = (scene.pos @ inv_lat)[:, self.clipping_axis]
        else:
            vals = scene.pos[:, self.clipping_axis]
            
        mask = (vals >= self.clipping_range[0]) & (vals <= self.clipping_range[1])
        if self.clipping_invert: mask = ~mask
        return np.where(mask)[0]

    def _calc_cn_distribution(self):
        """Calculates coordination numbers based on a simple distance cutoff."""
        self.export_status = "Calculating CN..."
        N = self.scene.N
        if N == 0: return
        
        pos = self.scene.pos
        cutoff = self.cn_cutoff
        
        # Simple N^2 for now, or usage of existing graph
        from scipy.spatial.distance import pdist, squareform
        dists = squareform(pdist(pos))
        cn = np.sum((dists > 0.0001) & (dists <= cutoff), axis=1)
        self.cn_results = cn
        self.export_status = "CN Computed."

    def _run_tb_calculation(self):
        """Runs the tight-binding pipeline on the current scene."""
        from .analysis.tight_binding import run_tight_binding
        N = self.scene.N
        if N == 0:
            self.tb_status = "No atoms loaded."
            return
        if N > 2000:
            self.tb_status = f"Warning: N={N} > 2000, may be slow..."
        
        self.tb_status = f"Running TB (N={N})..."
        try:
            lat = self.scene.lat if self.scene.is_periodic else None
            self.tb_results = run_tight_binding(
                self.scene.pos, self.scene.elements, self.tb_params, lat
            )
            res = self.tb_results
            self.tb_status = (
                f"Done. HOMO={res.homo:+.2f}eV  Gap={res.gap:.2f}eV"
            )
            # Auto-show DOS
            self.tb_show_dos = True
        except Exception as e:
            self.tb_status = f"Error: {e}"
            self.tb_results = None

    def _reconstruct_tb_field(self):
        """Phase 4/5: Reconstructs 3D scalar fields and uploads them to the VolumeRenderer."""
        if self.tb_results is None: return
        
        from .analysis.tb_visualization import (
            make_grid_from_atoms, build_state_field, build_total_density_field
        )
        
        mode = self.tb_viz_mode
        self.tb_status = f"Reconstructing {mode} field..."
        
        try:
            grid = make_grid_from_atoms(self.scene.pos, self.tb_params)
            res = self.tb_results
            
            # Extract atom_to_orbs and orb_info (rebuild if not stored)
            if res.orb_info and res.atom_to_orbs:
                ato, oi = res.atom_to_orbs, res.orb_info
            else:
                from .analysis.tight_binding import build_hamiltonian
                lat = self.scene.lat if self.scene.is_periodic else None
                _, _, ato, oi = build_hamiltonian(self.scene.pos, self.scene.elements, self.tb_params, lat)
            
            field = None
            if mode == "HOMO":
                idx = res.n_occ_states - 1
                if idx >= 0:
                    field = build_state_field(self.scene.pos, self.scene.elements, oi, res.eigenvectors, idx, grid)
                    field.values = np.abs(field.values) # Show both lobes
            elif mode == "LUMO":
                idx = res.n_occ_states
                if idx < res.eigenvectors.shape[1]:
                    field = build_state_field(self.scene.pos, self.scene.elements, oi, res.eigenvectors, idx, grid)
                    field.values = np.abs(field.values) # Show both lobes
            elif mode == "Total Density":
                field = build_total_density_field(self.scene.pos, self.scene.elements, oi, 
                                                 res.eigenvectors, res.occupations, grid)
            
            if field is not None:
                # Upload to GPU Volume Renderer
                self.rnd_volume.upload_data(field.values, field.grid.origin, np.array([field.grid.spacing]*3))
                self.tb_status = f"Done: {mode} reconstructed."
                self.tb_show_volume = True
            else:
                self.tb_status = "Field reconstruction failed (Invalid mode or index)."
                
        except Exception as e:
            self.tb_status = f"Error: {e}"
            print(f"TB Viz Error: {e}")

    def _apply_tb_coloring(self):
        """Maps TB charge or bond order onto atom sphere colors."""
        if self.tb_results is None or self.tb_color_mode == "None":
            self._update_coloring()
            return
        
        N = self.scene.N
        rgba = np.ones((N, 4), dtype=np.float32)
        rgba[:, :3] = self._get_element_colors(self.scene)[:, :3]
        
        if self.tb_color_mode == "TB Charge" and len(self.tb_results.local_charge) == N:
            q = self.tb_results.local_charge.copy()
            q_min, q_max = q.min(), q.max()
            if q_max > q_min:
                t = (q - q_min) / (q_max - q_min)
            else:
                t = np.full(N, 0.5)
            # Blue=low → Red=high
            rgba[:, 0] = t
            rgba[:, 1] = 0.3 * (1 - t)
            rgba[:, 2] = 1.0 - t
        
        elif self.tb_color_mode == "TB Bond Order":
            # Color atoms by their max bond order to any neighbor
            bo_map = np.zeros(N, dtype=np.float32)
            for (i, j), bo in self.tb_results.bond_order.items():
                bo_map[i] = max(bo_map[i], abs(bo))
                if j < N:
                    bo_map[j] = max(bo_map[j], abs(bo))
            bo_max = bo_map.max()
            if bo_max > 0:
                t = bo_map / bo_max
                rgba[:, 0] = 0.2 + 0.8 * t
                rgba[:, 1] = 0.8 * t
                rgba[:, 2] = 0.4 * (1 - t)
        
        self.rnd_sphere.upload_data(self.scene.pos, self.scene.radii, rgba)

    def _update_clipping_alpha(self):
        """Updates the alpha channel of atom colors based on clipping range and alpha."""
        if not self.clipping_enabled:
            # Revert to full opaque if clipping was disabled
            if hasattr(self, "_clipping_active") and self._clipping_active:
                self._update_coloring() # Reset to original colors
                self._clipping_active = False
            return

        self._clipping_active = True
        
        # 1. Calculate clipping mask
        if self.clipping_fractional:
            ext_lat = self.scene.lat
            inv_lat = np.linalg.inv(ext_lat)
            pos_frac = self.scene.pos @ inv_lat
            vals = pos_frac[:, self.clipping_axis]
        else:
            vals = self.scene.pos[:, self.clipping_axis]
        
        mask_in_range = (vals >= self.clipping_range[0]) & (vals <= self.clipping_range[1])
        if self.clipping_invert: 
            mask_inner = mask_in_range
        else:
            mask_inner = mask_in_range

        # The user said "alpha de los que NO se deben mostrar"
        # If invert=False: In-range is visible (alpha=1), out-of-range is ghosted (alpha=clipping_alpha)
        # If invert=True: Out-of-range is visible, in-range is ghosted
        if self.clipping_invert:
            mask_ghost = mask_in_range
        else:
            mask_ghost = ~mask_in_range

        # 2. Apply colors
        # Start from base element colors or current analysis colors
        # Ensure we don't accidentally recurse if _update_coloring calls us
        rgba = np.ones((self.scene.N, 4), dtype=np.float32)
        base_colors = self._get_element_colors(self.scene)
        rgba[:, :3] = base_colors[:, :3]
        
        rgba[:, 3] = np.where(mask_ghost, self.clipping_alpha, 1.0)
        
        # 3. Synchronize with SphereRenderer
        self.rnd_sphere.upload_data(self.scene.pos, self.scene.radii, rgba)
        
        # 4. Synchronize with BondRenderer
        if self.show_bonds and self.scene.bond_count > 0:
            self.rnd_bond.upload_data(self.scene.pos, self.scene.bonds, rgba, self.scene.lat)

    def _generate_gaussian_density(self):
        """Triggers density generation and uploads to GPU."""
        from .analysis.volumetric import generate_density_from_atoms
        self.export_status = "Generating Density..."
        
        res = getattr(self, "vol_res", 0.25)
        sigma = getattr(self, "vol_sigma", 0.5)
        
        self.vol_data = generate_density_from_atoms(
            self.scene.pos, self.scene.elements, self.scene.lat,
            resolution=res, sigma=sigma
        )
        self.vol_min = float(np.min(self.vol_data.grid))
        self.vol_max = float(np.max(self.vol_data.grid))
        self.vol_iso_level = self.vol_min + 0.2 * (self.vol_max - self.vol_min)
        
        self.rnd_volume.upload_data(self.vol_data.grid, self.vol_data.origin, self.vol_data.spacing)
        self.vol_enabled = True
        self.export_status = f"Generated {self.vol_data.shape} grid. Max={self.vol_max:.2f}"

    def _load_volumetric_cube(self):
        """Opens a file dialog (simulated) or loads from a preset path."""
        # In a real app, this would use a file picker. 
        # For now, we'll try to load a default or prompt for path if possible.
        from .analysis.volumetric import load_cube_file
        try:
            self.vol_data = load_cube_file(self._cube_path)
            self.vol_min = float(np.min(self.vol_data.grid))
            self.vol_max = float(np.max(self.vol_data.grid))
            self.vol_iso_level = self.vol_min + 0.2 * (self.vol_max - self.vol_min)
            
            self.rnd_volume.upload_data(self.vol_data.grid, self.vol_data.origin, self.vol_data.spacing)
            self.vol_enabled = True
            self.export_status = f"Loaded {self._cube_path}. Range=[{self.vol_min:.2f}, {self.vol_max:.2f}]"
        except Exception as e:
            self.export_status = f"Error loading .cube: {e}"
