#!/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
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.picking_renderer import PickingRenderer

# UI Subsystem
from .ui.plots import draw_rdf_plot, draw_bad_plot

@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.5, 'headlight': 0.8, 
            'shadow_intensity': 0.6, 'shadow_softness': 1.5
        }
        self.element_materials = {} # Dict[Element, MaterialDict]
        self._mat_tbo_dirty = True
        self.light_rotating = False
        self.show_shadows = False
        
        # 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.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
        
        # Output & Export Settings
        self.gif_frame_count = 10
        self.gif_random_selection = False
        self.export_status = ""
        
        # 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_shadows = False
        self.show_ground_shadow = False
        self.shadow_res_idx = 3
        self.win_size = (1200, 800)
        
        # Caches
        self._bond_cache: Dict[Tuple[int, float], np.ndarray] = {}
        
        # Initialize OpenGL
        self._init_gl()
        self._init_renderers()
        
        # Initial Load
        self.provider2: Optional[Any] = None
        self.scene2: Optional[Scene] = None
        
        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.rnd_pick = PickingRenderer()

    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 _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.scene.set_replicas(*self.scene.replicas)
            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.scene2.set_replicas(*self.scene.replicas)
            self._sync_gpu(self.scene2)
            # 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 _sync_gpu(self, scene: Optional[Scene] = None):
        """Updates GPU buffers with specified scene state."""
        if scene is None: scene = self.scene
        
        self.rnd_sphere.upload_data(scene.pos, scene.radii, scene.colors, scene.selected_atoms)
        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:
                # We need a local context check here
                if scene == self.scene: self._calc_bonds_simple(scene)
                else: self._calc_bonds_simple(scene)
            
            self.rnd_bond.upload_data(scene.pos, np.array(scene.bonds), 
                                      scene.colors, lattice=scene.lat)

    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()
            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 _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 = cam.get_projection_matrix(aspect, is_ortho=self.is_ortho)
            
        # 1. Shadow Depth Pass
        shadow_matrix = None
        if self.pretty_render and self.lights and self.show_shadows:
            l_dir = np.array(self.lights[0]['pos'])
            l_dir = l_dir / (np.linalg.norm(l_dir) + 1e-6)
            mn, mx = scene.get_full_bounds()
            center = (mn + mx) * 0.5
            size = np.linalg.norm(mx - mn) * 0.8
            l_pos = center + l_dir * size
            up = np.array([0,1,0], dtype=np.float32) if abs(l_dir[1]) < 0.9 else np.array([1,0,0], dtype=np.float32)
            z = l_dir.astype(np.float32); xi = np.cross(up, z); xi /= (np.linalg.norm(xi) + 1e-6)
            yi = np.cross(z, xi)
            l_view = np.eye(4, dtype=np.float32)
            l_view[:3,:3] = np.array([xi, yi, z], dtype=np.float32)
            l_view[:3, 3] = -l_view[:3,:3] @ l_pos.astype(np.float32)
            s_ortho = float(size) if size > 0 else 10.0
            l_proj = np.diag([1/s_ortho, 1/s_ortho, -2/(s_ortho*4), 1]).astype(np.float32)
            l_proj[2, 3] = 0.0
            
            # Use current scene data for shadow pass
            self._sync_gpu(scene)
            self.rnd_sphere.draw_shadows(l_view, l_proj, scene.N, scene.replica_offsets, self.atom_scale)
            if self.show_bonds:
                self.rnd_bond.draw_shadows(l_view, l_proj, scene.replica_offsets, 
                                           scene.bond_count, self.bond_scale,
                                           self.rnd_sphere.shadow_fbo, self.rnd_sphere.shadow_size)
            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

        # 2. Main Color Pass
        glViewport(x, y, w, h)
        self._sync_gpu(scene) # Ensure data matches the scene
        self._update_material_tbo()
        
        self.rnd_sphere.draw(view, proj, scene.N, scene.replica_offsets, 
                             ortho_scale=cam.ortho_scale, viewport_height=h, 
                             lights=self.lights, materials=self.materials,
                             atom_scale=self.atom_scale, is_ortho=self.is_ortho,
                             scene_bounds=scene.get_full_bounds(),
                             skip_px=self.quality.skip_px, point_px=self.quality.point_px,
                             pretty_render=self.pretty_render, shadow_matrix=shadow_matrix)
        if self.show_bonds:
            self.rnd_bond.draw(view, proj, scene.replica_offsets, scene.bond_count, 
                               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)
        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])

    def _update_material_tbo(self):
        """Syncs material and color TBOs with per-element settings."""
        if not self._mat_tbo_dirty: return
        
        from .constants import CPK_COLORS
        for e in set(self.scene.elements):
            if e not in self.element_materials:
                self.element_materials[e] = {
                    'col': list(CPK_COLORS.get(e, [0.5, 0.5, 0.5])),
                    'ambient': 0.4, 'specular': 0.6, 'shininess': 48.0, 'rim': 0.8
                }
        
        # Build Buffers
        mat_arr = np.zeros((self.scene.N, 4), dtype=np.float32)
        col_arr = np.zeros((self.scene.N, 3), dtype=np.float32)
        
        for i, el in enumerate(self.scene.elements):
            m = self.element_materials[el]
            mat_arr[i] = [m['ambient'], m['specular'], m['shininess'], m['rim']]
            col_arr[i] = m['col']
            
        self.rnd_sphere.upload_materials(mat_arr)
        # We need an upload_colors but can use upload_data(colors=...)
        # Instead, let's just make it simpler: upload_data re-uploads all.
        self.rnd_sphere.upload_data(self.scene.pos, self.scene.radii, col_arr, self.scene.selected_atoms)
        self._mat_tbo_dirty = False

    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.show_shadows = imgui.checkbox("Shadow Mapping (High Quality)", self.show_shadows)
            self.materials['show_shadows'] = self.show_shadows
            
            # 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.show_ground_shadow = imgui.checkbox("Ground Reflection/Shadow", self.show_ground_shadow)
            self.materials['ground_shadow'] = self.show_ground_shadow
            
            _, self.materials['exposure'] = imgui.slider_float("Exposure / Gamma", self.materials['exposure'], 0.1, 20.0)
            _, self.materials['headlight'] = imgui.slider_float("Headlight", self.materials['headlight'], 0.0, 5.0)
            if self.show_shadows:
                _, 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 Softness (Blur)", self.materials['shadow_softness'], 0.0, 5.0)
            self._render_gui_lighting()

        if imgui.collapsing_header("Material Library (Per Element)")[0]:
            for el, m in self.element_materials.items():
                if imgui.tree_node(f"[{el}] Properties"):
                    ch1, m['col'] = imgui.color_edit3(f"Color##{el}", *m['col'])
                    ch2, m['ambient'] = imgui.slider_float(f"Ambient##{el}", m['ambient'], 0.0, 5.0)
                    ch3, m['specular'] = imgui.slider_float(f"Specular##{el}", m['specular'], 0.0, 10.0)
                    ch4, m['shininess'] = imgui.slider_float(f"Shininess##{el}", m['shininess'], 1.0, 256.0)
                    ch5, m['rim'] = imgui.slider_float(f"Rim##{el}", m['rim'], 0.0, 10.0)
                    if any([ch1, ch2, ch3, ch4, ch5]): 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()

        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}")
            imgui.text(f"Energy: {self.scene.energy:.6f}")

        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()

            # 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)

        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):")
            if not hasattr(self, "_comp_path"): self._comp_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)
            _, self.pretty_render = imgui.checkbox("Pretty Render (HQ)", self.pretty_render)
            
            imgui.separator()
            _, 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, 5)
            changed_ny, ry = imgui.slider_int("ny", ry, 0, 5)
            changed_nz, rz = imgui.slider_int("nz", rz, 0, 5)
            if changed_nx or changed_ny or changed_nz:
                self.scene.set_replicas(rx, ry, rz)
                self._sync_gpu()
            
            imgui.separator()
            
            changed_b, self.show_bonds = imgui.checkbox("Show Bonds", self.show_bonds)
            if changed_b: self._sync_gpu()
            if self.show_bonds:
                _, self.bond_scale = imgui.slider_float("Bond Width", self.bond_scale, 0.05, 1.0)
            
            _, self.show_cell = imgui.checkbox("Show Unit Cell", self.show_cell)
            _, self.bg_color = imgui.color_edit3("Background", *self.bg_color)

        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._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._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._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]
            
            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()

            # --- 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()

        # 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("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()
            
            if self.export_status:
                imgui.text_colored(self.export_status, 0.4, 1.0, 0.4)

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

    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:
                             if idx in self.scene.selected_atoms: self.scene.selected_atoms.remove(idx)
                             else: self.scene.selected_atoms.append(idx)
                         else: self.scene.selected_atoms = []
                    else:
                         self._select_atoms_in_box()
                    self.box_selecting = False
                    self._sync_gpu()
                self.orbiting = self.panning = self.light_rotating = False

    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:
            target_cam.rotate_yaw(0.005 * dx)
            target_cam.rotate_pitch(-0.005 * dy)
            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.center = np.copy(target_cam.center)
        
        if 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:
            if key == glfw.KEY_ESCAPE: glfw.set_window_should_close(win, True)
            if key == glfw.KEY_H: self.show_gui = not self.show_gui

    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
        # Re-render matrices for picking (must match main render)
        view = look_at(self.cam.eye(), self.cam.target(), WORLD_UP)
        aspect = w / h if h > 0 else 1.0
        if self.is_ortho:
            s = self.cam.ortho_scale / 2.0
            proj = ortho(-s*aspect, s*aspect, -s, s, -2000.0, 2000.0)
        else:
            proj = perspective(self.cam.fov, aspect, 0.1, 5000.0)
            
        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)
        
        # Current Matrices
        aspect = w / h if h > 0 else 1.0
        view = look_at(self.cam.eye(), self.cam.target(), WORLD_UP)
        if self.is_ortho:
            s = self.cam.ortho_scale / 2.0
            proj = ortho(-s*aspect, s*aspect, -s, s, -2000.0, 2000.0)
        else:
            proj = perspective(self.cam.fov, aspect, 0.1, 5000.0)
        
        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)
        
        added_any = False
        # Project each atom and check if within bounds
        # Note: self.scene.pos is (N, 3).
        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 # Behind camera?
            # Normalize device coordinates (NDC)
            ndc_x, ndc_y = pos_v[0] / pos_v[3], pos_v[1] / pos_v[3]
            # Device coordinates ((-1, -1) to (1, 1)) -> Screen coordinates ((0,0) to (win_w, win_h))
            # Screen origin is top-left in GLFW callbacks
            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:
                if i not in self.scene.selected_atoms:
                    self.scene.selected_atoms.append(i)
                    added_any = True
        
        if not added_any:
            self.scene.selected_atoms = [] # Clear if nothing found? Usually behavior is keep or clear.
            # 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 _calc_bonds_simple(self, scene: Optional[Scene] = None):
        """Ultra-fast simple bonding based on CPK-like cutoffs with caching."""
        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 _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)
        r, g = compute_rdf(self.scene.pos, self.scene.pos, 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, e in enumerate(self.scene.elements) if e == 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, e in enumerate(self.scene.elements) if e == sb]
                    if 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
        x, y = compute_z_profile(self.scene.pos)
        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
        if self.scene.N > 0:
            res = compute_voronoi(self.scene.pos, 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)
        
        v_min = np.min(scalars) if self.color_range_auto else self.color_range[0]
        v_max = np.max(scalars) if self.color_range_auto else self.color_range[1]
        
        if v_max == v_min: v_max += 1e-7
        t = np.clip((scalars - v_min) / (v_max - v_min), 0.0, 1.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)
