"""
GNDC Reader - Model loading and reconstruction.

This module provides the GNDCReader class for loading compressed
models and reconstructing data from them.
"""

import torch
import struct
import numpy as np
import json
import os
import rasterio
from tqdm import tqdm
from datetime import datetime

from pygndc.utils.compression import fast_decompress
from pygndc.utils.date_utils import normalize_time


class GNDCReader:
    """
    Reader for GNDC compressed model files.
    
    Handles loading models, decoding masks/residuals, and
    reconstructing data at arbitrary time points.
    """
    
    # Valid modes: "auto", "tcnn_cuda", "torch_gpu", "torch_cpu"
    _MODE_MAP = {
        "tcnn_cuda": ("tcnn", "cuda"),
        "torch_gpu": ("torch", "cuda"),
        "torch_cpu": ("torch", "cpu"),
    }

    def __init__(self, mode="auto"):
        import torch

        if mode == "auto":
            has_cuda = torch.cuda.is_available()
            if has_cuda:
                try:
                    import tinycudann
                    mode = "tcnn_cuda"
                except ImportError:
                    mode = "torch_gpu"
            else:
                mode = "torch_cpu"

        if mode not in self._MODE_MAP:
            raise ValueError(
                f"Invalid mode '{mode}'. Choose from: auto, tcnn_cuda, torch_gpu, torch_cpu"
            )

        # Validate environment for requested mode
        if mode == "tcnn_cuda":
            try:
                import tinycudann
            except ImportError:
                raise ImportError(
                    "mode='tcnn_cuda' requires tinycudann. Install it with:\n"
                    "  pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch\n"
                    "Or use mode='torch_gpu' / mode='torch_cpu'."
                )
            if not torch.cuda.is_available():
                raise RuntimeError(
                    "mode='tcnn_cuda' requires CUDA, but torch.cuda.is_available() is False. "
                    "Install CUDA-enabled PyTorch or use mode='torch_cpu'."
                )
        elif mode == "torch_gpu" and not torch.cuda.is_available():
            raise RuntimeError(
                "mode='torch_gpu' requires CUDA, but torch.cuda.is_available() is False. "
                "Use mode='torch_cpu' instead."
            )

        self.backend, self.device = self._MODE_MAP[mode]
        self.mode = mode

        if mode == "torch_gpu":
            print("Note: using torch_gpu mode. Install tinycudann for faster inference (tcnn_cuda).")
        elif mode == "torch_cpu":
            print("Note: using torch_cpu mode. For better performance, use a CUDA-enabled GPU.")

        self.model = None
        self.model_config = None
        self.dataset_meta = None

        self.mask_vol = None
        self.mask_is_static = True
        self.residuals = None
        self.has_residuals = False

        # Lazy decode attributes (v0.4: flat packed bits)
        self._mask_packed = None      # 存储 packed bits 的原始字节
        self._mask_shape = None       # shape (T, H, W)
        self._mask_type = None        # mask 类型标志 (0/1)
        self._bytes_per_frame = None  # 每帧字节数

        # v0.5: per-frame chunked mask
        self._mask_version = "0.4"    # file version
        self._mask_raw_bytes = None   # raw mask section bytes (for chunked access)
        self._mask_frame_index = None # List[(offset, comp_len)] for each frame

    def load(self, ndc_path):
        """
        Load a .gndc model file.
        
        Args:
            ndc_path: Path to the .gndc file
        """
        if not os.path.exists(ndc_path):
            raise FileNotFoundError(f"Model file not found: {ndc_path}")

        with open(ndc_path, "rb") as f:
            header_len = struct.unpack("I", f.read(4))[0]
            header = json.loads(f.read(header_len))
            
            self.model_config = header['model_config']
            self.dataset_meta = header.get('dataset_meta', None)

            meta = self.dataset_meta
            if 't_norms' not in meta or not meta['t_norms']:
                try:
                    gs_str = str(meta['global_start_date'])
                    ge_str = str(meta['global_end_date'])
                    
                    def parse_dt(s):
                        try: return datetime.strptime(s, "%Y-%m-%d %H:%M:%S")
                        except: return datetime.fromisoformat(s)

                    g_start = parse_dt(gs_str)
                    g_end = parse_dt(ge_str)
                    
                    timestamps = meta['timestamps']
                    meta['t_norms'] = []
                    for ts_str in timestamps:
                        dt = parse_dt(ts_str)
                        t_val = normalize_time(dt, g_start, g_end)
                        meta['t_norms'].append(t_val)
                        
                    
                except Exception as e:
                    print(f"Warning: Failed to reconstruct t_norms dynamically: {e}")
            
            self.has_residuals = header.get('has_residuals', False)
            self.residual_info = header.get('residual_info', None)
            self._mask_version = header.get('version', '0.4')

            w_len = struct.unpack("I", f.read(4))[0]
            compressed_bytes = f.read(w_len)
            weights_raw = fast_decompress(compressed_bytes)

            q_info = header.get("quantization_info", {"mode": "float16"})
            self._load_weights_from_raw(weights_raw, q_info)

            try:
                chunk_len_bytes = f.read(4)
                if len(chunk_len_bytes) == 4:
                    mask_len = struct.unpack("I", chunk_len_bytes)[0]
                    mask_bytes = f.read(mask_len)
                    self._decode_mask(mask_bytes)
                else:
                    print("Warning: No Mask block found (Legacy file?).")
            except Exception as e:
                print(f"Mask Load Error: {e}")

            if self.has_residuals:
                try:
                    res_len_bytes = f.read(4)
                    if len(res_len_bytes) == 4:
                        res_len = struct.unpack("I", res_len_bytes)[0]
                        res_bytes = f.read(res_len)
                        self._decode_residuals(res_bytes)
                    else:
                        print("Warning: Header expects residuals but EOF reached.")
                except Exception as e:
                    print(f"Residual Load Error: {e}")

        self._init_model()

    def _load_weights_from_raw(self, weights_raw, q_info):
        """Load and dequantize weights."""
        q_mode = q_info.get("mode", "float16")
        if q_mode == "int16":
            scale = q_info['scale']
            w_int = np.frombuffer(weights_raw, dtype=np.int16)
            self.w_tensor = torch.from_numpy(w_int.astype(np.float32) / scale).to(self.device)
        elif q_mode == "int8":
            scale = q_info['scale']
            w_int = np.frombuffer(weights_raw, dtype=np.int8)
            self.w_tensor = torch.from_numpy(w_int.astype(np.float32) / scale).to(self.device)
        else:
            self.w_tensor = torch.from_numpy(np.frombuffer(weights_raw, dtype=np.float16).astype(np.float32)).to(self.device)

    def _init_model(self):
        """Initialize GNDCNetwork model."""
        cfg = self.model_config
        n_bands = cfg.get('n_bands', 1)
        spatial_dims = cfg.get('spatial_dims', 3)
        if self.backend == "torch":
            from pygndc.models.GNDCNetworkTorch import GNDCNetworkTorch as GNDCNetwork
        else:
            from pygndc.models.GNDCNetwork import GNDCNetwork

        self.model = GNDCNetwork(
            n_output_dims=n_bands,
            xy_hash_size=cfg.get('xy_hash_size', 19), 
            xy_base_resolution=cfg.get('xy_base_resolution', 512),
            xy_n_levels=cfg.get('xy_n_levels', 9),
            xy_per_level_scale=cfg.get('xy_per_level_scale', 1.43),
            xy_features_per_level=cfg.get('xy_features_per_level', 4),
            t_hash_size=cfg.get('t_hash_size', 21),
            t_base_resolution=cfg.get('t_base_resolution', 16),
            t_n_levels=cfg.get('t_n_levels', 9),
            t_per_level_scale=cfg.get('t_per_level_scale', 1.43),
            t_features_per_level=cfg.get('t_features_per_level', 4),
            spatial_scale=cfg.get('spatial_scale', 1.0),
            decoder_hidden=cfg.get('decoder_hidden', 128),
            n_hidden_layers=cfg.get('n_hidden_layers', 6),
            spatial_dims=spatial_dims,
        ).to(self.device)

        self.model.load_weights_from_flat_buffer(self.w_tensor)
        del self.w_tensor

    def reset(self):
        """Reset reader and free GPU memory."""
        if self.model is not None:
            del self.model
            self.model = None

        self.mask_vol = None
        self.mask_is_static = True
        self.residuals = None
        self.dataset_meta = None
        self.model_config = None

        # Lazy decode cleanup
        self._mask_packed = None
        self._mask_shape = None
        self._mask_type = None
        self._bytes_per_frame = None

        # v0.5 chunked mask cleanup
        self._mask_version = "0.4"
        self._mask_raw_bytes = None
        self._mask_frame_index = None

        if hasattr(self, '_spatial_coords_cache'):
            self._spatial_coords_cache = None
        if hasattr(self, '_base_lonlat_cache'):
            self._base_lonlat_cache = None
        if hasattr(self, '_cached_spatial_coords'):
            self._cached_spatial_coords = None

        import gc
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        

    def _decode_mask(self, mask_bytes):
        """Decode mask data block. Supports v0.4 (flat packed) and v0.5 (per-frame chunked)."""
        if not mask_bytes: return

        if self._mask_version >= "0.5":
            self._decode_mask_v05(mask_bytes)
        else:
            self._decode_mask_v04(mask_bytes)

    def _decode_mask_v04(self, mask_bytes):
        """Decode v0.4 mask: single zstd blob containing all frames packed together."""
        raw = fast_decompress(mask_bytes)

        self._mask_type = raw[0]
        self._mask_shape = np.frombuffer(raw[1:13], dtype=np.int32)

        self._mask_packed = raw[13:]

        T, H, W = int(self._mask_shape[0]), int(self._mask_shape[1]), int(self._mask_shape[2])
        self._bytes_per_frame = (H * W + 7) // 8
        self._mask_shape = (T, H, W)

        self.mask_vol = True

        if self._mask_type == 0:
            self.mask_is_static = True
            print(f"  -> Mask decoded (v0.4): {self._mask_shape} (Static)")
        elif self._mask_type == 1:
            self.mask_is_static = False
            print(f"  -> Mask decoded (v0.4): {self._mask_shape} (Dynamic - Raw)")
        else:
            print(f"Warning: Unknown mask type {self._mask_type}, treating as Static.")
            self.mask_is_static = True

    def _decode_mask_v05(self, mask_bytes):
        """Decode v0.5 mask: per-frame chunked zstd format (no outer compression)."""
        raw = mask_bytes  # v0.5 mask section is NOT wrapped in an outer zstd layer

        self._mask_type = raw[0]
        self._mask_shape_arr = np.frombuffer(raw[1:13], dtype=np.int32)
        T, H, W = int(self._mask_shape_arr[0]), int(self._mask_shape_arr[1]), int(self._mask_shape_arr[2])
        self._mask_shape = (T, H, W)
        self._bytes_per_frame = (H * W + 7) // 8

        # Build frame index: scan [4B len][data] sequence
        self._mask_raw_bytes = raw
        self._mask_frame_index = []
        offset = 13
        n_frames = 1 if self._mask_type == 0 else T
        for _ in range(n_frames):
            if offset + 4 > len(raw):
                break
            comp_len = struct.unpack_from('<I', raw, offset)[0]
            offset += 4
            self._mask_frame_index.append((offset, comp_len))
            offset += comp_len

        self.mask_vol = True

        if self._mask_type == 0:
            self.mask_is_static = True
            print(f"  -> Mask decoded (v0.5): {self._mask_shape} (Static, 1 chunk)")
        elif self._mask_type == 1:
            self.mask_is_static = False
            print(f"  -> Mask decoded (v0.5): {self._mask_shape} (Dynamic, {len(self._mask_frame_index)} chunks)")
        else:
            print(f"Warning: Unknown mask type {self._mask_type}, treating as Static.")
            self.mask_is_static = True

    def _decode_residuals(self, res_bytes):
        """Decode residual data block."""
        if not res_bytes: return
        print("  -> Decompressing residuals (this may take memory)...")
        raw = fast_decompress(res_bytes)
        
        shape = np.frombuffer(raw[:16], dtype=np.int64) 
        
        if not hasattr(self, 'residual_info') or self.residual_info is None:
             raise RuntimeError("Missing 'residual_info' in GNDC Header. Cannot decode residuals.")

        r_dtype_str = self.residual_info.get('dtype', 'float16')
        scale = self.residual_info.get('scale', 1.0)
        
        if r_dtype_str == 'int8':
            dtype_np = np.int8
        elif r_dtype_str == 'int16':
            dtype_np = np.int16
        else:
            dtype_np = np.float16
            
        print(f"  -> Residual Format: {r_dtype_str} (Scale={scale})")

        res_flat = np.frombuffer(raw[16:], dtype=dtype_np)
        
        if dtype_np != np.float16:
            self.residuals = res_flat.astype(np.float32) * scale
        else:
            self.residuals = res_flat.astype(np.float32)

        meta = self.dataset_meta
        T, H, W, C = meta['shape']
        
        expected_size = T * H * W * C
        if self.residuals.size == expected_size:
            self.residuals = self.residuals.reshape(T, H, W, C)
            print(f"  -> Residuals loaded successfully: {self.residuals.shape}")
        else:
            self.residuals = None
            raise ValueError(f"Residual size mismatch! Expected {expected_size}, got {self.residuals.size}.")

    def get_mask(self, t_idx):
        """Get mask for specified time index (lazy decode, supports v0.4 and v0.5)."""
        if self.mask_vol is None: return None

        # v0.5 chunked path
        if self._mask_frame_index is not None:
            return self._get_mask_v05(t_idx)

        # v0.4 flat packed path
        if self._mask_packed is not None:
            return self._get_mask_v04(t_idx)

        # Legacy fallback (pre-lazy data)
        if self.mask_is_static:
            return self.mask_vol[0]
        return self.mask_vol[t_idx]

    def _get_mask_v04(self, t_idx):
        """Get mask frame from v0.4 flat packed bits."""
        T, H, W = self._mask_shape
        bytes_per_frame = int(self._bytes_per_frame)
        effective_t = 0 if self.mask_is_static else t_idx

        start_byte = effective_t * bytes_per_frame
        end_byte = start_byte + bytes_per_frame

        packed_arr = np.frombuffer(self._mask_packed, dtype=np.uint8)
        frame_packed = packed_arr[start_byte:end_byte]

        total_pixels = H * W
        frame_unpacked = np.unpackbits(frame_packed, count=total_pixels)
        return frame_unpacked.reshape(H, W).astype(bool)

    _zstd_dctx = None  # reusable decompressor

    def _get_zstd_dctx(self):
        if GNDCReader._zstd_dctx is None:
            import zstandard as zstd
            GNDCReader._zstd_dctx = zstd.ZstdDecompressor()
        return GNDCReader._zstd_dctx

    def _ensure_mask_packed_cache(self):
        """Decompress all v0.5 chunked frames into a flat packed buffer (like v0.4).

        This is done once on first pixel-series query. After this,
        ``_mask_packed`` and ``_bytes_per_frame`` are set so that
        ``get_pixel_mask_series`` can use the fast vectorised v0.4 path.
        """
        if self._mask_packed is not None:
            return  # already cached

        T, H, W = self._mask_shape
        bytes_per_frame = (H * W + 7) // 8
        n_frames = len(self._mask_frame_index)

        buf = bytearray(n_frames * bytes_per_frame)
        dctx = self._get_zstd_dctx()
        for i in range(n_frames):
            offset, comp_len = self._mask_frame_index[i]
            compressed = self._mask_raw_bytes[offset:offset + comp_len]
            frame_bytes = dctx.decompress(compressed)
            buf[i * bytes_per_frame:(i + 1) * bytes_per_frame] = frame_bytes

        self._mask_packed = bytes(buf)
        self._bytes_per_frame = bytes_per_frame

    def _get_mask_v05(self, t_idx):
        """Get mask frame from v0.5 per-frame chunked zstd."""
        T, H, W = self._mask_shape
        total_pixels = H * W
        effective_t = 0 if self.mask_is_static else t_idx

        offset, comp_len = self._mask_frame_index[effective_t]
        compressed = self._mask_raw_bytes[offset:offset + comp_len]

        dctx = self._get_zstd_dctx()
        frame_packed = np.frombuffer(dctx.decompress(compressed), dtype=np.uint8)

        frame_unpacked = np.unpackbits(frame_packed, count=total_pixels)
        return frame_unpacked.reshape(H, W).astype(bool)

    def get_pixel_mask_series(self, row, col):
        """
        Get mask time series for a single pixel (highly efficient).

        Args:
            row: Pixel row index
            col: Pixel column index

        Returns:
            np.ndarray, shape=(T,), dtype=bool
            True = nodata (masked), False = valid data
        """
        T = self.dataset_meta['shape'][0] if self.dataset_meta else 1

        if self.mask_vol is None:
            return np.zeros(T, dtype=bool)

        if self._mask_shape is not None:
            T, H, W = self._mask_shape
        else:
            return np.zeros(T, dtype=bool)

        # v0.5 chunked path: lazily inflate all frames into flat packed cache,
        # then use vectorized bit extraction (same as v0.4 fast path).
        # First call pays the decompression cost; subsequent calls are O(1).
        if self._mask_frame_index is not None:
            self._ensure_mask_packed_cache()
            # Fall through to v0.4 flat packed path below

        # v0.4 flat packed path
        if self._mask_packed is None:
            return np.zeros(T, dtype=bool)

        bytes_per_frame = int(self._bytes_per_frame)
        packed_arr = np.frombuffer(self._mask_packed, dtype=np.uint8)

        if self.mask_is_static:
            pixel_idx = row * W + col
            byte_offset = pixel_idx // 8
            bit_pos = 7 - (pixel_idx % 8)
            byte_val = packed_arr[byte_offset]
            mask_val = (byte_val >> bit_pos) & 1
            return np.full(T, bool(mask_val), dtype=bool)

        pixel_idx = row * W + col
        byte_offset = pixel_idx // 8
        bit_pos = 7 - (pixel_idx % 8)

        byte_indices = np.arange(T, dtype=np.int64) * bytes_per_frame + byte_offset
        bytes_vec = packed_arr[byte_indices]

        return ((bytes_vec >> bit_pos) & 1).astype(bool)

    def _lonlat_to_rowcol(self, lon: float, lat: float) -> tuple:
        """
        Convert geographic coordinates (lon, lat) to pixel (row, col).

        Args:
            lon: Longitude
            lat: Latitude

        Returns:
            (row, col) tuple of float pixel coordinates

        Raises:
            RuntimeError: If metadata or transform is missing
        """
        meta = self.dataset_meta
        if meta is None:
            raise RuntimeError("dataset_meta missing. Call load() first.")

        if 'transform' not in meta:
            raise RuntimeError("GeoTransform missing in metadata.")

        transform = meta['transform']
        # transform = [res_lon, 0, min_lon, 0, res_lat(<0), max_lat]
        min_lon = float(transform[2])
        res_lon = float(transform[0])
        max_lat = float(transform[5])
        res_lat = float(transform[4])  # negative

        col = (lon - min_lon) / res_lon - 0.5
        row = (lat - max_lat) / res_lat - 0.5

        return row, col

    def reconstruct_pixel_series(
        self, row: int, col: int, chunk_size: int = 2**16,
        apply_mask: bool = True, add_residuals: bool = True
    ) -> np.ndarray:
        """
        Efficiently reconstruct the full time series for a single pixel.

        Builds a (T, spatial_dims+1) coordinate tensor and runs one batched
        inference instead of T full-frame reconstructions.

        Args:
            row: Pixel row index
            col: Pixel column index
            chunk_size: Processing chunk size
            apply_mask: Whether to apply mask (sets masked values to NaN)
            add_residuals: Whether to add residuals

        Returns:
            np.ndarray of shape (T, C) — the time series at this pixel
        """
        if self.model is None:
            raise RuntimeError("Model not loaded.")

        meta = self.dataset_meta
        T, H, W, C_total = meta['shape']
        t_norms = meta['t_norms']
        data_ranges = meta['data_range']
        is_per_channel = isinstance(data_ranges[0], (list, tuple, np.ndarray))
        use_log = meta.get('use_log', False)
        model_bands = self.model_config.get('n_bands', 1)

        # Build spatial coords for this single pixel (replicated T times)
        row_t = torch.tensor([row], device=self.device)
        col_t = torch.tensor([col], device=self.device)
        s_coord = self._coords_from_rowcol(row_t, col_t, offset_x=0.5, offset_y=0.5)
        # s_coord: (1, spatial_dims), replicate to (T, spatial_dims)
        s_coords = s_coord.expand(T, -1)

        # Build time column (T, 1)
        t_vals = torch.tensor(t_norms, device=self.device, dtype=torch.float32).unsqueeze(1)

        # Full coordinate tensor (T, spatial_dims + 1)
        coords = torch.cat([s_coords, t_vals], dim=-1)

        # Batched inference
        results = []
        self.model.eval()
        with torch.no_grad():
            for i in range(0, T, chunk_size):
                end = min(i + chunk_size, T)
                preds, _ = self.model(coords[i:end])
                results.append(preds.cpu().numpy())

        series = np.concatenate(results, axis=0)  # (T, model_bands)

        # Denormalize
        if is_per_channel:
            dmin = np.array([r[0] for r in data_ranges], dtype=np.float32)
            dmax = np.array([r[1] for r in data_ranges], dtype=np.float32)
        else:
            dmin, dmax = float(data_ranges[0]), float(data_ranges[1])

        series = series * (dmax - dmin) + dmin
        if use_log:
            np.expm1(series, out=series)

        # Add residuals
        if add_residuals and self.has_residuals and self.residuals is not None:
            series += self.residuals[:, row, col, :].astype(np.float32)

        # Apply mask
        if apply_mask and self.mask_vol is not None:
            mask_series = self.get_pixel_mask_series(row, col)
            # mask_series is (T,) bool — True = masked
            if mask_series.shape[0] == T:
                series[mask_series] = np.nan
            elif mask_series.shape[0] == 1:
                # static mask returned single value
                if mask_series[0]:
                    series[:] = np.nan

        return series

    def get_residual(self, t_idx):
        """Get residuals for specified time index."""
        if self.residuals is None: return None
        return self.residuals[t_idx]

    def _coords_from_rowcol(self, rows, cols, offset_x=0.5, offset_y=0.5, require_lonlat_grad=False):
        """Convert row/col indices to spatial coordinates."""
        meta = self.dataset_meta
        if meta is None:
            raise RuntimeError("dataset_meta missing. Call load() first.")

        device = self.device

        rows = rows.to(device=device, dtype=torch.float32)
        cols = cols.to(device=device, dtype=torch.float32)

        if not isinstance(offset_x, torch.Tensor):
            ox = float(offset_x)
            offset_x = torch.full_like(cols, ox, dtype=torch.float32, device=device)
        else:
            offset_x = offset_x.to(device=device, dtype=torch.float32)

        if not isinstance(offset_y, torch.Tensor):
            oy = float(offset_y)
            offset_y = torch.full_like(rows, oy, dtype=torch.float32, device=device)
        else:
            offset_y = offset_y.to(device=device, dtype=torch.float32)
        
        norm_mode = meta['norm_mode']
        if norm_mode == 3:
            H, W = meta['shape'][1], meta['shape'][2]
            
            x = (cols + offset_x) / W
            y = (rows + offset_y) / H
            
            x = torch.clamp(x, 0.0, 1.0)
            y = torch.clamp(y, 0.0, 1.0)
            
            s_coords = torch.stack([x, y], dim=-1)
            
            if require_lonlat_grad:
                x.requires_grad_(True)
                y.requires_grad_(True)
                s_coords = torch.stack([x, y], dim=-1)
                return s_coords, x, y
            
            return s_coords

        if 'transform' not in meta:
             raise RuntimeError("GeoTransform missing in metadata for norm_mode 0/1/2.")
        transform = meta['transform']
        min_lon, res_lon = float(transform[2]), float(transform[0])
        max_lat, res_lat = float(transform[5]), float(transform[4])

        lon_deg = min_lon + (cols + offset_x) * res_lon
        lat_deg = max_lat + (rows + offset_y) * res_lat

        if require_lonlat_grad:
            lon_deg = lon_deg.clone().detach().requires_grad_(True)
            lat_deg = lat_deg.clone().detach().requires_grad_(True)

        if 'norm_mode' in meta:
            norm_mode = meta['norm_mode']
        else:
            norm_mode = 1 if meta.get('local_xyz_params') is not None else 0

        if norm_mode == 2:
            planar_params = meta.get('local_planar_params')
            if planar_params is None:
                 raise RuntimeError("norm_mode=2 but 'local_planar_params' missing in metadata.")
            
            x = (lon_deg - planar_params['lon_min']) / planar_params['lon_span']
            y = (lat_deg - planar_params['lat_min']) / planar_params['lat_span']
            
            x = torch.clamp(x, 0.0, 1.0)
            y = torch.clamp(y, 0.0, 1.0)
            
            s_coords = torch.stack([x, y], dim=-1)
            
        else:
            lon_rad = torch.deg2rad(lon_deg)
            lat_rad = torch.deg2rad(lat_deg)
            
            cos_phi = torch.cos(lat_rad)
            raw_x = cos_phi * torch.cos(lon_rad)
            raw_y = cos_phi * torch.sin(lon_rad)
            raw_z = torch.sin(lat_rad)
            
            local_params = meta.get('local_xyz_params', None)

            if local_params is not None and norm_mode == 1:
                x = (raw_x - local_params['x_min']) / local_params['x_span']
                y = (raw_y - local_params['y_min']) / local_params['y_span']
                z = (raw_z - local_params['z_min']) / local_params['z_span']
                
                x = torch.clamp(x, 0.0, 1.0)
                y = torch.clamp(y, 0.0, 1.0)
                z = torch.clamp(z, 0.0, 1.0)
            else:
                x = (raw_x + 1.0) * 0.5
                y = (raw_y + 1.0) * 0.5
                z = (raw_z + 1.0) * 0.5

            s_coords = torch.stack([x, y, z], dim=-1)

        if require_lonlat_grad:
            return s_coords, lon_deg, lat_deg
        
        return s_coords

    def reconstruct_single_frame(self, t_idx, band_idx=-1, chunk_size=2**18, apply_mask=True, add_residuals=True):
        """
        Reconstruct a single frame.
        
        Args:
            t_idx: Time index
            band_idx: Band index (-1 for all bands)
            chunk_size: Processing chunk size
            apply_mask: Whether to apply mask
            add_residuals: Whether to add residuals
        """
        if self.model is None: raise RuntimeError("Model not loaded.")
        meta = self.dataset_meta
        
        T, H, W, C_total = meta['shape']
        t_val = meta['t_norms'][t_idx]
        data_ranges = meta['data_range']
        is_per_channel = isinstance(data_ranges[0], (list, tuple, np.ndarray))

        use_log = meta.get('use_log', False)
        
        model_bands = self.model_config.get('n_bands', 1)
        target_band_idx = -1
        extract_single_band = False
        
        if band_idx != -1:
            target_band_idx = band_idx - 1 
            extract_single_band = True
            frame = np.zeros((H, W), dtype=np.float32)
            if is_per_channel:
                dmin = float(data_ranges[target_band_idx][0])
                dmax = float(data_ranges[target_band_idx][1])
            else:
                dmin, dmax = data_ranges
        else:
            frame = np.zeros((H, W, model_bands), dtype=np.float32)
            if is_per_channel:
                dmin = np.array([r[0] for r in data_ranges], dtype=np.float32)
                dmax = np.array([r[1] for r in data_ranges], dtype=np.float32)
            else:
                dmin, dmax = data_ranges

        rows = torch.arange(H, device=self.device)
        cols = torch.arange(W, device=self.device)
        g_row, g_col = torch.meshgrid(rows, cols, indexing='ij')
        spatial_xyz = self._coords_from_rowcol(g_row.flatten(), g_col.flatten(), offset_x=0.5, offset_y=0.5)
        
        self.model.eval()
        with torch.no_grad():
            for i in range(0, H*W, chunk_size):
                end = min(i + chunk_size, H*W)
                s_batch = spatial_xyz[i:end]
                t_batch = torch.full((s_batch.shape[0], 1), t_val, device=self.device)
                coords = torch.cat([s_batch, t_batch], dim=-1)
                preds, _ = self.model(coords)
                
                if extract_single_band:
                    frame.reshape(-1)[i:end] = preds[:, target_band_idx].cpu().numpy()
                else:
                    frame.reshape(-1, model_bands)[i:end] = preds.cpu().numpy()
                    
        frame = frame * (dmax - dmin) + dmin
        if use_log: np.expm1(frame, out=frame)
        
        if add_residuals and self.has_residuals:
            res_frame = self.get_residual(t_idx)
            if res_frame is not None:
                if extract_single_band:
                    res_band = res_frame[:, :, target_band_idx].astype(np.float32)
                    frame += res_band
                else:
                    frame += res_frame.astype(np.float32)

        if apply_mask and self.mask_vol is not None:
            mask_frame = self.get_mask(t_idx)
            if mask_frame is not None:
                if extract_single_band:
                    frame[mask_frame] = np.nan
                else:
                    frame[mask_frame] = np.nan

        return frame

    def reconstruct_and_save_all(self, output_dir, src_scale=1.0, out_dtype='float32', 
                               name_prefix='Frame',
                               apply_mask=True, add_residuals=True,
                               recons_dates=None):
        """
        Export all frames as TIF files.
        
        Args:
            output_dir: Output directory
            src_scale: Source scale factor
            out_dtype: Output data type
            name_prefix: Filename prefix
            apply_mask: Whether to apply mask
            add_residuals: Whether to add residuals
            recons_dates: Custom date range for interpolation
        """
        import os
        import rasterio
        from rasterio.transform import Affine
        from datetime import datetime, timedelta
        
        if not os.path.exists(output_dir): os.makedirs(output_dir)
        
        meta = self.dataset_meta
        H, W = meta['shape'][1], meta['shape'][2]
        
        print(f"Exporting Sequence to {output_dir}...")

        export_tasks = []
        
        if recons_dates is not None:
            start_str, end_str, span_days = recons_dates
            
            def parse_date(d):
                if isinstance(d, str):
                    if len(d) == 10: return datetime.strptime(d, "%Y-%m-%d")
                    return datetime.strptime(d, "%Y-%m-%d %H:%M:%S")
                return d

            curr_date = parse_date(start_str)
            end_date = parse_date(end_str)
            
            g_start = parse_date(meta['global_start_date'])
            g_end = parse_date(meta['global_end_date'])
            
            print(f"Interpolation Mode: {curr_date} to {end_date}, step={span_days} days")
            
            while curr_date <= end_date:
                tnorm = normalize_time(curr_date, g_start, g_end)
                
                export_tasks.append({
                    'type': 'interp',
                    'date': curr_date,
                    'tnorm': tnorm,
                    'fname_suffix': curr_date.strftime("%Y%m%d")
                })
                curr_date += timedelta(days=span_days)
                
            if add_residuals:
                print("Warning: 'add_residuals' is disabled in interpolation mode.")
                add_residuals = False
                
        else:
            timestamps = meta.get('timestamps', [])
            T = len(timestamps)
            for t_idx in range(T):
                ts_str = timestamps[t_idx]
                safe_date = ts_str.split(" ")[0].replace("-", "")
                export_tasks.append({
                    'type': 'obs',
                    't_idx': t_idx,
                    'fname_suffix': safe_date
                })

        save_nodata = None
        if 'int' in out_dtype: save_nodata = 0
        elif 'float' in out_dtype: save_nodata = np.nan

        profile = {
            'driver': 'GTiff', 'height': H, 'width': W, 
            'count': self.model_config.get('n_bands', 1),
            'dtype': out_dtype, 'crs': meta['crs'], 
            'transform': Affine(*meta['transform']),
            'compress': 'lzw', 'nodata': save_nodata,
            'bigtiff': 'YES'
        }

        for task in tqdm(export_tasks, desc="Exporting"):
            
            if task['type'] == 'obs':
                frame = self.reconstruct_single_frame(
                    task['t_idx'], band_idx=-1, 
                    apply_mask=apply_mask, add_residuals=add_residuals
                )
            else:
                frame = self._reconstruct_arbitrary_time(
                    task['tnorm'], 
                    apply_static_mask=apply_mask
                )

            frame_raw = frame / (src_scale + 1e-9)
            
            if 'int' in out_dtype:
                frame_raw = np.nan_to_num(frame_raw, nan=save_nodata)
                frame_raw = np.round(frame_raw).astype(out_dtype)
            else:
                frame_raw = frame_raw.astype(out_dtype)
            
            fname = f"{name_prefix}_{task['fname_suffix']}.tif"
            out_path = os.path.join(output_dir, fname)
            
            with rasterio.open(out_path, 'w', **profile) as dst:
                if profile['count'] == 1:
                    dst.write(frame_raw.squeeze(), indexes=1)
                else:
                    for b in range(profile['count']):
                        dst.write(frame_raw[:, :, b], indexes=b+1)

    def _reconstruct_arbitrary_time(self, tnorm, apply_static_mask=True, chunk_size=1024*256):
        """Reconstruct at arbitrary time point."""
        meta = self.dataset_meta
        H, W = meta['shape'][1], meta['shape'][2]
        
        if not hasattr(self, '_cached_spatial_coords'):
            device = self.device
            rows, cols = torch.meshgrid(
                torch.arange(H, device=device), 
                torch.arange(W, device=device), 
                indexing='ij'
            )
            self._cached_spatial_coords = self._coords_from_rowcol(rows.flatten(), cols.flatten())
            
        spatial_flat = self._cached_spatial_coords
        N = spatial_flat.shape[0]
        
        t_vec = torch.full((N, 1), tnorm, device=self.device, dtype=torch.float32)
        
        input_coords = torch.cat([spatial_flat, t_vec], dim=-1)
        
        model_out_list = []
        with torch.no_grad():
            for i in range(0, N, chunk_size):
                batch = input_coords[i : i+chunk_size]
                pred, _  = self.model(batch)
                model_out_list.append(pred)
                
        full_pred = torch.cat(model_out_list, dim=0)
        img_pred = full_pred.view(H, W, -1).cpu().numpy()

        data_ranges = meta['data_range']
        is_per_channel = isinstance(data_ranges[0], (list, tuple, np.ndarray))
        use_log = meta.get('use_log', False)

        if is_per_channel:
            dmin = np.array([r[0] for r in data_ranges], dtype=np.float32)
            dmax = np.array([r[1] for r in data_ranges], dtype=np.float32)
        else:
            dmin, dmax = data_ranges

        img_pred = img_pred * (dmax - dmin) + dmin
        if use_log: np.expm1(img_pred, out=img_pred)
        
        if apply_static_mask and self.mask_vol is not None:
            if self.mask_is_static:
                mask_frame = self.get_mask(0)
                if mask_frame is not None:
                    mask_expanded = np.repeat(mask_frame[:, :, np.newaxis], img_pred.shape[2], axis=2)
                    img_pred[mask_expanded] = np.nan
        
        return img_pred

    def compute_spatial_derivatives(self, target_t_norm, mode="dx", chunk_size=2**16):
        """
        Compute spatial derivatives.
        
        Args:
            target_t_norm: Normalized time point [0, 1]
            mode: Derivative mode ('dx', 'dy', 'mag', 'dir')
            chunk_size: Processing chunk size
            
        Returns:
            result_map: Derivative image
        """
        if self.model is None: raise RuntimeError("Model not loaded.")
        
        meta = self.dataset_meta
        H, W = meta['shape'][1], meta['shape'][2]
        C = self.model_config.get('n_bands', 1) 
        transform = meta['transform']

        result_map = np.zeros((H * W, C), dtype=np.float32)

        rows = torch.arange(H, device=self.device)
        cols = torch.arange(W, device=self.device)
        g_row, g_col = torch.meshgrid(rows, cols, indexing='ij')
        flat_rows = g_row.flatten()
        flat_cols = g_col.flatten()

        min_lon, res_lon = transform[2], transform[0]
        max_lat, res_lat = transform[5], transform[4]
        
        self.model.eval()
        
        print(f"Computing Spatial Derivative (Mode='{mode}') at t={target_t_norm:.4f}...")
        
        with torch.enable_grad():
            for i in range(0, H*W, chunk_size):
                end = min(i + chunk_size, H*W)
                current_chunk_size = end - i
                
                chunk_rows = flat_rows[i:end]
                chunk_cols = flat_cols[i:end]
                
                s_coords, lon_var, lat_var = self._coords_from_rowcol(
                    chunk_rows, chunk_cols, offset_x=0.5, offset_y=0.5, require_lonlat_grad=True
                )
                
                t_chunk = torch.full((current_chunk_size, 1), target_t_norm, device=self.device)
                coords = torch.cat([s_coords, t_chunk], dim=-1)
                
                preds, _ = self.model(coords)
                
                chunk_results = torch.zeros((current_chunk_size, C), device=self.device)
                
                for band in range(C):
                    output_slice = preds[:, band]
                    
                    gradients = torch.autograd.grad(
                        outputs=output_slice,
                        inputs=[lon_var, lat_var],
                        grad_outputs=torch.ones_like(output_slice),
                        create_graph=False,
                        retain_graph=True
                    )
                    
                    d_lon = gradients[0]
                    d_lat = gradients[1]
                    
                    if mode == "dx":
                        chunk_results[:, band] = d_lon
                    elif mode == "dy":
                        chunk_results[:, band] = d_lat
                    elif mode == "mag":
                        mag = torch.sqrt(d_lon**2 + d_lat**2)
                        chunk_results[:, band] = mag
                    elif mode == "dir":
                        direction = torch.atan2(d_lat, d_lon)
                        chunk_results[:, band] = direction
                
                result_map[i:end] = chunk_results.detach().cpu().numpy()
                
                del coords, s_coords, preds, gradients, lon_var, lat_var

        return result_map.reshape(H, W, C)

    def _calc_ndvi_torch(self, preds, red_idx=0, nir_idx=1):
        """Calculate NDVI from predictions."""
        red = preds[:, red_idx]
        nir = preds[:, nir_idx]
        ndvi = (nir - red) / (nir + red + 1e-8)
        return ndvi

    def compute_finite_temporal_derivatives_NDVI(self, target_t_norm, order=1, delta_t=1e-3, 
                                                  red_idx=0, nir_idx=1, chunk_size=2**16):
        """
        Compute NDVI temporal derivatives using finite difference.
        
        Args:
            target_t_norm: Normalized time point
            order: 1 for first derivative, 2 for second
            delta_t: Time step for finite difference
            red_idx: Red band index
            nir_idx: NIR band index
            chunk_size: Processing chunk size
            
        Returns:
            result_map: Derivative image
        """
        if self.model is None: raise RuntimeError("Model not loaded.")
        if order not in [1, 2]: raise ValueError("Order must be 1 or 2")
        
        meta = self.dataset_meta
        H, W = meta['shape'][1], meta['shape'][2]
        transform = meta['transform']

        result_map = np.zeros((H * W), dtype=np.float32)

        rows = torch.arange(H, device=self.device)
        cols = torch.arange(W, device=self.device)
        g_row, g_col = torch.meshgrid(rows, cols, indexing='ij')
        flat_rows = g_row.flatten()
        flat_cols = g_col.flatten()
        
        self.model.eval()
        
        mode_str = "Velocity (1st)" if order == 1 else "Acceleration (2nd)"
        print(f"Computing NDVI {mode_str} using Finite Difference (dt={delta_t})...")
        
        with torch.no_grad():
            for i in range(0, H*W, chunk_size):
                end = min(i + chunk_size, H*W)
                current_chunk_size = end - i
                
                chunk_rows = flat_rows[i:end]
                chunk_cols = flat_cols[i:end]
                
                s_coords = self._coords_from_rowcol(chunk_rows, chunk_cols, offset_x=0.5, offset_y=0.5)

                t_minus = target_t_norm - delta_t
                t_plus = target_t_norm + delta_t
                t_center = target_t_norm
                
                batch_t_minus = torch.full((current_chunk_size, 1), t_minus, device=self.device)
                batch_t_plus = torch.full((current_chunk_size, 1), t_plus, device=self.device)
                
                coords_minus = torch.cat([s_coords, batch_t_minus], dim=-1)
                coords_plus = torch.cat([s_coords, batch_t_plus], dim=-1)
                
                preds_minus, _ = self.model(coords_minus)
                ndvi_minus = self._calc_ndvi_torch(preds_minus, red_idx, nir_idx)
                
                preds_plus, _ = self.model(coords_plus)
                ndvi_plus = self._calc_ndvi_torch(preds_plus, red_idx, nir_idx)
                
                val = None
                
                if order == 1:
                    val = (ndvi_plus - ndvi_minus) / (2 * delta_t)
                    
                elif order == 2:
                    batch_t_center = torch.full((current_chunk_size, 1), t_center, device=self.device)
                    coords_center = torch.cat([s_coords, batch_t_center], dim=-1)
                    preds_center, _ = self.model(coords_center)
                    ndvi_center = self._calc_ndvi_torch(preds_center, red_idx, nir_idx)
                    
                    val = (ndvi_plus - 2 * ndvi_center + ndvi_minus) / (delta_t ** 2)
                
                result_map[i:end] = val.cpu().numpy()

        return result_map.reshape(H, W)
