"""
Interactive Viewer module for GNDC.

This module provides a full-featured GUI viewer for visualizing
compressed geospatial time-series data.
"""

import tkinter as tk
from tkinter import filedialog, ttk, messagebox
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import matplotlib.dates as mdates
import matplotlib as mpl
import torch
import numpy as np
import os
import rasterio
from rasterio.enums import Resampling
from datetime import datetime, timedelta
from PIL import Image, ImageTk, ImageDraw, ImageFont

from pygndc.reader.reader import GNDCReader
from pygndc.utils.date_utils import normalize_time

try:
    import tinycudann as _tcnn  # noqa: F401
    _HAS_TCNN = True
except ImportError:
    _HAS_TCNN = False

try:
    from scipy.signal import savgol_filter
except ImportError:
    savgol_filter = None


class GNDCInteractiveViewer:
    """Full-featured interactive viewer for GNDC compressed data."""
    
    def __init__(self, master=None):
        if master is None:
            self.root = tk.Tk()
            self.root.title("GeoNDC Neural Earth Viewer")
        else:
            self.root = master
            self.root.title("GeoNDC Neural Earth Viewer")

        self.mode_var = tk.StringVar(value="auto")
        self.reader = GNDCReader(mode="auto")
        self.mode_var.set(self.reader.mode)
        self.device = self.reader.device
        
        self.is_loaded = False
        self.meta = {}
        self.H, self.W, self.C = 0, 0, 0
        
        self.full_flat_rows = None 
        self.full_flat_cols = None 
        self.active_flat_rows = None 
        self.active_flat_cols = None 
        self.has_mask = False
        self.valid_indices = None 
        self.cached_spatial_coords = None 
        
        self.tk_image_ref = None  
        self.canvas_image_id = None 
        self.global_vmax = None   
        self.bands_idx = None
        self._cbar_cache = {}
        self.show_colorbar = tk.BooleanVar(value=True)
        
        self.img_scale = 1.0
        self.img_offset_x = 0
        self.img_offset_y = 0
        
        self.crop_mode = False
        self.crop_rect = None
        self.crop_rect_id = None
        self.crop_start = None
        self.cached_pil_img = None
        self.cropped_region = None
        
        self.valid_dates_set = set()
        self.obs_datetimes = [] 
        self.current_dt = None

        self.ts_band_var = tk.StringVar(value='ALL')
        self.display_mode = tk.StringVar(value='RGB')
        self.res_var = tk.StringVar(value="0")
        self.smooth_method = tk.StringVar(value='None')
        self.last_rendered_ts = 0
        self.external_mask_bool = None
        self.show_current_time_on_ts = tk.BooleanVar(value=True)
        self.apply_mask = tk.BooleanVar(value=True)

        self.mask_is_static = True
        self.last_mask_t_idx = -1

        self.setup_ui()

    def setup_ui(self):
        BASE_FONT_SIZE = 12
        UI_FONT = ("Arial", BASE_FONT_SIZE)
        BOLD_FONT = ("Arial", BASE_FONT_SIZE + 2, "bold")

        self.root.option_add("*Font", UI_FONT)
        style = ttk.Style()
        style.configure(".", font=UI_FONT)
        style.configure("TButton", font=UI_FONT)
        style.configure("TCombobox", font=UI_FONT)
        style.configure("TRadiobutton", font=UI_FONT)

        self.main_canvas = tk.Canvas(self.root, bg="#202020", highlightthickness=0)
        self.main_canvas.pack(side=tk.TOP, fill=tk.BOTH, expand=True)

        menubar = tk.Menu(self.root)
        
        file_menu = tk.Menu(menubar, tearoff=0, font=("Arial", 12))
        file_menu.add_command(label="Open Model...", command=self.load_model)
        file_menu.add_command(label="Load Mask (Optional)...", command=self.load_mask)
        file_menu.add_separator()
        file_menu.add_command(label="Save Current View as TIF (Visual)...", command=self.save_current_view_as_tif)
        file_menu.add_separator()
        file_menu.add_command(label="Export Current Frame (Raw Data)...", command=self.save_current_frame_raw)
        file_menu.add_command(label="Export Time Series (Raw Data)...", command=self.save_time_series_raw)
        file_menu.add_separator()
        file_menu.add_command(label="Export GIF Animation...", command=self.export_gif_animation)
        file_menu.add_separator()
        file_menu.add_command(label="Exit", command=self.root.quit)
        menubar.add_cascade(label="File", menu=file_menu)

        option_menu = tk.Menu(menubar, tearoff=0, font=("Arial", 12))

        # Display Settings submenu
        display_settings_menu = tk.Menu(option_menu, tearoff=0, font=("Arial", 12))
        display_settings_menu.add_command(label="Display Settings...", command=self.show_display_settings_dialog)
        display_settings_menu.add_separator()
        display_settings_menu.add_checkbutton(
            label="Show Colorbar",
            variable=self.show_colorbar,
            onvalue=True,
            offvalue=False
        )
        display_settings_menu.add_checkbutton(
            label="Apply Mask",
            variable=self.apply_mask,
            onvalue=True,
            offvalue=False,
            command=self._on_apply_mask_changed
        )
        option_menu.add_cascade(label="Display", menu=display_settings_menu)

        # TimeSeries Settings submenu
        ts_settings_menu = tk.Menu(option_menu, tearoff=0, font=("Arial", 12))
        ts_settings_menu.add_command(label="TimeSeries Settings...", command=self.show_timeseries_settings_dialog)
        ts_settings_menu.add_separator()
        ts_settings_menu.add_checkbutton(
            label="Show Current Time Line",
            variable=self.show_current_time_on_ts,
            onvalue=True,
            offvalue=False
        )
        option_menu.add_cascade(label="TimeSeries", menu=ts_settings_menu)

        # Mode selection submenu
        mode_menu = tk.Menu(option_menu, tearoff=0, font=("Arial", 12))
        has_cuda = torch.cuda.is_available()
        mode_menu.add_radiobutton(
            label="tcnn_cuda (fastest)",
            variable=self.mode_var,
            value="tcnn_cuda",
            state="normal" if (_HAS_TCNN and has_cuda) else "disabled",
        )
        mode_menu.add_radiobutton(
            label="torch_gpu",
            variable=self.mode_var,
            value="torch_gpu",
            state="normal" if has_cuda else "disabled",
        )
        mode_menu.add_radiobutton(
            label="torch_cpu",
            variable=self.mode_var,
            value="torch_cpu",
        )
        option_menu.add_cascade(label="Mode", menu=mode_menu)

        menubar.add_cascade(label="Option", menu=option_menu)
        self.root.config(menu=menubar)
        
        self.main_canvas.bind("<Button-1>", self.on_canvas_press)
        self.main_canvas.bind("<B1-Motion>", self.on_canvas_drag)
        self.main_canvas.bind("<ButtonRelease-1>", self.on_canvas_release)
        self.main_canvas.bind("<Configure>", self.on_window_resize)

        control_frame = tk.Frame(self.root)
        control_frame.pack(side=tk.BOTTOM, fill=tk.X, padx=10, pady=10)

        btn_frame = tk.Frame(control_frame)
        btn_frame.pack(side=tk.TOP, fill=tk.X, pady=5)

        self.time_label = tk.Label(btn_frame, text="Date: YYYY-MM-DD", font=("Arial", 14, "bold"))
        self.time_label.pack(side=tk.LEFT, padx=15)

        self.crop_btn = tk.Button(btn_frame, text="Select Region", command=self.toggle_crop_mode, bg="#e0e0e0")
        self.crop_btn.pack(side=tk.LEFT, padx=5)
        ttk.Button(btn_frame, text="Reset View", command=self.reset_to_full_image).pack(side=tk.LEFT, padx=5)
        
        slider_frame = tk.Frame(control_frame)
        slider_frame.pack(side=tk.TOP, fill=tk.X, pady=5)
        self.time_slider = ttk.Scale(slider_frame, from_=0.0, to=1.0, 
                                     orient=tk.HORIZONTAL, command=self.on_slider_change, 
                                     length=400)
        self.time_slider.set(0) 
        self.time_slider.pack(side=tk.LEFT, fill=tk.X, expand=True)
        self.time_slider.bind("<Button-1>", self.jump_to_click_slider)

        step_frame = tk.Frame(slider_frame)
        step_frame.pack(side=tk.RIGHT, padx=5)
        ttk.Button(step_frame, text="< Prev Obs", command=lambda: self.step_time(-1)).pack(side=tk.LEFT)
        ttk.Button(step_frame, text="Next Obs >", command=lambda: self.step_time(1)).pack(side=tk.LEFT)

    def jump_to_click_slider(self, event):
        width = self.time_slider.winfo_width()
        if width == 0: return 
        ratio = np.clip(event.x / width, 0.0, 1.0)
        min_val = self.time_slider.cget("from")
        max_val = self.time_slider.cget("to")
        new_val = min_val + (max_val - min_val) * ratio
        self.time_slider.set(new_val)

    def on_window_resize(self, event):
        if self.is_loaded and self.current_dt:
            self.render_image_fast()

    def load_model(self):
        self.root.focus_force()
        self.root.update()
        path = filedialog.askopenfilename(defaultextension=".gndc", filetypes=[("GNDC Files", "*.gndc"), ("All Files", "*.*")])
        if path:
            try:
                if self.is_loaded or self.reader.model is not None:
                    print("Unloading previous model...")
                    self._reset_app_state()

                mode = self.mode_var.get()
                self.reader = GNDCReader(mode=mode)
                self.mode_var.set(self.reader.mode)
                self.device = self.reader.device
                self.reader.load(path)
                self.root.title(f"GeoNDC Viewer - {os.path.basename(path)}")
                
                self.meta, self.full_flat_rows, self.full_flat_cols, self.H, self.W, self.C = self._prepare_spatial_coords(self.reader)
                
                self.mask_is_static = True
                self.last_mask_t_idx = -1

                if self.reader.mask_vol is not None:
                    self.mask_is_static = getattr(self.reader, 'mask_is_static', True)
                    self.has_mask = True
                    print(f"Mask available (Type: {'Static' if self.mask_is_static else 'Dynamic'})")
                    if self.apply_mask.get():
                        print("Apply Mask is ON — filtering pixels")
                        self._update_mask_indices(0)
                        self.last_mask_t_idx = 0
                    else:
                        print("Apply Mask is OFF — using all pixels")
                        self.active_flat_rows = self.full_flat_rows
                        self.active_flat_cols = self.full_flat_cols
                        self.valid_indices = None
                else:
                    self.active_flat_rows = self.full_flat_rows
                    self.active_flat_cols = self.full_flat_cols
                    self.has_mask = False
                    self.valid_indices = None
                
                self.bands_idx = [0, 3, 2] if self.C >= 3 else [0]
                timestamps_str = self.meta.get('timestamps', [])
                self.obs_datetimes = [datetime.strptime(ts, "%Y-%m-%d %H:%M:%S") for ts in timestamps_str]
                self.valid_dates_set = {dt.strftime("%Y-%m-%d") for dt in self.obs_datetimes}

                min_ts, max_ts = self.obs_datetimes[0].timestamp(), self.obs_datetimes[-1].timestamp()
                self.time_slider.config(from_=min_ts, to=max_ts)
                self.time_slider.set(min_ts)

                self.is_loaded = True

                self.precompute_coords()
                self.on_slider_change(min_ts)
                
            except Exception as e:
                import traceback
                traceback.print_exc()
                messagebox.showerror("Error", str(e))

    def _reset_app_state(self):
        self.is_loaded = False
        self.reader.reset()
        self.meta = {}
        self.H, self.W, self.C = 0, 0, 0
        
        self.full_flat_rows = None 
        self.full_flat_cols = None 
        self.active_flat_rows = None 
        self.active_flat_cols = None 
        self.has_mask = False
        self.valid_indices = None 
        self.cached_spatial_coords = None 
        self.external_mask_bool = None
        
        self.tk_image_ref = None  
        self.canvas_image_id = None 
        self.global_vmax = None
        self.cached_pil_img = None
        
        self.cropped_region = None
        self.crop_mode = False
        self.crop_rect = None
        if self.crop_rect_id:
            self.main_canvas.delete(self.crop_rect_id)
            self.crop_rect_id = None
        self.crop_btn.config(relief=tk.RAISED)
        
        self.main_canvas.delete("all")
        
        self.obs_datetimes = []
        self.valid_dates_set = set()
        self.last_rendered_ts = -1.0
        
        import gc
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    def _prepare_spatial_coords(self, reader: GNDCReader):
        meta = reader.dataset_meta
        if meta is None:
            raise RuntimeError("No dataset metadata available.")

        T, H, W, C = meta['shape']
        
        rows = torch.arange(H, device=self.reader.device)
        cols = torch.arange(W, device=self.reader.device)
        g_row, g_col = torch.meshgrid(rows, cols, indexing='ij')
        
        flat_rows = g_row.flatten()
        flat_cols = g_col.flatten()
        
        return meta, flat_rows, flat_cols, H, W, C

    def _update_mask_indices(self, t_idx):
        if self.reader.mask_vol is None: return

        nodata_mask_np = self.reader.get_mask(t_idx) 
        if nodata_mask_np is None: return

        valid_mask_gpu = ~torch.from_numpy(nodata_mask_np).to(self.device).view(-1)
        self.valid_indices = torch.nonzero(valid_mask_gpu).squeeze()
        
        self.active_flat_rows = self.full_flat_rows[self.valid_indices]
        self.active_flat_cols = self.full_flat_cols[self.valid_indices]
        
        self.has_mask = True

    def _on_apply_mask_changed(self):
        """Re-render when Apply Mask toggle changes."""
        if not self.is_loaded:
            return
        if self.apply_mask.get():
            # Re-enable mask: recompute mask indices
            if self.reader.mask_vol is not None:
                t_idx = 0
                if self.obs_datetimes:
                    times = np.array([dt.timestamp() for dt in self.obs_datetimes])
                    curr_ts = self.time_slider.get()
                    t_idx = int((np.abs(times - curr_ts)).argmin())
                self._update_mask_indices(t_idx)
                self.last_mask_t_idx = t_idx
            elif self.external_mask_bool is not None:
                self.valid_indices = torch.tensor(
                    np.where(self.external_mask_bool.reshape(-1))[0],
                    device=self.device,
                )
                cpu_indices = self.valid_indices.cpu().numpy()
                self.active_flat_rows = self.full_flat_rows[cpu_indices]
                self.active_flat_cols = self.full_flat_cols[cpu_indices]
                self.has_mask = True
        else:
            # Disable mask: use all pixels
            self.active_flat_rows = self.full_flat_rows
            self.active_flat_cols = self.full_flat_cols
            self.valid_indices = None
        self.precompute_coords()
        self.global_vmax = None
        self.last_rendered_ts = -1.0
        self.on_slider_change(self.time_slider.get())

    def load_mask(self):
        if not self.is_loaded: return
        path = filedialog.askopenfilename(filetypes=[("GeoTIFF", "*.tif;*.tiff")])
        if not path: return
        try:
            with rasterio.open(path) as src:
                mask_data = src.read(1, out_shape=(self.H, self.W), resampling=Resampling.nearest)
                nodata = src.nodata
            
            is_valid = (mask_data != nodata) if nodata is not None else (mask_data != 0)
            self.external_mask_bool = is_valid.astype(bool)
            self.has_mask = True

            if self.apply_mask.get():
                self.valid_indices = torch.tensor(np.where(is_valid.reshape(-1))[0], device=self.device)
                cpu_indices = self.valid_indices.cpu().numpy()
                self.active_flat_rows = self.full_flat_rows[cpu_indices]
                self.active_flat_cols = self.full_flat_cols[cpu_indices]
            else:
                self.valid_indices = None
                self.active_flat_rows = self.full_flat_rows
                self.active_flat_cols = self.full_flat_cols
            
            self.precompute_coords()
            self.last_rendered_ts = -1.0
            self.on_slider_change(self.time_slider.get())
            
            print(f"External Mask Loaded: {path}")
            
        except Exception as e:
            messagebox.showerror("Mask Error", str(e))

    def precompute_coords(self):
        if self.active_flat_rows is None: return        
        rows_t = self.active_flat_rows if isinstance(self.active_flat_rows, torch.Tensor) else torch.tensor(self.active_flat_rows)
        cols_t = self.active_flat_cols if isinstance(self.active_flat_cols, torch.Tensor) else torch.tensor(self.active_flat_cols)
        
        self.cached_spatial_coords = self.reader._coords_from_rowcol(
            rows_t.to(self.device),
            cols_t.to(self.device),
            offset_x=0.5, 
            offset_y=0.5
        )

    def step_time(self, step_dir):
        if not self.is_loaded or not self.obs_datetimes:
            return

        all_timestamps = np.array([dt.timestamp() for dt in self.obs_datetimes])
        current_ts = self.time_slider.get()
        idx = (np.abs(all_timestamps - current_ts)).argmin()
        new_idx = np.clip(idx + step_dir, 0, len(all_timestamps) - 1)
        target_ts = all_timestamps[new_idx]
        
        if target_ts != current_ts:
            self.time_slider.set(target_ts)
            print(f"Jump to Frame {new_idx+1}/{len(all_timestamps)}: {self.obs_datetimes[new_idx]}")

    def on_slider_change(self, val):
        if not self.is_loaded: return
        ts = float(val)

        if hasattr(self, 'last_rendered_ts'):
            if abs(ts - self.last_rendered_ts) < 10.0:
                return
        self.last_rendered_ts = ts
        self.current_dt = datetime.fromtimestamp(ts)
        d_str = self.current_dt.strftime("%Y-%m-%d")
        color = "blue" if d_str in self.valid_dates_set else "black"
        self.time_label.config(text=f"Date: {d_str}", fg=color)
        
        self.render_image_fast()

    def render_image_fast(self):
        if not self.is_loaded: return
        
        mode = self.display_mode.get() if hasattr(self, 'display_mode') else 'RGB'

        GLOBAL_START_DATE = datetime.fromisoformat(self.reader.dataset_meta['global_start_date'])
        GLOBAL_END_DATE = datetime.fromisoformat(self.reader.dataset_meta['global_end_date'])
        t_norm = normalize_time(self.current_dt, GLOBAL_START_DATE, GLOBAL_END_DATE)

        if self.obs_datetimes:
            times = np.array([dt.timestamp() for dt in self.obs_datetimes])
            curr_ts = self.current_dt.timestamp()
            t_idx = (np.abs(times - curr_ts)).argmin()
        else:
            t_idx = 0
        
        use_mask = self.apply_mask.get() and self.has_mask

        if use_mask and not self.mask_is_static:
            if t_idx != self.last_mask_t_idx:
                self._update_mask_indices(t_idx)
                self.precompute_coords()
                self.last_mask_t_idx = t_idx

        render_H, render_W = self.H, self.W

        current_spatial_coords = None
        current_valid_indices = None
        is_cropped_render = False

        if self.cropped_region:
            is_cropped_render = True
            left, top, right, bottom = map(int, self.cropped_region)
            render_W = right - left
            render_H = bottom - top

            rows = torch.arange(top, bottom, device=self.device)
            cols = torch.arange(left, right, device=self.device)
            g_row, g_col = torch.meshgrid(rows, cols, indexing='ij')

            flat_crop_rows = g_row.flatten()
            flat_crop_cols = g_col.flatten()

            local_valid_mask_gpu = None

            if use_mask and getattr(self, 'external_mask_bool', None) is not None:
                crop_mask_np = self.external_mask_bool[top:bottom, left:right]
                local_valid_mask_gpu = torch.from_numpy(crop_mask_np).to(self.device).reshape(-1)

            elif use_mask and self.reader.mask_vol is not None:
                mask_frame = self.reader.get_mask(t_idx)
                if mask_frame is not None:
                    crop_mask = mask_frame[top:bottom, left:right]
                    crop_mask_gpu = torch.from_numpy(crop_mask).to(self.device).reshape(-1)
                    local_valid_mask_gpu = ~crop_mask_gpu

            if local_valid_mask_gpu is not None:
                flat_crop_rows = flat_crop_rows[local_valid_mask_gpu]
                flat_crop_cols = flat_crop_cols[local_valid_mask_gpu]
                current_valid_indices = torch.nonzero(local_valid_mask_gpu).squeeze()
            
            current_spatial_coords = self.reader._coords_from_rowcol(
                flat_crop_rows, flat_crop_cols, offset_x=0.5, offset_y=0.5
            )
            
        else:
            if self.cached_spatial_coords is None: return
            current_spatial_coords = self.cached_spatial_coords
            current_valid_indices = self.valid_indices

        num_pixels = current_spatial_coords.shape[0]
        
        if num_pixels == 0:
            pil_img = Image.new('RGB', (render_W, render_H), (0, 0, 0))
            
        else:
            if mode == 'RGB':
                C = self.C
                result_gpu = torch.empty((num_pixels, C), device=self.device, dtype=torch.float32)
                chunk = 2**21 
                t_val = torch.tensor(t_norm, device=self.device)
                
                with torch.no_grad():
                    with torch.amp.autocast('cuda', enabled=False):
                        for i in range(0, num_pixels, chunk):
                            end = min(i + chunk, num_pixels)
                            s_batch = current_spatial_coords[i:end]
                            t_batch = t_val.view(1, 1).expand(s_batch.shape[0], 1)
                            result_gpu[i:end], _ = self.reader.model(torch.cat([s_batch, t_batch], dim=-1))

                dmin, dmax = self._get_data_range(to_tensor=True)
                result_gpu = result_gpu * (dmax - dmin) + dmin
                if self.meta.get('use_log', False):
                    torch.expm1(result_gpu, out=result_gpu)

                if self.reader.has_residuals:
                    res_frame = self.reader.get_residual(t_idx) 
                    if res_frame is not None:
                        if is_cropped_render:
                            left, top, right, bottom = map(int, self.cropped_region)
                            res_crop = res_frame[top:bottom, left:right].reshape(-1, self.C)
                            res_gpu = torch.from_numpy(res_crop).to(self.device)
                            
                            if current_valid_indices is not None:
                                result_gpu += res_gpu[current_valid_indices]
                            else:
                                result_gpu += res_gpu
                        else:
                            res_flat = torch.from_numpy(res_frame.reshape(-1, self.C)).to(self.device)
                            if self.valid_indices is not None:
                                result_gpu += res_flat[self.valid_indices]
                            else:
                                result_gpu += res_flat

                full_img_gpu = torch.zeros((render_H * render_W, C), device=self.device, dtype=torch.float32)
                
                if current_valid_indices is not None:
                    full_img_gpu[current_valid_indices] = result_gpu
                else:
                    full_img_gpu = result_gpu 

                full_img_gpu = full_img_gpu.view(render_H, render_W, C)
                
                if C >= 3:
                    bands = self.bands_idx if self.bands_idx else [0, 3, 2]
                    rgb_gpu = full_img_gpu[:, :, bands]
                    
                    if self.global_vmax is None:
                        flat_sample = rgb_gpu[::10, ::10, :].reshape(-1)
                        if flat_sample.numel() > 0:
                            k = int(flat_sample.numel() * 0.98)
                            val, _ = torch.kthvalue(flat_sample, k)
                            self.global_vmax = val.item()
                            print(f"Auto Vmax: {self.global_vmax:.4f}")
                    
                    vmax_val = self.global_vmax if self.global_vmax else 1.0
                    scale = 1.5 / (vmax_val + 1e-6)
                    rgb_gpu = rgb_gpu.mul_(scale).clamp_(0, 1).mul_(255)
                    final_tensor = rgb_gpu.to(torch.uint8)
                    img_np = final_tensor.cpu().numpy()
                    pil_img = Image.fromarray(img_np, mode='RGB')
                else:
                    gray = full_img_gpu[:, :, 0]
                    vmin, vmax = gray.min(), gray.max()
                    gray = (gray - vmin) / (vmax - vmin + 1e-6)
                    final_tensor = (gray * 255).to(torch.uint8)
                    img_np = final_tensor.cpu().numpy()
                    pil_img = Image.fromarray(img_np, mode='L')

            elif mode in ['NDVI', 'NDVI_dt', 'NDVI_spatial']:
                result_map_flat = torch.zeros(num_pixels, device=self.device, dtype=torch.float32)
                chunk_size = 2**20 
                
                req_grad = (mode == 'NDVI_dt')
                req_spatial_grad = (mode == 'NDVI_spatial')
                
                t_val = torch.tensor([[t_norm]], device=self.device, dtype=torch.float32, requires_grad=req_grad)
                
                self.reader.model.eval()
                context = torch.enable_grad() if (req_grad or req_spatial_grad) else torch.no_grad()
                
                with context:
                    for i in range(0, num_pixels, chunk_size):
                        end = min(i + chunk_size, num_pixels)
                        curr_chunk = end - i
                        
                        s_batch = current_spatial_coords[i:end]
                        if req_spatial_grad:
                            s_batch = s_batch.clone().detach().requires_grad_(True)
                        
                        t_batch = t_val.expand(curr_chunk, 1)
                        coords = torch.cat([s_batch, t_batch], dim=-1)
                        
                        preds, _ = self.reader.model(coords)
                        
                        red = preds[:, 0]
                        nir = preds[:, 1]
                        ndvi = (nir - red) / (nir + red + 1e-8)
                        
                        chunk_res = None
                        if mode == 'NDVI':
                            chunk_res = ndvi
                        elif mode == 'NDVI_dt':
                            grad_t = torch.autograd.grad(ndvi, t_batch, torch.ones_like(ndvi), create_graph=False)[0]
                            chunk_res = grad_t.view(-1)
                        elif mode == 'NDVI_spatial':
                            grad_s = torch.autograd.grad(ndvi, s_batch, torch.ones_like(ndvi), create_graph=False)[0]
                            chunk_res = torch.norm(grad_s, dim=1)
                        
                        result_map_flat[i:end] = chunk_res.detach()
                        del coords, preds, ndvi
                
                full_map = np.full(render_H * render_W, np.nan, dtype=np.float32)
                res_cpu = result_map_flat.cpu().numpy()
                
                if current_valid_indices is not None:
                    idx_cpu = current_valid_indices.cpu().numpy()
                    full_map[idx_cpu] = res_cpu
                else:
                    full_map[:] = res_cpu
                    
                full_map = full_map.reshape(render_H, render_W)
                
                show_cbar = self.show_colorbar.get() if isinstance(self.show_colorbar, tk.BooleanVar) else True
                
                if mode == 'NDVI':
                    if show_cbar:
                        pil_img = self._render_with_colorbar(full_map, cmap_name='RdYlGn', vmin=-1.0, vmax=1.0)
                    else:
                        pil_img = self._render_without_colorbar(full_map, cmap_name='RdYlGn', vmin=-1.0, vmax=1.0)
                else:
                    valid_vals = full_map[~np.isnan(full_map)]
                    clim = np.percentile(np.abs(valid_vals), 98) if valid_vals.size > 0 else 1.0
                    if clim == 0: clim = 1e-5
                    if show_cbar:
                        pil_img = self._render_with_colorbar(full_map, cmap_name='coolwarm', vmin=-clim, vmax=clim)
                    else:
                        pil_img = self._render_without_colorbar(full_map, cmap_name='coolwarm', vmin=-clim, vmax=clim)

            else:
                raise ValueError(f"Unknown mode {mode}")

        self.cached_pil_img = pil_img
        if is_cropped_render:
            self.display_cropped_image(pil_img)
        else:
            self.cached_pil_img = pil_img
            if self.cropped_region:
                self.display_cropped_image(pil_img.crop(self.cropped_region))
            else:
                self.display_cropped_image(pil_img)

    def toggle_crop_mode(self):
        self.crop_mode = not self.crop_mode
        if self.crop_mode:
            self.crop_btn.config(relief=tk.SUNKEN)
            self.crop_rect = None
        else:
            self.crop_btn.config(relief=tk.RAISED)
            if self.crop_rect_id:
                self.main_canvas.delete(self.crop_rect_id)
                self.crop_rect_id = None

    def on_canvas_press(self, event):
        if self.crop_mode:
            self.crop_start = (event.x, event.y)
        else:
            self.on_click_canvas(event)

    def on_canvas_drag(self, event):
        if not self.crop_mode or not self.crop_start:
            return
        
        if self.crop_rect_id:
            self.main_canvas.delete(self.crop_rect_id)
        
        x0, y0 = self.crop_start
        x1, y1 = event.x, event.y
        self.crop_rect_id = self.main_canvas.create_rectangle(
            x0, y0, x1, y1, outline='yellow', width=2
        )

    def on_canvas_release(self, event):
        if not self.crop_mode or not self.crop_start:
            return
        
        x0, y0 = self.crop_start
        x1, y1 = event.x, event.y
        
        left = min(x0, x1)
        top = min(y0, y1)
        right = max(x0, x1)
        bottom = max(y0, y1)
        
        if right - left < 10 or bottom - top < 10:
            return 
        
        if self.cached_pil_img is None: return

        view_left = (left - self.img_offset_x) / self.img_scale
        view_top = (top - self.img_offset_y) / self.img_scale
        view_right = (right - self.img_offset_x) / self.img_scale
        view_bottom = (bottom - self.img_offset_y) / self.img_scale
        
        if self.cropped_region:
            prev_global_left, prev_global_top, _, _ = self.cropped_region
            
            global_left = prev_global_left + view_left
            global_top = prev_global_top + view_top
            global_right = prev_global_left + view_right
            global_bottom = prev_global_top + view_bottom
        else:
            global_left = view_left
            global_top = view_top
            global_right = view_right
            global_bottom = view_bottom

        final_left = max(0, int(global_left))
        final_top = max(0, int(global_top))
        final_right = min(self.W, int(global_right))
        final_bottom = min(self.H, int(global_bottom))
        
        if final_right <= final_left or final_bottom <= final_top:
            print("Invalid crop region.")
            return
        
        self.cropped_region = (final_left, final_top, final_right, final_bottom)
        
        self.render_image_fast()
        
        self.crop_mode = False
        self.crop_btn.config(relief=tk.RAISED)
        if self.crop_rect_id:
            self.main_canvas.delete(self.crop_rect_id)
            self.crop_rect_id = None

    def display_cropped_image(self, pil_img):
        canvas_w = self.main_canvas.winfo_width()
        canvas_h = self.main_canvas.winfo_height()
        if canvas_w < 10: canvas_w = 800
        if canvas_h < 10: canvas_h = 600
        
        img_w, img_h = pil_img.size
        
        scale_w = canvas_w / img_w
        scale_h = canvas_h / img_h
        img_scale = min(scale_w, scale_h)
        
        new_w = int(img_w * img_scale)
        new_h = int(img_h * img_scale)
        
        img_offset_x = (canvas_w - new_w) // 2
        img_offset_y = (canvas_h - new_h) // 2
        
        pil_img_resized = pil_img.resize((new_w, new_h), Image.Resampling.BILINEAR)
        self.tk_image_ref = ImageTk.PhotoImage(image=pil_img_resized)
        
        if self.canvas_image_id is None:
            self.canvas_image_id = self.main_canvas.create_image(
                img_offset_x, img_offset_y, 
                image=self.tk_image_ref, anchor=tk.NW
            )
        else:
            self.main_canvas.itemconfig(self.canvas_image_id, image=self.tk_image_ref)
            self.main_canvas.coords(self.canvas_image_id, img_offset_x, img_offset_y)
        
        self.img_scale = img_scale
        self.img_offset_x = img_offset_x
        self.img_offset_y = img_offset_y

    def reset_to_full_image(self):
        if self.cropped_region is None:
            return

        self.cropped_region = None
        self.crop_mode = False
        self.crop_btn.config(relief=tk.RAISED)
        
        if self.crop_rect_id:
            self.main_canvas.delete(self.crop_rect_id)
            self.crop_rect_id = None
        
        self.render_image_fast()

    def on_click_canvas(self, event):
        if not self.is_loaded: return
        
        screen_x = self.main_canvas.canvasx(event.x)
        screen_y = self.main_canvas.canvasy(event.y)
        
        img_x = (screen_x - self.img_offset_x) / self.img_scale
        img_y = (screen_y - self.img_offset_y) / self.img_scale
        
        if self.cropped_region:
            img_left, img_top, img_right, img_bottom = self.cropped_region
            col = int(img_left + img_x)
            row = int(img_top + img_y)
        else:
            col = int(img_x)
            row = int(img_y)

        if 0 <= col < self.W and 0 <= row < self.H:
            print(f"Clicked Pixel: {row}, {col}")
            self.popup_trend_window(row, col)

    def popup_trend_window(self, row, col):
        try:
            val_str = self.res_var.get()
            step_days = float(val_str) if val_str else -1.0
        except ValueError:
            step_days = -1.0

        show_interpolation = (step_days > 0)

        transform = self.meta['transform']
        lon = transform[2] + (col + 0.5) * transform[0]
        lat = transform[5] + (row + 0.5) * transform[4]
        
        meta = self.reader.dataset_meta
        norm_mode = meta.get('norm_mode', 2) 
        device = self.device
        spatial_vec = None
        
        if norm_mode == 3:
            H, W = meta['shape'][1], meta['shape'][2]
            u = (col + 0.5) / W 
            v = (row + 0.5) / H 
            spatial_vec = torch.tensor([u, v], device=device, dtype=torch.float32)
        
        elif norm_mode == 2:
            params = meta.get('local_planar_params')
            if params:
                x = (lon - params['lon_min']) / params['lon_span']
                y = (lat - params['lat_min']) / params['lat_span']
                x = max(0.0, min(1.0, x))
                y = max(0.0, min(1.0, y))
                spatial_vec = torch.tensor([x, y], device=device, dtype=torch.float32)
            else:
                spatial_vec = torch.tensor([0.5, 0.5], device=device, dtype=torch.float32)

        else:
            lon_r, lat_r = np.deg2rad(lon), np.deg2rad(lat)
            cl = np.cos(lat_r)
            raw_x = cl * np.cos(lon_r)
            raw_y = cl * np.sin(lon_r)
            raw_z = np.sin(lat_r)
            
            if norm_mode == 1:
                params = meta.get('local_xyz_params')
                if params:
                    x = (raw_x - params['x_min']) / params['x_span']
                    y = (raw_y - params['y_min']) / params['y_span']
                    z = (raw_z - params['z_min']) / params['z_span']
                else:
                    x, y, z = 0.5, 0.5, 0.5
            else:
                x = (raw_x + 1.0) * 0.5
                y = (raw_y + 1.0) * 0.5
                z = (raw_z + 1.0) * 0.5
            
            x = max(0.0, min(1.0, x))
            y = max(0.0, min(1.0, y))
            z = max(0.0, min(1.0, z))
            spatial_vec = torch.tensor([x, y, z], device=device, dtype=torch.float32)

        GLOBAL_START_DATE = datetime.fromisoformat(meta['global_start_date'])
        GLOBAL_END_DATE = datetime.fromisoformat(meta['global_end_date'])

        obs_dates = self.obs_datetimes
        obs_tnorms = [normalize_time(dt, GLOBAL_START_DATE, GLOBAL_END_DATE) for dt in obs_dates]
        
        vals_obs = self._query_pixel_series(spatial_vec, obs_tnorms)

        # 使用高效的像素时序 mask 查询（仅在 Apply Mask 启用时）
        if self.apply_mask.get() and self.reader.mask_vol is not None:
            vals_obs = vals_obs.astype(np.float32)
            mask_series = self.reader.get_pixel_mask_series(row, col)
            vals_obs[mask_series, :] = np.nan

        vals_dense = None
        dense_dates = []
        
        if show_interpolation:
            plot_start_dt = min(obs_dates)
            plot_end_dt = max(obs_dates)
            
            curr_dt = plot_start_dt
            while curr_dt <= plot_end_dt:
                dense_dates.append(curr_dt)
                curr_dt += timedelta(days=step_days)

            if not dense_dates or dense_dates[-1] != plot_end_dt:
                dense_dates.append(plot_end_dt)

            if dense_dates:
                dense_tnorms = [normalize_time(dt, GLOBAL_START_DATE, GLOBAL_END_DATE) for dt in dense_dates]
                vals_dense = self._query_pixel_series(spatial_vec, dense_tnorms)
                
                method = self.smooth_method.get()
                if method == 'Simple':
                    all_dates = []
                    for dt in dense_dates:
                        all_dates.extend([dt - timedelta(days=1), dt, dt + timedelta(days=1)])
                    all_tnorms = [normalize_time(dt, GLOBAL_START_DATE, GLOBAL_END_DATE) for dt in all_dates]
                    vals_all = self._query_pixel_series(spatial_vec, all_tnorms)
                    
                    vals_dense_smooth = []
                    for i in range(0, len(vals_all), 3):
                        val_minus = vals_all[i]
                        val_center = vals_all[i + 1]
                        val_plus = vals_all[i + 2]
                        smoothed = 0.2 * val_minus + 0.6 * val_center + 0.2 * val_plus
                        vals_dense_smooth.append(smoothed)
                    vals_dense = np.array(vals_dense_smooth)
                elif method == 'Savitzky-Golay':
                    if savgol_filter and len(vals_dense) >= 5:
                        vals_dense = savgol_filter(vals_dense, window_length=5, polyorder=2, axis=0)

        top = tk.Toplevel(self.root)
        mode_str = f"Interpolation (step={step_days}d)" if show_interpolation else "Raw Observation"
        top.title(f"{mode_str} @ ({row}, {col})")
        
        fig_pop = plt.Figure(figsize=(9, 5), dpi=100)
        ax_pop = fig_pop.add_subplot(111)
        
        selected_band = self.ts_band_var.get()
        
        if selected_band == 'ALL':
            bands_to_plot = min(self.C, 10)
            band_indices = list(range(bands_to_plot))
            labels = [f"B{b+1}" for b in band_indices]
        elif selected_band.startswith('Band '):
            # Handle "Band 1", "Band 2", etc.
            try:
                band_idx = int(selected_band.split()[1]) - 1
                if 0 <= band_idx < self.C:
                    bands_to_plot = 1
                    band_indices = [band_idx]
                    labels = [selected_band]
                else:
                    bands_to_plot = 0
                    band_indices = []
            except (ValueError, IndexError):
                bands_to_plot = 0
                band_indices = []
        elif selected_band == 'NDVI':
            bands_to_plot = 1
            band_indices = ['NDVI']
            labels = ['NDVI']
        else:
            bands_to_plot = 0
            band_indices = []
        
        base_colors = ['r', 'g', 'b', 'c', 'm', 'y', 'k']
        all_plotted_values = []
        
        for i, b in enumerate(band_indices):
            if b == 'NDVI':
                red_obs = vals_obs[:, 0]
                nir_obs = vals_obs[:, 1]
                ndvi_obs = (nir_obs - red_obs) / (nir_obs + red_obs + 1e-8)
                all_plotted_values.append(ndvi_obs)
                
                if show_interpolation and vals_dense is not None:
                    red_dense = vals_dense[:, 0]
                    nir_dense = vals_dense[:, 1]
                    ndvi_dense = (nir_dense - red_dense) / (nir_dense + red_dense + 1e-8)
                
                color = 'purple'
                label_prefix = 'NDVI'
            else:
                ndvi_obs = vals_obs[:, b]
                all_plotted_values.append(ndvi_obs)
                if show_interpolation and vals_dense is not None:
                    ndvi_dense = vals_dense[:, b]
                color = base_colors[i % len(base_colors)]
                label_prefix = labels[i]
            
            if show_interpolation:
                if vals_dense is not None:
                    ax_pop.plot(dense_dates, ndvi_dense if b == 'NDVI' else vals_dense[:, b], 
                                color=color, linestyle='-', linewidth=1.5, alpha=0.6,
                                label=f"{label_prefix} Fit")
                
                ax_pop.plot(obs_dates, ndvi_obs, 
                            color=color, linestyle='None', marker='o', markersize=6, alpha=1.0,
                            markeredgecolor='white', markeredgewidth=0.5,
                            label=f"{label_prefix} Obs")
            else:
                ax_pop.plot(obs_dates, ndvi_obs, 
                            color=color, linestyle='-', marker='o', markersize=4, alpha=0.9,
                            linewidth=1.2,
                            label=f"{label_prefix} Obs")

        ax_pop.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
        ax_pop.xaxis.set_major_locator(mdates.AutoDateLocator())
        fig_pop.autofmt_xdate()
        
        ax_pop.set_ylabel("Reflectance / Value", fontsize=14)
        ax_pop.set_title(f"Time Series: Pixel ({row}, {col})", fontsize=16, fontweight='bold')
        ax_pop.tick_params(axis='both', which='major', labelsize=12)

        final_min, final_max = 0.0, 1.0
        
        if all_plotted_values:
            flat_vals = np.concatenate([v.flatten() for v in all_plotted_values])
            valid_vals = flat_vals[~np.isnan(flat_vals)]
            
            if len(valid_vals) > 0:
                y_min, y_max = np.min(valid_vals), np.max(valid_vals)
                span = y_max - y_min
                
                MIN_SPAN = 0.1
                if span < MIN_SPAN:
                    center = (y_max + y_min) / 2
                    y_min = center - MIN_SPAN / 2
                    y_max = center + MIN_SPAN / 2
                    span = MIN_SPAN
                
                padding = span * 0.3
                final_min = y_min - padding
                final_max = y_max + padding
                
                if selected_band == 'NDVI':
                    final_max = min(1.0, final_max)
                    final_min = max(-1.0, final_min)
                else:
                    final_min = max(0.0, final_min)
        
        ax_pop.set_ylim(final_min, final_max)
        ax_pop.grid(True, linestyle='--', alpha=0.5)

        if self.show_current_time_on_ts.get() and self.current_dt:
            ax_pop.axvline(
                x=self.current_dt, 
                color='black', 
                linestyle='--', 
                linewidth=1.5, 
                alpha=0.8,
                label='Current Date'
            )
        
        ax_pop.legend(loc='best', ncol=min(len(band_indices)+1, 3), fontsize='large')

        canvas_pop = FigureCanvasTkAgg(fig_pop, master=top)
        canvas_pop.draw()
        canvas_pop.get_tk_widget().pack(fill=tk.BOTH, expand=True)

    def _query_pixel_series(self, spatial_vec, t_list):
        N = len(t_list)
        
        spatial_dim = spatial_vec.shape[-1] 
        s_batch = spatial_vec.unsqueeze(0).expand(N, spatial_dim)
        
        t_batch = torch.tensor(t_list, device=self.device, dtype=torch.float32).unsqueeze(1)
        
        coords = torch.cat([s_batch, t_batch], dim=-1)
        
        with torch.no_grad():
            with torch.amp.autocast('cuda', enabled=False):
                preds, _ = self.reader.model(coords.float())
        
        vals = preds.cpu().numpy()
        dmin, dmax = self._get_data_range(to_tensor=False)
        use_log = self.meta.get('use_log', False)
        
        vals = vals * (dmax - dmin) + dmin
        if use_log: np.expm1(vals, out=vals)
        return vals

    def _render_without_colorbar(self, arr, cmap_name='viridis', vmin=None, vmax=None):
        H, W = arr.shape

        if vmin is None:
            try:
                vmin = np.nanmin(arr)
            except Exception:
                vmin = 0.0
        if vmax is None:
            try:
                vmax = np.nanmax(arr)
            except Exception:
                vmax = 1.0

        pil_rgb = None
        try:
            import torch
            if torch.cuda.is_available() and getattr(self, 'device', '').startswith('cuda'):
                pil_rgb = self._map_array_to_rgb_gpu(arr, cmap_name, vmin, vmax)
        except Exception:
            pil_rgb = None

        if pil_rgb is None:
            cmap = mpl.cm.get_cmap(cmap_name)
            norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax, clip=True)
            with np.errstate(invalid='ignore'):
                mapped = cmap(norm(arr))
            rgb = (mapped[..., :3] * 255.0).astype(np.uint8)
            nan_mask = np.isnan(arr)
            if nan_mask.any():
                rgb[nan_mask, :] = 0
            pil_rgb = Image.fromarray(rgb, mode='RGB')

        return pil_rgb

    def _render_with_colorbar(self, arr, cmap_name='viridis', vmin=None, vmax=None):
        H, W = arr.shape

        if vmin is None:
            try:
                vmin = np.nanmin(arr)
            except Exception:
                vmin = 0.0
        if vmax is None:
            try:
                vmax = np.nanmax(arr)
            except Exception:
                vmax = 1.0

        pil_rgb = None
        try:
            import torch
            if torch.cuda.is_available() and getattr(self, 'device', '').startswith('cuda'):
                pil_rgb = self._map_array_to_rgb_gpu(arr, cmap_name, vmin, vmax)
        except Exception:
            pil_rgb = None

        if pil_rgb is None:
            cmap = mpl.cm.get_cmap(cmap_name)
            norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax, clip=True)
            with np.errstate(invalid='ignore'):
                mapped = cmap(norm(arr))
            rgb = (mapped[..., :3] * 255.0).astype(np.uint8)
            nan_mask = np.isnan(arr)
            if nan_mask.any():
                rgb[nan_mask, :] = 0
            pil_rgb = Image.fromarray(rgb, mode='RGB')

        if not getattr(self, 'show_colorbar', True):
            return pil_rgb

        cbar_img = self._get_cached_colorbar(cmap_name, vmin, vmax, H)

        pad = max(4, int(W * 0.01))
        cb_w, cb_h = cbar_img.size
        out_w = W + pad + cb_w
        out_h = max(H, cb_h)

        out = Image.new('RGB', (out_w, out_h), (0, 0, 0))
        out.paste(pil_rgb, (0, 0))
        out.paste(cbar_img, (W + pad, (out_h - cb_h) // 2))
        return out

    def _get_cached_colorbar(self, cmap_name, vmin, vmax, height):
        h_bucket = max(64, int(round(height / 32) * 32))
        key = (cmap_name, float(vmin), float(vmax), h_bucket)
        if key in self._cbar_cache:
            return self._cbar_cache[key]

        cmap = mpl.cm.get_cmap(cmap_name)
        norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax, clip=True)

        grad = np.linspace(1.0, 0.0, h_bucket, dtype=np.float32)
        grad = grad[:, None]
        rgba = cmap(grad)
        rgb = (rgba[..., :3] * 255.0).astype(np.uint8)

        cb_width = max(12, int(h_bucket * 0.03))
        rgb = np.repeat(rgb, cb_width, axis=1)

        text_width = 120
        cbar_img = Image.new('RGB', (cb_width + text_width, h_bucket), (0, 0, 0))
        cbar_img.paste(Image.fromarray(rgb, mode='RGB'), (0, 0))

        try:
            draw = ImageDraw.Draw(cbar_img)
            try:
                font = ImageFont.truetype("arial.ttf", 42)
            except:
                font = ImageFont.load_default()
            
            max_str = f"{vmax:.3g}"
            min_str = f"{vmin:.3g}"
            
            draw.text((cb_width + 5, 5), max_str, fill=(255, 255, 255), font=font)
            draw.text((cb_width + 5, h_bucket - 50), min_str, fill=(255, 255, 255), font=font)
        except Exception as e:
            print(f"Warning: Could not add colorbar text labels: {e}")

        self._cbar_cache[key] = cbar_img
        return cbar_img

    def _map_array_to_rgb_gpu(self, arr, cmap_name, vmin, vmax):
        import torch
        import matplotlib as mpl

        device = self.device if hasattr(self, 'device') else ('cuda' if torch.cuda.is_available() else 'cpu')

        H, W = arr.shape
        t = torch.from_numpy(arr.astype(np.float32)).to(device)

        nan_mask = torch.isnan(t)
        denom = (vmax - vmin) if (vmax - vmin) != 0 else 1.0
        t_norm = (t - vmin) / (denom + 1e-12)
        t_norm = t_norm.clamp(0.0, 1.0)

        L = 256
        cmap = mpl.cm.get_cmap(cmap_name, L)
        lut = (np.asarray(cmap(np.linspace(0.0, 1.0, L)))[:, :3] * 255.0).astype(np.uint8)
        lut_t = torch.from_numpy(lut).to(device)

        idx = (t_norm * (L - 1)).round().to(torch.long)
        idx[nan_mask] = 0

        idx_flat = idx.view(-1)
        rgb_flat = lut_t[idx_flat]
        rgb = rgb_flat.view(H, W, 3).to(torch.uint8)

        if nan_mask.any():
            rgb[nan_mask, :] = 0

        arr_rgb = rgb.cpu().numpy()
        pil = Image.fromarray(arr_rgb, mode='RGB')
        return pil

    def _get_data_range(self, to_tensor=True):
        raw_range = self.meta['data_range']
        
        if len(raw_range) > 0 and isinstance(raw_range[0], (list, tuple)):
            arr = np.array(raw_range, dtype=np.float32)
            dmin_np = arr[:, 0]
            dmax_np = arr[:, 1]
        else:
            dmin_np, dmax_np = float(raw_range[0]), float(raw_range[1])
        
        if to_tensor:
            if isinstance(dmin_np, np.ndarray):
                dmin = torch.tensor(dmin_np, device=self.device, dtype=torch.float32)
                dmax = torch.tensor(dmax_np, device=self.device, dtype=torch.float32)
            else:
                dmin = dmin_np
                dmax = dmax_np
            return dmin, dmax
        else:
            return dmin_np, dmax_np

    def save_current_view_as_tif(self):
        if not self.is_loaded or self.cached_pil_img is None:
            messagebox.showwarning("Warning", "No image rendered yet.")
            return

        default_name = "output.tif"
        if self.current_dt:
            date_str = self.current_dt.strftime("%Y%m%d")
            mode = self.display_mode.get()
            default_name = f"GNDC_{mode}_{date_str}.tif"
            
        out_path = filedialog.asksaveasfilename(
            defaultextension=".tif",
            initialfile=default_name,
            filetypes=[("GeoTIFF", "*.tif;*.tiff")]
        )
        if not out_path: return

        try:
            from rasterio.transform import Affine

            img_pil = self.cached_pil_img
            img_np = np.array(img_pil)

            if img_np.ndim == 3:
                count = 3
                data_to_write = img_np.transpose(2, 0, 1)
            else:
                count = 1
                data_to_write = img_np[np.newaxis, :, :]

            height, width = img_np.shape[:2]

            meta = self.reader.dataset_meta
            base_transform_list = meta['transform']
            src_transform = Affine(*base_transform_list)

            current_transform = src_transform
            
            if self.cropped_region:
                left, top, right, bottom = map(int, self.cropped_region)
                crop_w = right - left
                crop_h = bottom - top
                
                if width == crop_w and height == crop_h:
                    current_transform = src_transform * Affine.translation(left, top)

            profile = {
                'driver': 'GTiff',
                'dtype': 'uint8',
                'nodata': None,
                'width': width,
                'height': height,
                'count': count,
                'crs': meta['crs'],
                'transform': current_transform,
                'compress': 'lzw',
                'bigtiff': 'yes'
            }

            with rasterio.open(out_path, 'w', **profile) as dst:
                dst.write(data_to_write)
            
            print(f"Saved view to {out_path}")
            
        except Exception as e:
            import traceback
            traceback.print_exc()
            messagebox.showerror("Error", f"Failed to save TIF:\n{e}")

    def _generate_raw_data_for_t(self, t_idx):
        from rasterio.transform import Affine

        meta = self.reader.dataset_meta
        GLOBAL_START_DATE = datetime.fromisoformat(meta['global_start_date'])
        GLOBAL_END_DATE = datetime.fromisoformat(meta['global_end_date'])
        
        target_dt = self.obs_datetimes[t_idx]
        t_norm = normalize_time(target_dt, GLOBAL_START_DATE, GLOBAL_END_DATE)

        render_H, render_W = self.H, self.W
        current_spatial_coords = None
        current_valid_indices = None
        
        src_transform = Affine(*meta['transform'])
        current_transform = src_transform
        
        is_cropped = False
        if self.cropped_region:
            is_cropped = True
            left, top, right, bottom = map(int, self.cropped_region)
            render_W = right - left
            render_H = bottom - top
            
            current_transform = src_transform * Affine.translation(left, top)
            
            rows = torch.arange(top, bottom, device=self.device)
            cols = torch.arange(left, right, device=self.device)
            g_row, g_col = torch.meshgrid(rows, cols, indexing='ij')
            flat_rows, flat_cols = g_row.flatten(), g_col.flatten()
            
            local_valid_mask = None
            if self.has_mask:
                if getattr(self, 'external_mask_bool', None) is not None:
                    crop_mask = self.external_mask_bool[top:bottom, left:right]
                    local_valid_mask = torch.from_numpy(crop_mask).to(self.device).reshape(-1)
                elif self.reader.mask_vol is not None:
                    mask_frame = self.reader.get_mask(t_idx)
                    if mask_frame is not None:
                        crop_mask = mask_frame[top:bottom, left:right]
                        local_valid_mask = ~torch.from_numpy(crop_mask).to(self.device).reshape(-1)

            if local_valid_mask is not None:
                flat_rows = flat_rows[local_valid_mask]
                flat_cols = flat_cols[local_valid_mask]
                current_valid_indices = torch.nonzero(local_valid_mask).squeeze()
            
            current_spatial_coords = self.reader._coords_from_rowcol(flat_rows, flat_cols, offset_x=0.5, offset_y=0.5)

        else:
            if self.cached_spatial_coords is None: self.precompute_coords()
            current_spatial_coords = self.cached_spatial_coords
            current_valid_indices = self.valid_indices
        
        num_pixels = current_spatial_coords.shape[0]
        C = self.C
        
        result_gpu = torch.empty((num_pixels, C), device=self.device, dtype=torch.float32)
        chunk = 2**20
        t_val = torch.tensor(t_norm, device=self.device)
        
        self.reader.model.eval()
        with torch.no_grad():
            with torch.amp.autocast('cuda', enabled=False):
                for i in range(0, num_pixels, chunk):
                    end = min(i + chunk, num_pixels)
                    s_batch = current_spatial_coords[i:end]
                    t_batch = t_val.view(1, 1).expand(s_batch.shape[0], 1)
                    result_gpu[i:end], _ = self.reader.model(torch.cat([s_batch, t_batch], dim=-1).float())

        dmin, dmax = self._get_data_range(to_tensor=True)
        result_gpu = result_gpu * (dmax - dmin) + dmin
        if self.meta.get('use_log', False):
            torch.expm1(result_gpu, out=result_gpu)

        if self.reader.has_residuals:
            res_frame = self.reader.get_residual(t_idx)
            if res_frame is not None:
                if is_cropped:
                    left, top, right, bottom = map(int, self.cropped_region)
                    res_crop = res_frame[top:bottom, left:right].reshape(-1, C)
                    res_gpu = torch.from_numpy(res_crop).to(self.device)
                    if current_valid_indices is not None:
                        result_gpu += res_gpu[current_valid_indices]
                    else:
                        result_gpu += res_gpu
                else:
                    res_flat = torch.from_numpy(res_frame.reshape(-1, C)).to(self.device)
                    if self.valid_indices is not None:
                        result_gpu += res_flat[self.valid_indices]
                    else:
                        result_gpu += res_flat

        full_img_gpu = torch.full((render_H * render_W, C), float('nan'), device=self.device, dtype=torch.float32)
        
        if current_valid_indices is not None:
            full_img_gpu[current_valid_indices] = result_gpu
        else:
            full_img_gpu = result_gpu

        img_np = full_img_gpu.view(render_H, render_W, C).cpu().numpy()
        
        return img_np, current_transform

    def save_current_frame_raw(self):
        if not self.is_loaded: return
        
        times = np.array([dt.timestamp() for dt in self.obs_datetimes])
        curr_ts = self.current_dt.timestamp()
        t_idx = (np.abs(times - curr_ts)).argmin()
        date_str = self.obs_datetimes[t_idx].strftime("%Y%m%d")
        
        mode_str = "Crop" if self.cropped_region else "Full"
        default_name = f"Raw_{date_str}_{mode_str}.tif"
        out_path = filedialog.asksaveasfilename(
            defaultextension=".tif",
            initialfile=default_name,
            filetypes=[("GeoTIFF", "*.tif")]
        )
        if not out_path: return
        
        try:
            print("Generating Raw Data...")
            img_np, transform = self._generate_raw_data_for_t(t_idx)
            H, W, C = img_np.shape
            
            profile = {
                'driver': 'GTiff',
                'dtype': 'float32',
                'nodata': float('nan'),
                'width': W,
                'height': H,
                'count': C,
                'crs': self.meta['crs'],
                'transform': transform,
                'compress': 'lzw',
                'predictor': 2
            }
            
            with rasterio.open(out_path, 'w', **profile) as dst:
                for b in range(C):
                    dst.write(img_np[:, :, b], indexes=b+1)
                    dst.set_band_description(b+1, f"Band_{b+1}")
            
            print(f"Saved Raw Frame to {out_path}")
            messagebox.showinfo("Success", f"Saved:\n{out_path}")
            
        except Exception as e:
            import traceback
            traceback.print_exc()
            messagebox.showerror("Export Error", str(e))

    def save_time_series_raw(self):
        if not self.is_loaded: return
        
        out_dir = filedialog.askdirectory(title="Select Output Directory for Time Series")
        if not out_dir: return
        
        prefix = "GNDC_Raw"
        
        try:
            count = len(self.obs_datetimes)
            print(f"Starting Batch Export of {count} frames...")
            
            mode_str = "Crop" if self.cropped_region else "Full"
            
            for t_idx, dt in enumerate(self.obs_datetimes):
                date_str = dt.strftime("%Y%m%d")
                fname = f"{prefix}_{date_str}_{mode_str}.tif"
                out_path = os.path.join(out_dir, fname)
                
                img_np, transform = self._generate_raw_data_for_t(t_idx)
                H, W, C = img_np.shape
                
                profile = {
                    'driver': 'GTiff',
                    'dtype': 'float32',
                    'nodata': float('nan'),
                    'width': W,
                    'height': H,
                    'count': C,
                    'crs': self.meta['crs'],
                    'transform': transform,
                    'compress': 'lzw',
                    'predictor': 2
                }
                
                with rasterio.open(out_path, 'w', **profile) as dst:
                    for b in range(C):
                        dst.write(img_np[:, :, b], indexes=b+1)
                
                self.root.update()
                
            print("Batch Export Done.")
            messagebox.showinfo("Success", f"Exported {count} files to:\n{out_dir}")
            
        except Exception as e:
            import traceback
            traceback.print_exc()
            messagebox.showerror("Export Error", str(e))

    def export_gif_animation(self):
        """Export current view as GIF animation."""
        if not self.is_loaded:
            messagebox.showwarning("Warning", "No model loaded.")
            return

        if not self.obs_datetimes:
            messagebox.showwarning("Warning", "No time series data available.")
            return

        # Create dialog for GIF settings
        dialog = tk.Toplevel(self.root)
        dialog.title("Export GIF Animation")
        dialog.transient(self.root)
        dialog.grab_set()

        FONT = ("Arial", 10)
        LABEL_FONT = ("Arial", 10, "bold")

        content_frame = tk.Frame(dialog)
        content_frame.pack(fill=tk.BOTH, expand=True, padx=15, pady=10)

        # Date range selection
        tk.Label(content_frame, text="Start Date:", font=LABEL_FONT).pack(anchor=tk.W, pady=(5, 3))
        date_strings = [dt.strftime("%Y-%m-%d") for dt in self.obs_datetimes]

        start_var = tk.StringVar(value=date_strings[0] if date_strings else "")
        start_combo = ttk.Combobox(content_frame, textvariable=start_var, values=date_strings, state='readonly', width=20)
        start_combo.pack(anchor=tk.W, pady=(0, 10))

        tk.Label(content_frame, text="End Date:", font=LABEL_FONT).pack(anchor=tk.W, pady=(5, 3))
        end_var = tk.StringVar(value=date_strings[-1] if date_strings else "")
        end_combo = ttk.Combobox(content_frame, textvariable=end_var, values=date_strings, state='readonly', width=20)
        end_combo.pack(anchor=tk.W, pady=(0, 10))

        # Frame interval
        tk.Label(content_frame, text="Frame Interval (ms):", font=LABEL_FONT).pack(anchor=tk.W, pady=(5, 3))
        interval_var = tk.StringVar(value="500")
        interval_entry = ttk.Entry(content_frame, textvariable=interval_var, width=22)
        interval_entry.pack(anchor=tk.W, pady=(0, 10))

        # Resolution scale
        tk.Label(content_frame, text="Resolution Scale (0.1-1.0):", font=LABEL_FONT).pack(anchor=tk.W, pady=(5, 3))
        scale_var = tk.StringVar(value="0.5")
        scale_entry = ttk.Entry(content_frame, textvariable=scale_var, width=22)
        scale_entry.pack(anchor=tk.W, pady=(0, 10))

        # Current settings info
        mode = self.display_mode.get() if hasattr(self, 'display_mode') else 'RGB'
        info_text = f"Current Mode: {mode}"
        if mode == 'RGB' and hasattr(self, 'bands_idx'):
            info_text += f"\nRGB Bands: {self.bands_idx}"
        if self.cropped_region:
            info_text += "\nCrop Region: Active"
        tk.Label(content_frame, text=info_text, font=FONT, fg="gray").pack(anchor=tk.W, pady=(10, 5))

        # Buttons
        btn_frame = tk.Frame(content_frame)
        btn_frame.pack(side=tk.BOTTOM, pady=(15, 0))

        def start_export():
            start_date_str = start_var.get()
            end_date_str = end_var.get()
            try:
                interval_ms = int(interval_var.get())
                if interval_ms <= 0:
                    raise ValueError("Interval must be positive")
            except ValueError:
                messagebox.showerror("Error", "Invalid interval value. Must be a positive integer.")
                return

            try:
                scale = float(scale_var.get())
                if scale <= 0 or scale > 1:
                    raise ValueError("Scale must be between 0.1 and 1.0")
            except ValueError:
                messagebox.showerror("Error", "Invalid scale value. Must be between 0.1 and 1.0.")
                return

            dialog.destroy()

            # Find date indices
            try:
                start_idx = date_strings.index(start_date_str)
                end_idx = date_strings.index(end_date_str)
            except ValueError:
                messagebox.showerror("Error", "Invalid date selection.")
                return

            if start_idx > end_idx:
                messagebox.showerror("Error", "Start date must be before end date.")
                return

            # Select output file
            default_name = f"GNDC_{mode}_{start_date_str}_{end_date_str}.gif"
            out_path = filedialog.asksaveasfilename(
                defaultextension=".gif",
                initialfile=default_name,
                filetypes=[("GIF", "*.gif")]
            )
            if not out_path:
                return

            try:
                # Get current settings
                original_dt = self.current_dt
                original_display_mode = self.display_mode.get() if hasattr(self, 'display_mode') else 'RGB'
                original_bands_idx = self.bands_idx if hasattr(self, 'bands_idx') else [0, 3, 2]

                # Prepare frames
                frames = []
                selected_dates = self.obs_datetimes[start_idx:end_idx+1]

                print(f"Generating GIF for {len(selected_dates)} frames...")

                for i, target_dt in enumerate(selected_dates):
                    self.current_dt = target_dt
                    self.render_image_fast()

                    if self.cached_pil_img is None:
                        print(f"Skipping frame {i+1}: no image rendered")
                        continue

                    # Convert to RGB and scale
                    img = self.cached_pil_img.convert('RGB')
                    if scale != 1.0:
                        new_size = (int(img.width * scale), int(img.height * scale))
                        img = img.resize(new_size, Image.LANCZOS)

                    frames.append(img)

                    # Update progress
                    self.root.update()
                    if i % 5 == 0:
                        print(f"Progress: {i+1}/{len(selected_dates)}")

                # Restore original state
                self.current_dt = original_dt
                self.display_mode.set(original_display_mode)
                self.bands_idx = original_bands_idx
                self.render_image_fast()

                if not frames:
                    messagebox.showwarning("Warning", "No frames generated.")
                    return

                # Save GIF
                print(f"Saving GIF with {len(frames)} frames...")
                frames[0].save(
                    out_path,
                    save_all=True,
                    append_images=frames[1:],
                    duration=interval_ms,
                    loop=0
                )

                print(f"GIF saved to {out_path}")
                messagebox.showinfo("Success", f"GIF saved to:\n{out_path}")

            except Exception as e:
                import traceback
                traceback.print_exc()
                messagebox.showerror("Export Error", str(e))

        ttk.Button(btn_frame, text="Export", command=start_export).pack(side=tk.LEFT, padx=5)
        ttk.Button(btn_frame, text="Cancel", command=dialog.destroy).pack(side=tk.LEFT, padx=5)

    def run(self):
        self.root.mainloop()


    def show_display_settings_dialog(self):
        """Show display settings dialog."""
        dialog = tk.Toplevel(self.root)
        dialog.title("Display Settings")
        dialog.transient(self.root)
        dialog.grab_set()

        FONT = ("Arial", 10)
        LABEL_FONT = ("Arial", 10, "bold")

        # Content frame - vertical layout
        content_frame = tk.Frame(dialog)
        content_frame.pack(fill=tk.BOTH, expand=True, padx=15, pady=10)

        # Display mode selection
        tk.Label(content_frame, text="Display Mode:", font=LABEL_FONT).pack(anchor=tk.W, pady=(5, 3))

        modes = [
            ("RGB", "RGB"),
            ("NDVI", "NDVI"),
            ("NDVI/dt", "NDVI_dt"),
            ("NDVI Spatial Grad", "NDVI_spatial")
        ]
        for label, value in modes:
            ttk.Radiobutton(content_frame, text=label, variable=self.display_mode, value=value).pack(anchor=tk.W, padx=15)

        # RGB Band selection
        tk.Label(content_frame, text="RGB Bands:", font=LABEL_FONT).pack(anchor=tk.W, pady=(15, 3))

        band_values = [f"Band {i+1}" for i in range(self.C)] if self.C > 0 else ["Band 1"]

        row_frame = tk.Frame(content_frame)
        row_frame.pack(fill=tk.X, padx=10)

        ttk.Label(row_frame, text="R:").grid(row=0, column=0, sticky=tk.W, padx=5, pady=2)
        self.rgb_r_var = tk.StringVar(value=band_values[min(0, len(band_values)-1)])
        r_combo = ttk.Combobox(row_frame, textvariable=self.rgb_r_var, values=band_values, state='readonly', width=12)
        r_combo.grid(row=0, column=1, padx=5, pady=2)

        ttk.Label(row_frame, text="G:").grid(row=0, column=2, sticky=tk.W, padx=5, pady=2)
        self.rgb_g_var = tk.StringVar(value=band_values[min(2, len(band_values)-1)] if self.C >= 3 else band_values[0])
        g_combo = ttk.Combobox(row_frame, textvariable=self.rgb_g_var, values=band_values, state='readonly', width=12)
        g_combo.grid(row=0, column=3, padx=5, pady=2)

        ttk.Label(row_frame, text="B:").grid(row=0, column=4, sticky=tk.W, padx=5, pady=2)
        self.rgb_b_var = tk.StringVar(value=band_values[min(1, len(band_values)-1)] if self.C >= 3 else band_values[0])
        b_combo = ttk.Combobox(row_frame, textvariable=self.rgb_b_var, values=band_values, state='readonly', width=12)
        b_combo.grid(row=0, column=5, padx=5, pady=2)

        # Buttons
        btn_frame = tk.Frame(dialog)
        btn_frame.pack(side=tk.BOTTOM, fill=tk.X, padx=15, pady=10)
        ttk.Button(btn_frame, text="Apply", command=lambda: [self.apply_display_settings(r_combo.current(), g_combo.current(), b_combo.current()), dialog.destroy()]).pack(side=tk.LEFT, padx=5)
        ttk.Button(btn_frame, text="Cancel", command=dialog.destroy).pack(side=tk.LEFT, padx=5)

    def apply_display_settings(self, r_idx, g_idx, b_idx):
        """Apply display settings and re-render."""
        if self.display_mode.get() == 'RGB':
            self.bands_idx = [r_idx, g_idx, b_idx] if self.C >= 3 else [r_idx]
        self.render_image_fast()

    def show_timeseries_settings_dialog(self):
        """Show time series settings dialog."""
        dialog = tk.Toplevel(self.root)
        dialog.title("TimeSeries Settings")
        dialog.transient(self.root)
        dialog.grab_set()

        # Resolution
        ttk.Label(dialog, text="Time Series Resolution:").pack(anchor=tk.W, padx=10, pady=(10, 5))

        res_frame = tk.Frame(dialog)
        res_frame.pack(fill=tk.X, padx=10, pady=5)
        ttk.Label(res_frame, text="Res (Days):").pack(side=tk.LEFT)
        ttk.Entry(res_frame, textvariable=self.res_var, width=10).pack(side=tk.LEFT, padx=5)
        ttk.Label(res_frame, text="(0 = Original observation dates)").pack(side=tk.LEFT)

        # Smooth method
        smooth_frame = tk.Frame(dialog)
        smooth_frame.pack(fill=tk.X, padx=10, pady=5)
        tk.Label(smooth_frame, text="Smooth:").pack(side=tk.LEFT)
        smooth_combo = ttk.Combobox(smooth_frame, textvariable=self.smooth_method, values=['None', 'Simple', 'Savitzky-Golay'], state='readonly', width=15)
        smooth_combo.pack(side=tk.LEFT, padx=5)

        # TS Band
        ts_band_frame = tk.Frame(dialog)
        ts_band_frame.pack(fill=tk.X, padx=10, pady=5)
        tk.Label(ts_band_frame, text="TS Band:").pack(side=tk.LEFT)
        self.ts_band_var = tk.StringVar(value='ALL')
        band_values = ['ALL'] + [f"Band {i+1}" for i in range(self.C)] if self.C > 0 else ['ALL']
        ts_band_combo = ttk.Combobox(ts_band_frame, textvariable=self.ts_band_var, values=band_values, state='readonly', width=15)
        ts_band_combo.pack(side=tk.LEFT, padx=5)

        # Info
        info_frame = tk.Frame(dialog, relief=tk.SUNKEN, bd=1)
        info_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        info_text = (
            "TimeSeries Settings:\n\n"
            "• Res (Days): Time interval for time series reconstruction.\n"
            "  0 = use original observation dates only.\n\n"
            "• Smooth: Smoothing method for time series.\n"
            "  - None: No smoothing\n"
            "  - Simple: Moving average\n"
            "  - Savitzky-Golay: Polynomial smoothing\n\n"
            "• TS Band: Band to display in time series chart."
        )
        tk.Label(info_frame, text=info_text, justify=tk.LEFT, font=("Arial", 9)).pack(anchor=tk.W, padx=5, pady=5)

        # Buttons
        btn_frame = tk.Frame(dialog)
        btn_frame.pack(side=tk.BOTTOM, pady=10)
        ttk.Button(btn_frame, text="Apply", command=dialog.destroy).pack(side=tk.LEFT, padx=5)
        ttk.Button(btn_frame, text="Cancel", command=dialog.destroy).pack(side=tk.LEFT, padx=5)


def start_viewer():
    """Start the interactive viewer."""
    try:
        from ctypes import windll
        windll.shcore.SetProcessDpiAwareness(1)
    except: pass
    
    root = tk.Tk()
    root.geometry("1800x1200")
    app = GNDCInteractiveViewer(root)
    root.mainloop()


if __name__ == "__main__":
    start_viewer()
