# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.

import os
import random
from collections.abc import Sequence
from pathlib import Path
from typing import Any, Union

import albumentations as A
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import rasterio
import torch
from matplotlib.figure import Figure
from torch import Tensor

from terratorch.datasets.generic_multimodal_dataset import MultimodalToTensor
from terratorch.datasets.transforms import MultimodalTransforms
from terratorch.datasets.utils import default_transform
from torchgeo.datasets import BioMassters
from torchgeo.datasets.utils import percentile_normalization


class BioMasstersNonGeo(BioMassters):
    """[BioMassters Dataset](https://huggingface.co/datasets/ibm-nasa-geospatial/BioMassters) for Aboveground Biomass prediction.

    Dataset intended for Aboveground Biomass (AGB) prediction
    over Finnish forests based on Sentinel 1 and 2 data with
    corresponding target AGB mask values generated by Light Detection
    and Ranging (LiDAR).

    Dataset Format:

    * .tif files for Sentinel 1 and 2 data
    * .tif file for pixel wise AGB target mask
    * .csv files for metadata regarding features and targets

    Dataset Features:

    * 13,000 target AGB masks of size (256x256px)
    * 12 months of data per target mask
    * Sentinel 1 and Sentinel 2 data for each location
    * Sentinel 1 available for every month
    * Sentinel 2 available for almost every month
      (not available for every month due to ESA acquisition halt over the region
      during particular periods)

    If you use this dataset in your research, please cite the following paper:

    * https://nascetti-a.github.io/BioMasster/

    .. versionadded:: 0.5
    """

    S1_BAND_NAMES = ["VV_Asc", "VH_Asc", "VV_Desc", "VH_Desc", "RVI_Asc", "RVI_Desc"]
    S2_BAND_NAMES = [
        "BLUE",
        "GREEN",
        "RED",
        "RED_EDGE_1",
        "RED_EDGE_2",
        "RED_EDGE_3",
        "NIR_BROAD",
        "NIR_NARROW",
        "SWIR_1",
        "SWIR_2",
        "CLOUD_PROBABILITY",
    ]

    all_band_names = {
        "S1": S1_BAND_NAMES,
        "S2": S2_BAND_NAMES,
    }

    rgb_bands = {
        "S1": [],
        "S2": ["RED", "GREEN", "BLUE"],
    }

    valid_splits = ("train", "test")
    valid_sensors = ("S1", "S2")

    BAND_SETS = {"all": all_band_names, "rgb": rgb_bands}

    default_metadata_filename = "The_BioMassters_-_features_metadata.csv.csv"

    def __init__(
        self,
        root = "data",
        split: str = "train",
        bands: dict[str, Sequence[str]] | Sequence[str] = BAND_SETS["all"],
        transform: A.Compose | None = None,
        mask_mean: float | None = 63.4584,
        mask_std: float | None = 72.21242,
        sensors: Sequence[str] = ["S1", "S2"],
        as_time_series: bool = False,
        metadata_filename: str = default_metadata_filename,
        max_cloud_percentage: float | None = None,
        max_red_mean: float | None = None,
        include_corrupt: bool = True,
        subset: float = 1,
        seed: int = 42,
        use_four_frames: bool = False
    ) -> None:
        """Initialize a new instance of BioMassters dataset.

        If ``as_time_series=False`` (the default), each time step becomes its own
        sample with the target being shared across multiple samples.

        Args:
            root: root directory where dataset can be found
            split: train or test split
            sensors: which sensors to consider for the sample, Sentinel 1 and/or
                Sentinel 2 ('S1', 'S2')
            as_time_series: whether or not to return all available
                time-steps or just a single one for a given target location
            metadata_filename: metadata file to be used
            max_cloud_percentage: maximum allowed cloud percentage for images
            max_red_mean: maximum allowed red_mean value for images
            include_corrupt: whether to include images marked as corrupted

        Raises:
            AssertionError: if ``split`` or ``sensors`` is invalid
            DatasetNotFoundError: If dataset is not found.
        """
        self.root = root
        self.sensors = sensors
        self.bands = bands
        assert (
            split in self.valid_splits
        ), f"Please choose one of the valid splits: {self.valid_splits}."
        self.split = split

        assert set(sensors).issubset(
            set(self.valid_sensors)
        ), f"Please choose a subset of valid sensors: {self.valid_sensors}."

        if len(self.sensors) == 1:
            sens = self.sensors[0]
            self.band_indices = [
                self.all_band_names[sens].index(band) for band in self.bands[sens]
            ]
        else:
            self.band_indices = {
                sens: [self.all_band_names[sens].index(band) for band in self.bands[sens]]
                for sens in self.sensors
            }

        self.mask_mean = mask_mean
        self.mask_std = mask_std
        self.as_time_series = as_time_series
        self.metadata_filename = metadata_filename
        self.max_cloud_percentage = max_cloud_percentage
        self.max_red_mean = max_red_mean
        self.include_corrupt = include_corrupt
        self.subset = subset
        self.seed = seed
        self.use_four_frames = use_four_frames

        self._verify()

        # open metadata csv files
        self.df = pd.read_csv(os.path.join(self.root, self.metadata_filename))

        # Filter sensors
        self.df = self.df[self.df["satellite"].isin(self.sensors)]

        # Filter split
        self.df = self.df[self.df["split"] == self.split]

        # Optional filtering
        self._filter_and_select_data()

        # Optional subsampling
        self._random_subsample()

        # generate numerical month from filename since first month is September
        # and has numerical index of 0
        self.df["num_month"] = (
            self.df["filename"]
            .str.split("_", expand=True)[2]
            .str.split(".", expand=True)[0]
            .astype(int)
        )

        # Set dataframe index depending on the task for easier indexing
        if self.as_time_series:
            self.df["num_index"] = self.df.groupby(["chip_id"]).ngroup()
        else:
            filter_df = (
                self.df.groupby(["chip_id", "month"])["satellite"].count().reset_index()
            )
            filter_df = filter_df[
                filter_df["satellite"] == len(self.sensors)
            ].drop("satellite", axis=1)
            # Guarantee that each sample has corresponding number of images available
            self.df = self.df.merge(filter_df, on=["chip_id", "month"], how="inner")

            self.df["num_index"] = self.df.groupby(["chip_id", "month"]).ngroup()

        # Adjust transforms based on the number of sensors
        if len(self.sensors) == 1:
            self.transform = transform if transform else default_transform
        elif transform is None:
            self.transform = MultimodalToTensor(self.sensors)
        else:
            transform = {
                s: transform[s] if s in transform else default_transform
                for s in self.sensors
            }
            self.transform = MultimodalTransforms(transform, shared=False)

        if self.use_four_frames:
            self._select_4_frames()

    def __len__(self) -> int:
        return len(self.df["num_index"].unique())

    def _load_input(self, filenames: list[Path]) -> Tensor:
        """Load the input imagery at the index.

        Args:
            filenames: list of filenames corresponding to input

        Returns:
            input image
        """
        filepaths = [
            os.path.join(self.root, f"{self.split}_features", f) for f in filenames
        ]
        arr_list = [rasterio.open(fp).read() for fp in filepaths]

        if self.as_time_series:
            arr = np.stack(arr_list, axis=0) # (T, C, H, W)
        else:
            arr = np.concatenate(arr_list, axis=0)
        return arr.astype(np.int32)

    def _load_target(self, filename: Path) -> Tensor:
        """Load the target mask at the index.

        Args:
            filename: filename of target to index

        Returns:
            target mask
        """
        with rasterio.open(os.path.join(self.root, f"{self.split}_agbm", filename), "r") as src:
            arr: np.typing.NDArray[np.float64] = src.read()

        return arr

    def _compute_rvi(self, img: np.ndarray, linear: np.ndarray, sens: str) -> np.ndarray:
        """Compute the RVI indices for S1 data."""
        rvi_channels = []
        if self.as_time_series:
            if "RVI_Asc" in self.bands[sens]:
                try:
                    vv_asc_index = self.all_band_names["S1"].index("VV_Asc")
                    vh_asc_index = self.all_band_names["S1"].index("VH_Asc")
                except ValueError as e:
                    msg = f"RVI_Asc needs band: {e}"
                    raise ValueError(msg) from e

                VV = linear[:, vv_asc_index, :, :]
                VH = linear[:, vh_asc_index, :, :]
                rvi_asc = 4 * VH / (VV + VH + 1e-6)
                rvi_asc = np.expand_dims(rvi_asc, axis=1)
                rvi_channels.append(rvi_asc)
            if "RVI_Desc" in self.bands[sens]:
                try:
                    vv_desc_index = self.all_band_names["S1"].index("VV_Desc")
                    vh_desc_index = self.all_band_names["S1"].index("VH_Desc")
                except ValueError as e:
                    msg = f"RVI_Desc needs band: {e}"
                    raise ValueError(msg) from e

                VV_desc = linear[:, vv_desc_index, :, :]
                VH_desc = linear[:, vh_desc_index, :, :]
                rvi_desc = 4 * VH_desc / (VV_desc + VH_desc + 1e-6)
                rvi_desc = np.expand_dims(rvi_desc, axis=1)
                rvi_channels.append(rvi_desc)
            if rvi_channels:
                rvi_concat = np.concatenate(rvi_channels, axis=1)
                img = np.concatenate([img, rvi_concat], axis=1)
        else:
            if "RVI_Asc" in self.bands[sens]:
                if linear.shape[0] < 2:
                    msg = f"Not enough bands to calculate RVI_Asc. Available bands: {linear.shape[0]}"
                    raise ValueError(msg)
                VV = linear[0]
                VH = linear[1]
                rvi_asc = 4 * VH / (VV + VH + 1e-6)
                rvi_asc = np.expand_dims(rvi_asc, axis=0)
                rvi_channels.append(rvi_asc)
            if "RVI_Desc" in self.bands[sens]:
                if linear.shape[0] < 4:
                    msg = f"Not enough bands to calculate RVI_Desc. Available bands: {linear.shape[0]}"
                    raise ValueError(msg)
                VV_desc = linear[2]
                VH_desc = linear[3]
                rvi_desc = 4 * VH_desc / (VV_desc + VH_desc + 1e-6)
                rvi_desc = np.expand_dims(rvi_desc, axis=0) 
                rvi_channels.append(rvi_desc)
            if rvi_channels:
                rvi_concat = np.concatenate(rvi_channels, axis=0)
                img = np.concatenate([linear, rvi_concat], axis=0)
        return img

    def _select_4_frames(self):
        """Filter the dataset to select only 4 frames per sample."""

        if "cloud_percentage" in self.df.columns:
            self.df = self.df.sort_values(by=["chip_id", "cloud_percentage"])
        else:
            self.df = self.df.sort_values(by=["chip_id", "num_month"])

        self.df = (
            self.df.groupby("chip_id")
            .head(4)  # Select the first 4 frames per chip
            .reset_index(drop=True)
        )

    def _process_sensor_images(self, sens: str, sens_filepaths: list[str]) -> np.ndarray:
        """Process images for a given sensor."""
        img = self._load_input(sens_filepaths)
        if sens == "S1":
            img = img.astype(np.float32)
            linear = 10 ** (img / 10)
            img = self._compute_rvi(img, linear, sens)
        if self.as_time_series:
            img = img.transpose(0, 2, 3, 1)  # (T, H, W, C)
        else:
            img = img.transpose(1, 2, 0)  # (H, W, C)
        if len(self.sensors) == 1:
            img = img[..., self.band_indices]
        else:
            img = img[..., self.band_indices[sens]]
        return img

    def __getitem__(self, index: int) -> dict:
        sample_df = self.df[self.df["num_index"] == index].copy()
        # Sort by satellite and month
        sample_df.sort_values(
            by=["satellite", "num_month"], inplace=True, ascending=True
        )

        filepaths = sample_df["filename"].tolist()
        output = {}

        if len(self.sensors) == 1:
            sens = self.sensors[0]
            sens_filepaths = [fp for fp in filepaths if sens in fp]
            img = self._process_sensor_images(sens, sens_filepaths)
            output["image"] = img.astype(np.float32)
        else:
            for sens in self.sensors:
                sens_filepaths = [fp for fp in filepaths if sens in fp]
                img = self._process_sensor_images(sens, sens_filepaths)
                output[sens] = img.astype(np.float32)

        # Load target
        target_filename = sample_df["corresponding_agbm"].unique()[0]
        target = np.array(self._load_target(Path(target_filename)))
        target = target.transpose(1, 2, 0)
        output["mask"] = target
        if self.transform:
            if len(self.sensors) == 1:
                output = self.transform(**output)
            else:
                output = self.transform(output)
        output["mask"] = output["mask"].squeeze().float()
        return output

    def _filter_and_select_data(self):
        if (
            self.max_cloud_percentage is not None
            and "cloud_percentage" in self.df.columns
        ):
            self.df = self.df[self.df["cloud_percentage"] <= self.max_cloud_percentage]

        if self.max_red_mean is not None and "red_mean" in self.df.columns:
            self.df = self.df[self.df["red_mean"] <= self.max_red_mean]

        if not self.include_corrupt and "corrupt_values" in self.df.columns:
            self.df = self.df[self.df["corrupt_values"] is False]

    def _random_subsample(self):
        if self.split == "train" and self.subset < 1.0:
            num_samples = int(len(self.df["num_index"].unique()) * self.subset)
            if self.seed is not None:
                random.seed(self.seed)
            selected_indices = random.sample(
                list(self.df["num_index"].unique()), num_samples
            )
            self.df = self.df[self.df["num_index"].isin(selected_indices)]
            self.df.reset_index(drop=True, inplace=True)

    def plot(
        self,
        sample: dict[str, Tensor],
        show_titles: bool = True,
        suptitle: str | None = None,
    ) -> Figure:
        """Plot a sample from the dataset.

        Args:
            sample: a sample returned by :meth:`__getitem__`
            show_titles: flag indicating whether to show titles above each panel
            suptitle: optional suptitle to use for figure

        Returns:
            a matplotlib Figure with the rendered sample
        """
        # Determine if the sample contains multiple sensors or a single sensor
        if isinstance(sample["image"], dict):
            ncols = len(self.sensors) + 1
        else:
            ncols = 2  # One for the image and one for the mask

        showing_predictions = "prediction" in sample
        if showing_predictions:
            ncols += 1

        fig, axs = plt.subplots(1, ncols=ncols, figsize=(5 * ncols, 10))

        if isinstance(sample["image"], dict):
            # Multiple sensors case
            for idx, sens in enumerate(self.sensors):
                img = sample["image"][sens].numpy()
                if self.as_time_series:
                    # Plot last time step
                    img = img[:, -1, ...]
                if sens == "S2":
                    img = img[[2, 1, 0], ...].transpose(1, 2, 0)
                    img = percentile_normalization(img)
                else:
                    co_polarization = img[0]  # transmit == receive
                    cross_polarization = img[1]  # transmit != receive
                    ratio = co_polarization / (cross_polarization + 1e-6)

                    co_polarization = np.clip(co_polarization / 0.3, 0, 1)
                    cross_polarization = np.clip(cross_polarization / 0.05, 0, 1)
                    ratio = np.clip(ratio / 25, 0, 1)

                    img = np.stack(
                        (co_polarization, cross_polarization, ratio), axis=0
                    )
                    img = img.transpose(1, 2, 0)  # Convert to (H, W, 3)

                axs[idx].imshow(img)
                axs[idx].axis("off")
                if show_titles:
                    axs[idx].set_title(sens)
            mask_idx = len(self.sensors)
        else:
            # Single sensor case
            sens = self.sensors[0]
            img = sample["image"].numpy()
            if self.as_time_series:
                # Plot last time step
                img = img[:, -1, ...]
            if sens == "S2":
                img = img[[2, 1, 0], ...].transpose(1, 2, 0)
                img = percentile_normalization(img)
            else:
                co_polarization = img[0]  # transmit == receive
                cross_polarization = img[1]  # transmit != receive
                ratio = co_polarization / (cross_polarization + 1e-6)

                co_polarization = np.clip(co_polarization / 0.3, 0, 1)
                cross_polarization = np.clip(cross_polarization / 0.05, 0, 1)
                ratio = np.clip(ratio / 25, 0, 1)

                img = np.stack(
                    (co_polarization, cross_polarization, ratio), axis=0
                )
                img = img.transpose(1, 2, 0)  # Convert to (H, W, 3)

            axs[0].imshow(img)
            axs[0].axis("off")
            if show_titles:
                axs[0].set_title(sens)
            mask_idx = 1

        # Plot target mask
        if "mask" in sample:
            target = sample["mask"].squeeze()
            target_im = axs[mask_idx].imshow(target, cmap="YlGn")
            plt.colorbar(target_im, ax=axs[mask_idx], fraction=0.046, pad=0.04)
            axs[mask_idx].axis("off")
            if show_titles:
                axs[mask_idx].set_title("Target")

        # Plot prediction if available
        if showing_predictions:
            pred_idx = mask_idx + 1
            prediction = sample["prediction"].squeeze()
            pred_im = axs[pred_idx].imshow(prediction, cmap="YlGn")
            plt.colorbar(pred_im, ax=axs[pred_idx], fraction=0.046, pad=0.04)
            axs[pred_idx].axis("off")
            if show_titles:
                axs[pred_idx].set_title("Prediction")

        if suptitle is not None:
            plt.suptitle(suptitle)

        return fig
