import io
import logging
import math
import re
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor

import geopandas as gpd

try:
    import imageio.v2 as imageio
except ImportError as imageio_import_error:  # pragma: no cover - fallback for environments without imageio
    imageio = None
import matplotlib
import matplotlib.pyplot as plt
import numpy as np

try:
    import odc
    import odc.stac
except ImportError as odc_import_error:
    class _ODCStub:
        def __getattr__(self, name):  # pragma: no cover - fallback path
            raise ImportError(
                "odc library is required for NDVI operations but could not be imported"
            ) from odc_import_error

    odc = _ODCStub()
    odc.stac = _ODCStub()
import pandas as pd
import planetary_computer
import pystac
import rioxarray  # noqa: F401
import xarray as xr
from matplotlib.ticker import ScalarFormatter
from shapely.geometry import Polygon
from tqdm.auto import tqdm

from environmental_risk_metrics.base import BaseEnvironmentalMetric
from environmental_risk_metrics.utils.planetary_computer import (
    get_planetary_computer_items,
)

matplotlib.use(backend="Agg")

logger = logging.getLogger(__name__)


class HarmonizedNDVI(BaseEnvironmentalMetric):
    """
    A class for accessing and analyzing NDVI data from multiple harmonized satellite collections
    including Sentinel-2, HLS2-S30, and HLS2-L30 from Planetary Computer.
    """
    
    # Band mappings for different collections
    BAND_MAPPINGS = {
        "sentinel-2-l2a": {
            "red": "B04",
            "green": "B03", 
            "blue": "B02",
            "nir": "B08",
            "cloud_mask": "SCL",
            "cloud_clear_values": [4, 5]  # SCL values for clear pixels
        },
        "hls2-s30": {
            "red": "B04",
            "green": "B03",
            "blue": "B02", 
            "nir": "B08",
            "cloud_mask": "Fmask",
            "cloud_clear_values": [0]  # Fmask value for clear pixels
        },
        "hls2-l30": {
            "red": "B04",
            "green": "B03",
            "blue": "B02",
            "nir": "B05",  # Different NIR band for Landsat
            "cloud_mask": "Fmask", 
            "cloud_clear_values": [0]  # Fmask value for clear pixels
        }
    }
    
    def __init__(
        self,
        start_date: str,
        end_date: str,
        gdf: gpd.GeoDataFrame,
        collections: list[str] = ["sentinel-2-l2a"],
        resolution: int = 30,
        max_entire_image_cloud_cover: int = 10,
        max_cropped_area_cloud_cover: int = 80,
        max_workers: int = 10,
        is_bare_soil_threshold: float = 0.25,
    ):
        sources = ["https://planetarycomputer.microsoft.com/api/stac/v1"]
        description = "Harmonized NDVI data from multiple satellite collections."

        super().__init__(sources=sources, description=description)

        self.collections = collections
        self.resolution = resolution
        self.max_entire_image_cloud_cover = max_entire_image_cloud_cover
        self.max_cropped_area_cloud_cover = max_cropped_area_cloud_cover
        self.max_workers = max_workers
        self.start_date = start_date
        self.end_date = end_date
        self.gdf = gdf
        self.is_bare_soil_threshold = is_bare_soil_threshold
        self.items = None
        self.xarray_data = None
        self.xarray_item_metadata = None
        self.ndvi_data = None
        self.ndvi_item_metadata = None
        self.mean_ndvi_data = None
        logger.debug("Initializing HarmonizedNDVI client")
        
        # Validate collections
        invalid_collections = set(collections) - set(self.BAND_MAPPINGS.keys())
        if invalid_collections:
            raise ValueError(f"Unsupported collections: {invalid_collections}. "
                           f"Supported collections: {list(self.BAND_MAPPINGS.keys())}")

    def get_items(
        self,
        max_entire_image_cloud_cover: int = None,
    ) -> dict[str, list[pystac.Item]]:
        """
        Search for items from all collections within a given date range and polygon.

        Args:
            max_entire_image_cloud_cover: Maximum cloud cover percentage allowed for entire image (excludes images with higher cloud cover)

        Returns:
            Dictionary with collection names as keys and list of pystac.Item objects as values
        """
        if self.items is not None:
            return self.items

        if max_entire_image_cloud_cover is None:
            max_entire_image_cloud_cover = self.max_entire_image_cloud_cover
            
        gdf = self.gdf.to_crs(epsg=4326)
        items_by_collection = {}

        def fetch_items_for_collection(collection_polygon_pair):
            collection, polygon = collection_polygon_pair
            return collection, get_planetary_computer_items(
                collections=[collection],
                start_date=self.start_date,
                end_date=self.end_date,
                polygon=polygon,
                max_entire_image_cloud_cover=max_entire_image_cloud_cover,
            )

        # Fetch items for each collection and geometry combination
        for collection in self.collections:
            items_by_collection[collection] = []
            
            collection_polygon_pairs = [(collection, geom) for geom in gdf.geometry]
            
            with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                results = list(tqdm(
                    executor.map(fetch_items_for_collection, collection_polygon_pairs), 
                    total=len(collection_polygon_pairs),
                    desc=f"Fetching {collection} items"
                ))
                
            # Organize results by geometry index
            for i, (_, items_list) in enumerate(results):
                if i >= len(items_by_collection[collection]):
                    items_by_collection[collection].extend([[] for _ in range(i + 1 - len(items_by_collection[collection]))])
                items_by_collection[collection][i] = items_list

        self.items = items_by_collection
        return self.items

    def _get_bands_for_collection(self, collection: str, include_rgb: bool = False) -> list[str]:
        """Get the required bands for a collection"""
        mapping = self.BAND_MAPPINGS[collection]
        bands = [mapping["nir"], mapping["red"]]
        if include_rgb:
            bands.extend([mapping["green"], mapping["blue"]])
        if mapping["cloud_mask"]:
            bands.append(mapping["cloud_mask"])
        return bands

    def load_xarray(
        self,
        bands: dict[str, list[str]] = None,
        show_progress: bool = True,
        filter_cloud_cover: bool = True,
        include_rgb: bool = False,
    ) -> dict[str, list[xr.Dataset]]:
        """Load data for all collections into xarray Datasets.

        Args:
            bands: Dictionary with collection names as keys and band lists as values
            show_progress: Whether to show a progress bar
            filter_cloud_cover: Whether to filter the data based on cloud cover
            include_rgb: Whether to include RGB bands for visualization

        Returns:
            Dictionary with collection names as keys and lists of xarray Datasets as values
        """
        if self.xarray_data is not None:
            return self.xarray_data
            
        logger.debug(f"Loading data for collections {self.collections} at {self.resolution}m resolution")
        
        items_dict = self.get_items()
        
        if not any(items_dict.values()):
            logger.error("No items found for any collection in the given date range and polygon")
            raise ValueError("No items found for any collection in the given date range and polygon")

        self.xarray_data = {}
        self.xarray_item_metadata = {}
        
        for collection in self.collections:
            self.xarray_data[collection] = []
            self.xarray_item_metadata[collection] = []
            items_list = items_dict.get(collection, [])
            
            if not items_list:
                logger.warning(f"No items found for collection {collection}")
                continue
                
            # Get bands for this collection
            if bands and collection in bands:
                collection_bands = bands[collection]
            else:
                collection_bands = self._get_bands_for_collection(collection, include_rgb)
            
            mapping = self.BAND_MAPPINGS[collection]
            
            for geometry_idx, (items, geometry) in enumerate(zip(items_list, self.gdf.geometry)):
                if not items:
                    logger.warning(f"No items for geometry {geometry_idx} in collection {collection}")
                    self.xarray_data[collection].append(None)
                    self.xarray_item_metadata[collection].append(None)
                    continue

                logger.debug(f"Signing {len(items)} items for collection {collection}")
                signed_items = [planetary_computer.sign(i) for i in items]

                metadata_map: dict[np.datetime64, list[dict]] = defaultdict(list)
                for item in items:
                    item_metadata = self._serialize_item_metadata(item)
                    item_datetime = item_metadata.get("datetime")
                    if item_datetime:
                        dt64 = np.datetime64(pd.to_datetime(item_datetime))
                        metadata_map[dt64].append(item_metadata)

                thread_pool = ThreadPoolExecutor(max_workers=self.max_workers)

                logger.debug(f"Loading data into xarray Dataset for {collection}")
                progress = tqdm if show_progress else None

                try:
                    ds = odc.stac.load(
                        items=signed_items,
                        bands=collection_bands,
                        resolution=self.resolution,
                        pool=thread_pool,
                        geopolygon=geometry,
                        progress=progress,
                    )
                except Exception as e:
                    logger.error(f"Failed to load data for {collection}, geometry {geometry_idx}: {e}")
                    self.xarray_data[collection].append(None)
                    self.xarray_item_metadata[collection].append(None)
                    continue

                # Apply cloud filtering if requested and cloud mask is available
                if self.max_cropped_area_cloud_cover and filter_cloud_cover and mapping["cloud_mask"]:
                    logger.debug(f"Filtering data based on cloud cover using {mapping['cloud_mask']} band")
                    
                    if collection == "sentinel-2-l2a":
                        # For Sentinel-2, use SCL band
                        cloud_clear_mask = (ds[mapping["cloud_mask"]] == 4) | (ds[mapping["cloud_mask"]] == 5)
                    else:
                        # For HLS collections, use Fmask band
                        cloud_clear_mask = ds[mapping["cloud_mask"]] == 0
                    
                    cloud_cover_pct = (1 - cloud_clear_mask.mean(dim=["x", "y"])) * 100
                    logger.debug(f"Cloud cover percentage: {cloud_cover_pct}. Filtering based on {self.max_cropped_area_cloud_cover}% threshold")
                    
                    logger.debug(f"Dataset time steps before filtering: {len(ds.time)}")
                    ds = ds.sel(time=cloud_cover_pct <= self.max_cropped_area_cloud_cover)
                    logger.debug(f"Filtered dataset to {len(ds.time)} time steps")

                    # Apply spatial cloud masking
                    if filter_cloud_cover:
                        ds = ds.where(cloud_clear_mask, drop=False)

                aligned_metadata: list[list[dict]] = []
                if metadata_map:
                    for time_val in ds.time.values:
                        np_time = np.datetime64(pd.to_datetime(time_val))
                        metadata_entries = metadata_map.get(np_time)
                        if metadata_entries is None:
                            for dt_key, value in metadata_map.items():
                                if pd.to_datetime(dt_key).date() == pd.to_datetime(np_time).date():
                                    metadata_entries = value
                                    break
                        aligned_metadata.append(metadata_entries or [])

                ds.attrs["item_metadata"] = aligned_metadata

                logger.debug(f"Successfully loaded {collection} data")
                self.xarray_data[collection].append(ds)
                self.xarray_item_metadata[collection].append(aligned_metadata)
                
        return self.xarray_data

    def load_ndvi_images(
        self,
        filter_cloud_cover: bool = True,
    ) -> dict[str, list[xr.DataArray]]:
        """Load NDVI data for all collections.

        Args:
            filter_cloud_cover: Whether to filter the data based on cloud cover

        Returns:
            Dictionary with collection names as keys and lists of NDVI DataArrays as values
        """
        if self.ndvi_data is not None:
            return self.ndvi_data
            
        logger.debug("Loading NDVI data for all collections")
        self.ndvi_data = {}
        self.ndvi_item_metadata = {}
        
        xarray_data = self.load_xarray(filter_cloud_cover=filter_cloud_cover)
        
        for collection in self.collections:
            self.ndvi_data[collection] = []
            self.ndvi_item_metadata[collection] = []
            datasets = xarray_data.get(collection, [])
            mapping = self.BAND_MAPPINGS[collection]
            collection_metadata = []
            if self.xarray_item_metadata and collection in self.xarray_item_metadata:
                collection_metadata = self.xarray_item_metadata[collection]
            
            for idx, ds in enumerate(datasets):
                if ds is None:
                    self.ndvi_data[collection].append(None)
                    self.ndvi_item_metadata[collection].append(None)
                    continue
                    
                logger.debug(f"Calculating NDVI for {collection} using bands {mapping['nir']} and {mapping['red']}")
                nir_band = ds[mapping["nir"]]
                red_band = ds[mapping["red"]]
                
                # Calculate NDVI
                ndvi = (nir_band - red_band) / (nir_band + red_band)
                # Mask out invalid values where denominator is zero
                ndvi = ndvi.where((nir_band + red_band) != 0)
                
                metadata: list[dict] | None = None
                if collection_metadata and idx < len(collection_metadata):
                    metadata = collection_metadata[idx]
                if metadata:
                    ndvi.attrs["item_metadata"] = metadata

                logger.debug(f"Successfully calculated NDVI for {collection}")
                self.ndvi_data[collection].append(ndvi)
                self.ndvi_item_metadata[collection].append(metadata)
                
        return self.ndvi_data

    def rgb_ndvi_images(
        self,
        vmin: float = -0.2,
        vmax: float = 0.8,
        boundary_color: str = "red",
        boundary_linewidth: float = 2,
        bbox_inches: str = "tight",
        pad_inches: float = 0.1,
        image_format: str = "png",
        timestamp_format: str = "%Y-%m-%d",
        figsize: tuple = (12, 6),
    ) -> dict[str, list[dict]]:
        """
        Generate side-by-side RGB and NDVI images for all collections and geometries.

        Args:
            vmin: Minimum value for NDVI color scale
            vmax: Maximum value for NDVI color scale
            boundary_color: Color of the polygon boundary
            boundary_linewidth: Line width of the polygon boundary
            bbox_inches: Bounding box setting for saving figure
            pad_inches: Padding when saving figure
            image_format: Format to save images in
            timestamp_format: Format string for timestamp keys
            figsize: Figure size as (width, height) tuple

        Returns:
            Dictionary with collection names as keys and lists of image dictionaries as values
        """
        logger.debug("Generating RGB+NDVI images for all collections")
        
        # Load data with RGB bands included
        xarray_data = self.load_xarray(include_rgb=True, filter_cloud_cover=True)
        ndvi_data = self.load_ndvi_images()
        
        outputs = {}
        
        for collection in self.collections:
            outputs[collection] = []
            datasets = xarray_data.get(collection, [])
            ndvi_datasets = ndvi_data.get(collection, [])
            mapping = self.BAND_MAPPINGS[collection]
            
            for geom_idx, (ds, ndvi, geometry) in enumerate(zip(datasets, ndvi_datasets, self.gdf.geometry)):
                metadata_list: list[dict] | None = None
                if ndvi is not None and hasattr(ndvi, "attrs"):
                    metadata_list = ndvi.attrs.get("item_metadata")
                geom_outputs = {
                    "images": {},
                    "metadata": metadata_list,
                }
                
                if ds is None or ndvi is None:
                    outputs[collection].append(geom_outputs)
                    continue
                
                crs = ds.coords["spatial_ref"].values.item()
                
                if imageio is None:
                    logger.warning("imageio not available; skipping RGB/NDVI image generation")
                    outputs[collection].append(geom_outputs)
                    continue

                for time in ds.time:
                    try:
                        # Create side-by-side plot
                        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)
                        
                        # RGB plot
                        rgb_data = ds.sel(time=time)[[mapping["red"], mapping["green"], mapping["blue"]]]
                        rgb_array = rgb_data.to_array().values
                        
                        # Normalize RGB values to 0-1 range for display
                        rgb_array = np.transpose(rgb_array, (1, 2, 0))
                        rgb_array = np.clip(rgb_array / np.percentile(rgb_array, 98), 0, 1)
                        
                        ax1.imshow(rgb_array)
                        ax1.set_title(f"RGB\n{pd.Timestamp(time.values).strftime(timestamp_format)}, {collection}")
                        ax1.set_axis_off()
                        
                        # Add boundary to RGB plot
                        gdf_crs = gpd.GeoDataFrame(
                            {"geometry": [geometry]}, geometry="geometry", crs=self.gdf.crs
                        ).to_crs(crs)
                        gdf_crs.boundary.plot(ax=ax1, color=boundary_color, linewidth=boundary_linewidth)
                        
                        # NDVI plot
                        ndvi_time = ndvi.sel(time=time)
                        ndvi_plot = ndvi_time.plot.imshow(
                            ax=ax2, cmap="RdYlGn", vmin=vmin, vmax=vmax, add_colorbar=False
                        )
                        ax2.set_title(f"NDVI\n{pd.Timestamp(time.values).strftime(timestamp_format)}, {collection}")
                        ax2.set_axis_off()
                        
                        # Add boundary to NDVI plot
                        gdf_crs.boundary.plot(ax=ax2, color=boundary_color, linewidth=boundary_linewidth)
                        
                        # Add colorbar for NDVI
                        cbar_ax = fig.add_axes([0.92, 0.25, 0.015, 0.5])
                        fig.colorbar(ndvi_plot, cax=cbar_ax, label="NDVI")
                        
                        # Format axes
                        for ax in (ax1, ax2):
                            ax.xaxis.set_major_formatter(ScalarFormatter(useMathText=True))
                            ax.ticklabel_format(style="sci", axis="x", scilimits=(0, 0))
                        
                        plt.tight_layout()
                        
                        # Save to bytes buffer
                        buf = io.BytesIO()
                        plt.savefig(buf, format=image_format, bbox_inches=bbox_inches, pad_inches=pad_inches, dpi=150)
                        buf.seek(0)
                        
                        # Add to dictionary with timestamp as key
                        timestamp = pd.Timestamp(time.values).strftime(timestamp_format)
                        geom_outputs["images"][timestamp] = buf.getvalue()
                        
                        plt.close(fig)
                        
                    except Exception as e:
                        logger.warning(f"Failed to generate image for {collection}, time {time}: {e}")
                        continue
                
                outputs[collection].append(geom_outputs)
                
        return outputs

    def generate_ndvi_gif(
        self,
        collection: str = None,
        geometry_index: int = 0,
        duration: float = 0.5,
        loop: int = 0,
        vmin: float = -0.2,  
        vmax: float = 0.8,
        figsize: tuple = (12, 6),
    ) -> bytes:
        """
        Generate a GIF showing the time series of RGB and NDVI images.

        Args:
            collection: Collection to use for GIF generation. If None, uses first available collection
            geometry_index: Index of geometry to use for GIF generation
            duration: Duration of each frame in seconds
            loop: Number of loops (0 for infinite)
            vmin: Minimum value for NDVI color scale
            vmax: Maximum value for NDVI color scale
            figsize: Figure size as (width, height) tuple

        Returns:
            GIF as bytes
        """
        if collection is None:
            collection = self.collections[0]
        
        if collection not in self.collections:
            raise ValueError(f"Collection {collection} not in available collections: {self.collections}")
        
        logger.debug(f"Generating GIF for collection {collection}, geometry {geometry_index}")
        
        # Get RGB+NDVI images
        rgb_ndvi_images = self.rgb_ndvi_images(vmin=vmin, vmax=vmax, figsize=figsize)
        
        if collection not in rgb_ndvi_images:
            raise ValueError(f"No data available for collection {collection}")
        
        if geometry_index >= len(rgb_ndvi_images[collection]):
            raise ValueError(f"Geometry index {geometry_index} out of range")
        
        images_dict = rgb_ndvi_images[collection][geometry_index]
        
        if not images_dict:
            raise ValueError(f"No images available for collection {collection}, geometry {geometry_index}")
        
        # Sort images by timestamp
        sorted_timestamps = sorted(images_dict.keys())
        frames = []
        
        for timestamp in sorted_timestamps:
            image_bytes = images_dict[timestamp]
            frame = imageio.imread(io.BytesIO(image_bytes))
            frames.append(frame)
        
        # Create GIF in memory
        gif_buffer = io.BytesIO()
        imageio.mimsave(gif_buffer, frames, format="GIF", duration=duration, loop=loop)
        gif_buffer.seek(0)
        
        logger.debug(f"Successfully generated GIF with {len(frames)} frames")
        return gif_buffer.getvalue()

    def ndvi_thumbnails(
        self,
        vmin: float = -0.2,
        vmax: float = 0.8,
        boundary_color: str = "red",
        boundary_linewidth: float = 2,
        add_colorbar: bool = False,
        add_labels: bool = False,
        bbox_inches: str = "tight",
        pad_inches: float = 0,
        image_format: str = "jpg",
        timestamp_format: str = "%Y-%m-%d",
    ) -> dict[str, list[dict]]:
        """
        Plot NDVI images and return them as jpgs in a dictionary

        Args:
            vmin: Minimum value for NDVI color scale
            vmax: Maximum value for NDVI color scale
            boundary_color: Color of the polygon boundary
            boundary_linewidth: Line width of the polygon boundary
            add_colorbar: Whether to add a colorbar to the plot
            add_labels: Whether to add labels to the plot
            bbox_inches: Bounding box setting for saving figure
            pad_inches: Padding when saving figure
            image_format: Format to save images in
            timestamp_format: Format string for timestamp keys

        Returns:
            Dictionary with collection names as keys and lists of image dictionaries as values
        """
        if self.ndvi_thumbnails_data is not None:
            return self.ndvi_thumbnails_data
            
        ndvi_data = self.load_ndvi_images()
        self.ndvi_thumbnails_data = {}

        for collection in self.collections:
            self.ndvi_thumbnails_data[collection] = []
            ndvi_list = ndvi_data.get(collection, [])
            
            for ndvi, geometry in zip(ndvi_list, self.gdf.geometry):
                images = {}
                
                if ndvi is None:
                    self.ndvi_thumbnails_data[collection].append(images)
                    continue
                
                crs = ndvi.coords["spatial_ref"].values.item()

                for time in ndvi.time:
                    try:
                        # Create new figure for each timestamp
                        fig, ax = plt.subplots()

                        # Plot NDVI data and polygon boundary
                        ndvi.sel(time=time).plot(
                            ax=ax,
                            vmin=vmin,
                            vmax=vmax,
                            add_colorbar=add_colorbar,
                            add_labels=add_labels,
                        )
                        gpd.GeoDataFrame(
                            {"geometry": [geometry]}, geometry="geometry", crs=self.gdf.crs
                        ).to_crs(crs).boundary.plot(
                            ax=ax, color=boundary_color, linewidth=boundary_linewidth
                        )
                        ax.set_axis_off()

                        # Save plot to bytes buffer
                        buf = io.BytesIO()
                        plt.savefig(
                            buf,
                            format=image_format,
                            bbox_inches=bbox_inches,
                            pad_inches=pad_inches,
                        )
                        buf.seek(0)

                        # Add to dictionary with timestamp as key
                        timestamp = pd.Timestamp(time.values).strftime(timestamp_format)
                        images[timestamp] = buf.getvalue()

                        # Close figure to free memory
                        plt.close(fig)
                    except Exception as e:
                        logger.warning(f"Failed to generate thumbnail for {collection}, time {time}: {e}")
                        continue
                        
                self.ndvi_thumbnails_data[collection].append(images)
                
        return self.ndvi_thumbnails_data

    def calculate_mean_ndvi(
        self,
        interpolate: bool = True,
        all_touched: bool = True,
    ) -> dict[str, list[pd.DataFrame]]:
        """
        Calculate mean NDVI value for the given polygon at each timestamp

        Args:
            interpolate (bool): Whether to interpolate missing values
            all_touched (bool): Whether to use all touched for clipping
        Returns:
            Dictionary with collection names as keys and lists of DataFrames with mean NDVI values as values
        """
        if self.mean_ndvi_data is not None:
            return self.mean_ndvi_data
            
        logger.debug("Calculating mean NDVI values for all collections")

        ndvi_data = self.load_ndvi_images()
        self.mean_ndvi_data = {}

        for collection in self.collections:
            self.mean_ndvi_data[collection] = []
            ndvi_images_list = ndvi_data.get(collection, [])
            
            for ndvi_images, geometry in zip(ndvi_images_list, self.gdf.geometry):
                if ndvi_images is None:
                    # Create empty dataframe for missing data
                    empty_df = pd.DataFrame(columns=["ndvi"])
                    if self.is_bare_soil_threshold:
                        empty_df["is_bare_soil"] = []
                    self.mean_ndvi_data[collection].append(empty_df)
                    continue
                    
                # Convert to rioxarray and clip once for all timestamps
                crs = ndvi_images.coords["spatial_ref"].values.item()
                ndvi_images = ndvi_images.rio.write_crs(crs)
                clipped_data = ndvi_images.rio.clip(
                    [geometry], self.gdf.crs, all_touched=all_touched
                )

                # Calculate means for all timestamps at once
                mean_values = clipped_data.mean(dim=["x", "y"]).values

                # Create dictionary mapping timestamps to means
                mean_ndvi = pd.DataFrame(
                    mean_values, columns=["ndvi"], index=clipped_data.time.values
                )
                mean_ndvi.index = pd.to_datetime(mean_ndvi.index)

                metadata_list = ndvi_images.attrs.get("item_metadata", []) if hasattr(ndvi_images, "attrs") else []
                metadata_map: dict[np.datetime64, list[dict]] = {}
                if metadata_list:
                    for metadata_entry in metadata_list:
                        if not metadata_entry:
                            continue
                        datetime_str = metadata_entry.get("datetime") or metadata_entry.get("properties", {}).get("datetime")
                        if not datetime_str:
                            continue
                        dt64 = np.datetime64(pd.to_datetime(datetime_str))
                        metadata_map.setdefault(dt64, []).append(metadata_entry)

                if metadata_map:
                    metadata_records = []
                    for time_val in mean_ndvi.index.values:
                        np_time = np.datetime64(pd.to_datetime(time_val))
                        metadata_records.append(metadata_map.get(np_time) or [])
                    mean_ndvi["metadata"] = metadata_records
                else:
                    mean_ndvi["metadata"] = [[] for _ in range(len(mean_ndvi))]

                if interpolate:
                    mean_ndvi = interpolate_ndvi(mean_ndvi, self.start_date, self.end_date)

                if self.is_bare_soil_threshold:
                    mean_ndvi["is_bare_soil"] = mean_ndvi["ndvi"] < self.is_bare_soil_threshold

                logger.debug(f"Calculated mean NDVI for {collection} with {len(mean_ndvi)} timestamps")
                self.mean_ndvi_data[collection].append(mean_ndvi)
                
        return self.mean_ndvi_data

    def get_data(
        self,
        all_touched: bool = True,
        interpolate: bool = True,
    ) -> dict[str, list[list[dict]]]:
        """Get mean NDVI values for all collections and geometries"""
        mean_ndvi_data = self.calculate_mean_ndvi(
            interpolate=interpolate,
            all_touched=all_touched,
        )
        output = {}
        
        for collection in self.collections:
            output[collection] = []
            mean_ndvi_df_list = mean_ndvi_data.get(collection, [])
            
            for mean_ndvi_df in mean_ndvi_df_list:
                if mean_ndvi_df.empty:
                    output[collection].append([])
                    continue
                    
                mean_ndvi_df_copy = mean_ndvi_df.reset_index(names="date").copy()
                mean_ndvi_dict = mean_ndvi_df_copy.to_dict(orient="records")
                
                for idx, record in enumerate(mean_ndvi_dict):
                    if 'ndvi' in record:
                        if pd.isna(record['ndvi']):
                            record.pop("ndvi")
                        else:
                            record['ndvi'] = round(record['ndvi'], 2)
                    if 'interpolated_ndvi' in record:
                        record['interpolated_ndvi'] = round(record['interpolated_ndvi'], 2)
                    if 'metadata' in record:
                        metadata_val = record['metadata']
                        if not metadata_val:
                            record.pop('metadata')
                        else:
                            cleaned_metadata = []
                            for metadata_entry in metadata_val:
                                if not isinstance(metadata_entry, dict):
                                    continue
                                cleaned_entry = metadata_entry.copy()
                                assets = cleaned_entry.get('assets')
                                if isinstance(assets, dict):
                                    cleaned_entry['assets'] = {
                                        asset_key: {
                                            'href': asset_value.get('href')
                                        }
                                        for asset_key, asset_value in assets.items()
                                        if isinstance(asset_value, dict) and 'href' in asset_value
                                    }
                                cleaned_metadata.append(cleaned_entry)
                            record['metadata'] = cleaned_metadata
                        
                output[collection].append(mean_ndvi_dict)
                
        return output

    @staticmethod
    def _serialize_item_metadata(item: pystac.Item) -> dict:
        """Serialize a STAC item preserving original asset hrefs and relevant metadata."""
        item_dict = item.to_dict(include_self_link=True)
        datetime_value = item_dict.get("properties", {}).get("datetime")
        if datetime_value:
            item_dict["datetime"] = datetime_value
        elif item.datetime:
            item_dict["datetime"] = item.datetime.isoformat()
        return item_dict


def interpolate_ndvi(df: pd.DataFrame, start_date: str, end_date: str):
    """
    Create a DataFrame from NDVI values, interpolate missing dates, and plot the results.

    Args:
        mean_ndvi_values (dict): Dictionary of dates and NDVI values
        start_date (str): Start date in YYYY-MM-DD format
        end_date (str): End date in YYYY-MM-DD format

    Returns:
        pd.DataFrame: DataFrame with interpolated daily NDVI values
    """
    date_range = pd.date_range(
        start=pd.to_datetime(start_date), end=pd.to_datetime(end_date), freq="D"
    )
    df = df.reindex(date_range)
    df["interpolated_ndvi"] = df["ndvi"].interpolate(method="linear", limit_direction="both")
    return df


# Backward compatibility alias
class Sentinel2(HarmonizedNDVI):
    """
    Backward compatibility class for Sentinel2. 
    This is now an alias for HarmonizedNDVI with Sentinel-2 as the default collection.
    """
    def __init__(
        self,
        start_date: str,
        end_date: str,
        gdf: gpd.GeoDataFrame,
        resolution: int = 10,
        max_entire_image_cloud_cover: int = 10,
        max_cropped_area_cloud_cover: int = 80,
        max_workers: int = 10,
        is_bare_soil_threshold: float = 0.25,
    ):
        # Call parent with Sentinel-2 collection only for backward compatibility
        super().__init__(
            start_date=start_date,
            end_date=end_date,
            gdf=gdf,
            collections=["sentinel-2-l2a"],
            resolution=resolution,
            max_entire_image_cloud_cover=max_entire_image_cloud_cover,
            max_cropped_area_cloud_cover=max_cropped_area_cloud_cover,
            max_workers=max_workers,
            is_bare_soil_threshold=is_bare_soil_threshold,
        )
        
    def get_data(self, all_touched: bool = True, interpolate: bool = True) -> list[list[dict]]:
        """Get mean NDVI values for Sentinel-2 data (backward compatibility method)"""
        data = super().get_data(all_touched=all_touched, interpolate=interpolate)
        # Return just the Sentinel-2 data for backward compatibility
        return data.get("sentinel-2-l2a", [])
