"""
GNDCDataset — high-level, GDAL/rasterio-style API for .gndc files.

Composes GNDCReader internally, providing familiar geospatial API patterns
with neural compression superpowers (continuous time interpolation,
analytic gradients, extreme compression ratios).
"""

import gc
import os
from collections import namedtuple
from datetime import datetime
from typing import List, Optional, Sequence, Tuple, Union

import numpy as np
import torch

from pygndc.exceptions import CoordinateError, GNDCError
from pygndc.reader.reader import GNDCReader
from pygndc.utils.date_utils import normalize_time

BoundingBox = namedtuple("BoundingBox", ["left", "bottom", "right", "top"])


class GNDCDataset:
    """
    High-level dataset interface for .gndc compressed geospatial files.

    Provides GDAL/rasterio-style metadata properties and data access methods,
    while exposing neural compression capabilities like continuous time
    interpolation and efficient point queries.

    Examples:
        >>> import pygndc
        >>> ds = pygndc.open("satellite.gndc")
        >>> print(ds.crs, ds.bounds, ds.shape)
        >>> frame = ds.read(t=0)                  # (H, W, C)
        >>> frame = ds.read_at_time("2024-06-15")  # continuous interpolation
        >>> ts = ds.sample(lon=116.5, lat=39.9)    # (T, C) point time series
        >>> ds.close()
    """

    def __init__(
        self,
        path: str,
        mode: str = "auto",
    ):
        """
        Open a .gndc file.

        Args:
            path: Path to the .gndc file
            mode: Inference mode ('auto', 'tcnn_cuda', 'torch_gpu', 'torch_cpu')
        """
        self._path = os.path.abspath(path)
        self._reader = GNDCReader(mode=mode)
        self._reader.load(self._path)
        self._closed = False

        # Cache parsed metadata
        self._meta = self._reader.dataset_meta
        self._model_config = self._reader.model_config
        self._timestamps_parsed: Optional[List[datetime]] = None
        self._global_start: Optional[datetime] = None
        self._global_end: Optional[datetime] = None

    # ------------------------------------------------------------------
    # Context manager
    # ------------------------------------------------------------------

    def __enter__(self) -> "GNDCDataset":
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.close()

    def close(self) -> None:
        """Release resources and free GPU memory."""
        if not self._closed:
            self._reader.reset()
            self._closed = True

    def _check_open(self) -> None:
        if self._closed:
            raise GNDCError("Dataset is closed.")

    # ------------------------------------------------------------------
    # Metadata properties
    # ------------------------------------------------------------------

    @property
    def mode(self) -> str:
        """Current inference mode (tcnn_cuda, torch_gpu, torch_cpu)."""
        return self._reader.mode

    @property
    def crs(self):
        """Coordinate reference system (rasterio.crs.CRS)."""
        self._check_open()
        from rasterio.crs import CRS

        crs_val = self._meta.get("crs")
        if crs_val is None:
            return None
        if isinstance(crs_val, CRS):
            return crs_val
        return CRS.from_user_input(crs_val)

    @property
    def transform(self):
        """Affine geo-transform."""
        self._check_open()
        from rasterio.transform import Affine

        t = self._meta.get("transform")
        if t is None:
            return None
        if isinstance(t, Affine):
            return t
        return Affine(*t)

    @property
    def bounds(self) -> BoundingBox:
        """Geographic bounding box (left, bottom, right, top)."""
        self._check_open()
        t = self.transform
        if t is None:
            return BoundingBox(0, 0, self.width, self.height)
        left = t.c
        top = t.f
        right = left + t.a * self.width
        bottom = top + t.e * self.height
        return BoundingBox(left, bottom, right, top)

    @property
    def shape(self) -> Tuple[int, int, int, int]:
        """Data shape as (T, H, W, C)."""
        self._check_open()
        s = self._meta["shape"]
        return (s[0], s[1], s[2], self._model_config.get("n_bands", 1))

    @property
    def height(self) -> int:
        """Spatial height in pixels."""
        return self._meta["shape"][1]

    @property
    def width(self) -> int:
        """Spatial width in pixels."""
        return self._meta["shape"][2]

    @property
    def n_bands(self) -> int:
        """Number of spectral bands."""
        return self._model_config.get("n_bands", 1)

    @property
    def n_timestamps(self) -> int:
        """Number of discrete time steps."""
        return self._meta["shape"][0]

    @property
    def timestamps(self) -> List[datetime]:
        """List of observation timestamps as datetime objects."""
        self._check_open()
        if self._timestamps_parsed is None:
            self._timestamps_parsed = []
            for ts_str in self._meta.get("timestamps", []):
                self._timestamps_parsed.append(self._parse_dt(str(ts_str)))
        return self._timestamps_parsed

    @property
    def resolution(self) -> Tuple[float, float]:
        """Spatial resolution (res_x, res_y) in CRS units."""
        self._check_open()
        t = self.transform
        if t is None:
            return (1.0, 1.0)
        return (abs(t.a), abs(t.e))

    @property
    def nodata(self) -> float:
        """NoData value (always NaN for neural reconstructions)."""
        return float("nan")

    @property
    def band_names(self) -> Optional[List[str]]:
        """Band names (None if not specified in header)."""
        return self._meta.get("band_names")

    @property
    def band_wavelengths(self) -> Optional[List[float]]:
        """Band center wavelengths in nm (None if not specified)."""
        return self._meta.get("band_wavelengths")

    @property
    def has_residuals(self) -> bool:
        """Whether residual correction data is present."""
        return self._reader.has_residuals

    @property
    def has_mask(self) -> bool:
        """Whether a validity mask is present."""
        return self._reader.mask_vol is not None

    @property
    def meta(self) -> dict:
        """Complete metadata dictionary."""
        self._check_open()
        return {
            "path": self._path,
            "crs": self.crs,
            "transform": self.transform,
            "bounds": self.bounds,
            "shape": self.shape,
            "height": self.height,
            "width": self.width,
            "n_bands": self.n_bands,
            "n_timestamps": self.n_timestamps,
            "resolution": self.resolution,
            "nodata": self.nodata,
            "band_names": self.band_names,
            "has_residuals": self.has_residuals,
            "has_mask": self.has_mask,
            "timestamps": self.timestamps,
            "dataset_meta": self._meta,
            "model_config": self._model_config,
        }

    # ------------------------------------------------------------------
    # Data access
    # ------------------------------------------------------------------

    def read(
        self,
        t: Union[int, str, datetime] = 0,
        bands: Optional[List[int]] = None,
        window: Optional[Tuple[int, int, int, int]] = None,
    ) -> np.ndarray:
        """
        Read a single frame.

        Args:
            t: Time index (int), date string ('YYYY-MM-DD'), or datetime.
            bands: List of 0-based band indices to read. None = all bands.
            window: Spatial window (col_off, row_off, width, height).

        Returns:
            np.ndarray of shape (H, W, C) or (h, w, len(bands)).
        """
        self._check_open()
        t_idx, _ = self._resolve_time(t)

        frame = self._reader.reconstruct_single_frame(t_idx, band_idx=-1)
        # frame: (H, W, C) or (H, W) if single band

        if frame.ndim == 2:
            frame = frame[:, :, np.newaxis]

        if window is not None:
            col_off, row_off, w, h = window
            frame = frame[row_off : row_off + h, col_off : col_off + w]

        if bands is not None:
            frame = frame[:, :, bands]

        return frame

    def read_at_time(
        self,
        time: Union[str, datetime],
        bands: Optional[List[int]] = None,
        window: Optional[Tuple[int, int, int, int]] = None,
    ) -> np.ndarray:
        """
        Reconstruct at an arbitrary time point via continuous interpolation.

        This is a neural compression superpower — the model can reconstruct
        at any time, not just observed timestamps.

        Args:
            time: Date string ('YYYY-MM-DD' or 'YYYY-MM-DD HH:MM:SS') or datetime.
            bands: List of 0-based band indices. None = all.
            window: Spatial window (col_off, row_off, width, height).

        Returns:
            np.ndarray of shape (H, W, C) or (h, w, len(bands)).
        """
        self._check_open()
        _, t_norm = self._resolve_time(time)
        frame = self._reader._reconstruct_arbitrary_time(t_norm)

        if frame.ndim == 2:
            frame = frame[:, :, np.newaxis]

        if window is not None:
            col_off, row_off, w, h = window
            frame = frame[row_off : row_off + h, col_off : col_off + w]

        if bands is not None:
            frame = frame[:, :, bands]

        return frame

    def sample(
        self,
        lon: float,
        lat: float,
        t: Optional[Union[int, str, datetime]] = None,
    ) -> np.ndarray:
        """
        Point query by geographic coordinates.

        Args:
            lon: Longitude
            lat: Latitude
            t: Time index/string/datetime. If None, returns full time series.

        Returns:
            np.ndarray of shape (C,) if t is given, or (T, C) for full series.

        Raises:
            CoordinateError: If coordinates are outside the dataset bounds.
        """
        self._check_open()
        row_f, col_f = self._reader._lonlat_to_rowcol(lon, lat)
        row, col = int(round(row_f)), int(round(col_f))

        H, W = self.height, self.width
        if row < 0 or row >= H or col < 0 or col >= W:
            raise CoordinateError(
                f"Coordinates ({lon}, {lat}) map to pixel ({row}, {col}) "
                f"which is outside bounds (0..{H-1}, 0..{W-1})."
            )

        if t is None:
            return self.pixel_series(row, col)

        t_idx, _ = self._resolve_time(t)
        frame = self._reader.reconstruct_single_frame(t_idx, band_idx=-1)
        if frame.ndim == 2:
            frame = frame[:, :, np.newaxis]
        return frame[row, col]  # (C,)

    def pixel_series(
        self,
        row: int,
        col: int,
    ) -> np.ndarray:
        """
        Get the full time series for a single pixel.

        Uses efficient batched inference (one forward pass for all T steps),
        not T full-frame reconstructions.

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

        Returns:
            np.ndarray of shape (T, C).
        """
        self._check_open()
        return self._reader.reconstruct_pixel_series(row, col)

    def read_region(
        self,
        bbox: Tuple[float, float, float, float],
        t: Union[int, str, datetime] = 0,
        bands: Optional[List[int]] = None,
    ) -> np.ndarray:
        """
        Read a spatial sub-region by geographic bounding box.

        Args:
            bbox: (min_lon, min_lat, max_lon, max_lat)
            t: Time index/string/datetime.
            bands: Band indices. None = all.

        Returns:
            np.ndarray of shape (h, w, C).
        """
        self._check_open()
        min_lon, min_lat, max_lon, max_lat = bbox

        # Convert corners to pixel coordinates
        row_top, col_left = self._reader._lonlat_to_rowcol(min_lon, max_lat)
        row_bottom, col_right = self._reader._lonlat_to_rowcol(max_lon, min_lat)

        # Clamp to valid range
        row_off = max(0, int(np.floor(row_top)))
        col_off = max(0, int(np.floor(col_left)))
        row_end = min(self.height, int(np.ceil(row_bottom)) + 1)
        col_end = min(self.width, int(np.ceil(col_right)) + 1)

        h = row_end - row_off
        w = col_end - col_off
        if h <= 0 or w <= 0:
            raise CoordinateError(f"Bounding box {bbox} maps to empty region.")

        return self.read(t=t, bands=bands, window=(col_off, row_off, w, h))

    # ------------------------------------------------------------------
    # String representations
    # ------------------------------------------------------------------

    def __repr__(self) -> str:
        if self._closed:
            return "<GNDCDataset (closed)>"
        T, H, W, C = self.shape
        crs_str = str(self.crs) if self.crs else "unknown"
        name = os.path.basename(self._path)
        return f"<GNDCDataset '{name}' (T={T}, {H}x{W}, {C}bands) {crs_str}>"

    def __str__(self) -> str:
        return self.info()

    def info(self) -> str:
        """Return a formatted summary string."""
        self._check_open()
        T, H, W, C = self.shape
        b = self.bounds
        lines = [
            f"GNDCDataset: {os.path.basename(self._path)}",
            f"  Mode:        {self.mode}",
            f"  Path:        {self._path}",
            f"  Shape:       T={T}, H={H}, W={W}, C={C}",
            f"  CRS:         {self.crs}",
            f"  Bounds:      ({b.left:.6f}, {b.bottom:.6f}, {b.right:.6f}, {b.top:.6f})",
            f"  Resolution:  {self.resolution}",
            f"  Timestamps:  {T} steps",
            f"  Residuals:   {self.has_residuals}",
            f"  Mask:        {self.has_mask}",
        ]
        if self.timestamps:
            lines.append(f"  Time range:  {self.timestamps[0]} -> {self.timestamps[-1]}")
        return "\n".join(lines)

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _parse_dt(s: str) -> datetime:
        """Parse a date/datetime string flexibly."""
        for fmt in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%d", "%Y%m%d"):
            try:
                return datetime.strptime(s, fmt)
            except ValueError:
                continue
        return datetime.fromisoformat(s)

    def _get_global_range(self) -> Tuple[datetime, datetime]:
        """Return (global_start, global_end) as datetimes."""
        if self._global_start is None:
            self._global_start = self._parse_dt(str(self._meta["global_start_date"]))
            self._global_end = self._parse_dt(str(self._meta["global_end_date"]))
        return self._global_start, self._global_end

    def _resolve_time(self, t: Union[int, str, datetime]) -> Tuple[int, float]:
        """
        Resolve a flexible time argument to (t_idx, t_norm).

        Args:
            t: int index, date string, or datetime

        Returns:
            (t_idx, t_norm): The nearest discrete index and the normalized [0,1] value.
        """
        t_norms = self._meta["t_norms"]

        if isinstance(t, int):
            if t < 0 or t >= len(t_norms):
                raise IndexError(f"Time index {t} out of range [0, {len(t_norms) - 1}].")
            return t, float(t_norms[t])

        if isinstance(t, str):
            t = self._parse_dt(t)

        # t is datetime — find nearest t_idx
        g_start, g_end = self._get_global_range()
        t_norm = normalize_time(t, g_start, g_end)
        t_idx = int(np.argmin([abs(tn - t_norm) for tn in t_norms]))
        return t_idx, t_norm

    # ------------------------------------------------------------------
    # Phase 2 convenience wrappers (delegate to interop module)
    # ------------------------------------------------------------------

    def to_rasterio(self, t: Union[int, str, datetime] = 0):
        """
        Export a frame as a rasterio MemoryFile.

        Args:
            t: Time index/string/datetime.

        Returns:
            rasterio.io.MemoryFile
        """
        from pygndc.interop.interop import frame_to_rasterio

        self._check_open()
        t_idx, _ = self._resolve_time(t)
        frame = self._reader.reconstruct_single_frame(t_idx, band_idx=-1)
        if frame.ndim == 2:
            frame = frame[:, :, np.newaxis]
        return frame_to_rasterio(frame, self.crs, self.transform, self.nodata)

    def to_xarray(
        self, times: Optional[List[int]] = None
    ):
        """
        Convert to an xarray Dataset with coordinates and CRS.

        Args:
            times: List of time indices to include. None = all.

        Returns:
            xarray.Dataset
        """
        from pygndc.interop.interop import dataset_to_xarray

        self._check_open()
        return dataset_to_xarray(self, times=times)

    def to_numpy(self, t: Union[int, str, datetime] = 0) -> np.ndarray:
        """
        Export a single frame as numpy array (H, W, C).

        Args:
            t: Time index/string/datetime.

        Returns:
            np.ndarray (H, W, C)
        """
        return self.read(t=t)

    def to_tif(
        self,
        path: str,
        t: Union[int, str, datetime] = 0,
        out_dtype: str = "float32",
    ) -> None:
        """
        Export a single frame as GeoTIFF.

        Args:
            path: Output file path.
            t: Time index/string/datetime.
            out_dtype: Output data type.
        """
        from pygndc.interop.interop import save_frame_as_tif

        self._check_open()
        t_idx, _ = self._resolve_time(t)
        frame = self._reader.reconstruct_single_frame(t_idx, band_idx=-1)
        if frame.ndim == 2:
            frame = frame[:, :, np.newaxis]
        save_frame_as_tif(frame, path, self.crs, self.transform, out_dtype)

    def to_tifs(
        self,
        output_dir: str,
        out_dtype: str = "float32",
        progress: bool = True,
        name_prefix: str = "Frame",
    ) -> None:
        """
        Export all frames as GeoTIFF files.

        Args:
            output_dir: Output directory.
            out_dtype: Output data type.
            progress: Show progress bar.
            name_prefix: Filename prefix.
        """
        from pygndc.interop.interop import save_all_tifs

        self._check_open()
        save_all_tifs(self, output_dir, out_dtype, progress, name_prefix)

    # ------------------------------------------------------------------
    # Phase 3: Analysis convenience methods
    # ------------------------------------------------------------------

    def ndvi(
        self,
        t: Union[int, str, datetime] = 0,
        red_band: int = 0,
        nir_band: int = 1,
    ) -> np.ndarray:
        """
        Compute NDVI for a single frame.

        Args:
            t: Time index/string/datetime.
            red_band: 0-based red band index.
            nir_band: 0-based NIR band index.

        Returns:
            np.ndarray of shape (H, W).
        """
        from pygndc.analysis import compute_ndvi

        frame = self.read(t=t)
        return compute_ndvi(frame[:, :, red_band], frame[:, :, nir_band])

    def evi(
        self,
        t: Union[int, str, datetime] = 0,
        blue_band: int = 0,
        green_band: int = 1,
        nir_band: int = 2,
    ) -> np.ndarray:
        """
        Compute EVI for a single frame.

        Args:
            t: Time index/string/datetime.
            blue_band: 0-based blue band index.
            green_band: 0-based green/red band index (passed as 'green' to compute_evi).
            nir_band: 0-based NIR band index.

        Returns:
            np.ndarray of shape (H, W).
        """
        from pygndc.analysis import compute_evi

        frame = self.read(t=t)
        return compute_evi(
            frame[:, :, blue_band], frame[:, :, green_band], frame[:, :, nir_band]
        )

    def gradient(
        self,
        t: Union[int, str, datetime] = 0,
        mode: str = "spatial",
    ) -> Union[Tuple[np.ndarray, np.ndarray], np.ndarray]:
        """
        Compute analytic gradients using the neural network's differentiability.

        Args:
            t: Time index/string/datetime.
            mode: 'spatial' returns (dx, dy); 'mag' returns gradient magnitude;
                  'dx'/'dy' returns a single component.

        Returns:
            Tuple (dx, dy) for mode='spatial', or single array for 'mag'/'dx'/'dy'.
        """
        self._check_open()
        _, t_norm = self._resolve_time(t)

        if mode == "spatial":
            dx = self._reader.compute_spatial_derivatives(t_norm, mode="dx")
            dy = self._reader.compute_spatial_derivatives(t_norm, mode="dy")
            return dx, dy
        return self._reader.compute_spatial_derivatives(t_norm, mode=mode)

    def temporal_derivative(
        self,
        t: Union[int, str, datetime] = 0,
        order: int = 1,
        delta_t: float = 1e-3,
        red_idx: int = 0,
        nir_idx: int = 1,
    ) -> np.ndarray:
        """
        Compute NDVI temporal derivative using finite differences.

        Args:
            t: Time index/string/datetime.
            order: 1 for velocity, 2 for acceleration.
            delta_t: Finite difference step in normalized time.
            red_idx: 0-based red band index.
            nir_idx: 0-based NIR band index.

        Returns:
            np.ndarray of shape (H, W).
        """
        self._check_open()
        _, t_norm = self._resolve_time(t)
        return self._reader.compute_finite_temporal_derivatives_NDVI(
            t_norm, order=order, delta_t=delta_t, red_idx=red_idx, nir_idx=nir_idx
        )
