import random

import geopandas as gpd
import numpy as np
import pandas as pd
from immutabledict import immutabledict

from agrigee_lite.sat.abstract_satellite import AbstractSatellite

try:
    import torch
    from torch.utils.data import Dataset

    _has_torch = True
except ImportError:
    Dataset = object
    _has_torch = False


class AgriGEELiteDataset(Dataset):
    default_emr_order = ("coastal", "blue", "green", "red", "re1", "re2", "re3", "nir", "re4", "swir1", "swir2")

    default_means = immutabledict({
        "blue": 0.0656,
        "green": 0.0948,
        "red": 0.1094,
        "re1": 0.1507,
        "re2": 0.2372,
        "re3": 0.2673,
        "nir": 0.2866,
        "re4": 0.2946,
        "swir1": 0.2679,
        "swir2": 0.1985,
    })

    # Means Landsat
    # blue                                     0.037266
    # green                                    0.071555
    # nir                                      0.307463
    # red                                      0.082291
    # swir1                                    0.260797
    # swir2                                    0.149657
    # tirs1                                   27.100378
    # tirs2                                   24.983678

    # Sentinel
    # blue                                     0.034531
    # green                                    0.068706
    # nir                                       0.27561
    # red                                      0.083885
    # swir1                                    0.270749
    # swir2                                    0.155194

    # Stds Landsat
    # blue                                   0.016576
    # green                                  0.023892
    # nir                                    0.072974
    # red                                    0.044285
    # swir1                                  0.081166
    # swir2                                  0.066548
    # tirs1                                  4.623496
    # tirs2                                  4.324476

    default_stds = immutabledict({
        "blue": 0.0363,
        "green": 0.0433,
        "red": 0.06476,
        "re1": 0.05795,
        "re2": 0.07416,
        "re3": 0.09644,
        "nir": 0.09784,
        "re4": 0.0984,
        "swir1": 0.08984,
        "swir2": 0.09784,
    })

    def __init__(
        self,
        gdf: str | gpd.GeoDataFrame,
        sits_df: str | pd.DataFrame,
        timestamp_processing="doy",
        transform=None,
        normalize=True,
        pytorch=True,
    ):
        if not _has_torch:
            raise ImportError("AgriGEELiteDataset requer 'torch' instalado. Instale com 'pip install torch'.")  # noqa: TRY003

        super().__init__()

        self.normalize = normalize
        self.pytorch = pytorch

        if isinstance(gdf, str):
            self.gdf = gpd.read_parquet(gdf)
        else:
            self.gdf = gdf.copy()

        if isinstance(sits_df, str):
            self.sits_df = pd.read_parquet(sits_df)
        else:
            self.sits_df = sits_df.copy()

        sits_cols = set(self.sits_df.columns)
        gdf_cols = set(self.gdf.columns)
        default_bands = set(AgriGEELiteDataset.default_emr_order)
        sits_extra_cols = sits_cols - default_bands - {"timestamp"}
        gdf_extra_cols = gdf_cols - default_bands - {"timestamp"}
        common_index_cols = sits_extra_cols & gdf_extra_cols

        if common_index_cols:
            print(f"Coluna(s) de índice comum encontrada(s): {common_index_cols}")
            index_col = list(common_index_cols)[0]  # noqa: RUF015
            self.sits_df = self.sits_df.rename(columns={index_col: "indexnum"})
            self.gdf = self.gdf.rename(columns={index_col: "indexnum"})
            self.gdf = self.gdf.reset_index(drop=True)
        elif "original_index" in sits_cols:
            self.sits_df = self.sits_df.rename(columns={"original_index": "indexnum"})
            self.gdf = self.gdf.reset_index().rename(columns={"index": "indexnum"})
        elif "indexnum" in sits_cols:
            self.gdf = self.gdf.reset_index().rename(columns={"index": "indexnum"})
        else:
            raise ValueError("Não foi possível identificar a coluna de índice comum entre sits_df e gdf.")  # noqa: TRY003

        # Keep only the intersection of indexes between sits_df and gdf
        print(f"GDF: {len(self.gdf)} parcelas, SITS: {len(self.sits_df)} amostras")
        sits_indexes = set(self.sits_df["indexnum"].unique())
        gdf_indexes = set(self.gdf["indexnum"].unique())

        valid_indexes = sits_indexes & gdf_indexes  # intersection
        print(f"Interseção: {len(valid_indexes)} parcelas válidas")

        # Filter both dataframes to keep only valid indexes
        self.gdf = self.gdf[self.gdf["indexnum"].isin(valid_indexes)].reset_index(drop=True)
        self.sits_df = self.sits_df[self.sits_df["indexnum"].isin(valid_indexes)].reset_index(drop=True)

        print(f"Após filtro - GDF: {len(self.gdf)} parcelas, SITS: {len(self.sits_df)} amostras")

        # Remap indexnum to be sequential from 0 to N
        unique_indexes = sorted(self.gdf["indexnum"].unique())
        index_mapping = {old_idx: new_idx for new_idx, old_idx in enumerate(unique_indexes)}
        self.gdf["indexnum"] = self.gdf["indexnum"].map(index_mapping)
        self.sits_df["indexnum"] = self.sits_df["indexnum"].map(index_mapping)

        self.transform = transform
        self.bands = [band for band in AgriGEELiteDataset.default_emr_order if band in self.sits_df.columns]
        self.timestamp_processing = timestamp_processing
        self.max_seq_len = self.sits_df.groupby("indexnum").size().max()

        self.sits_df["timestamp"] = pd.to_datetime(self.sits_df["timestamp"])
        self.sequence_length = self.max_seq_len

        if self.timestamp_processing == "doy":
            self.sits_df["timestamp"] = self.sits_df["timestamp"].dt.dayofyear
        elif self.timestamp_processing == "days_after_start":
            start_map = self.gdf.set_index("indexnum")["start_date"]
            starts = self.sits_df["indexnum"].map(start_map)
            self.sits_df["timestamp"] = (self.sits_df["timestamp"] - starts).dt.days
            self.sits_df["timestamp"] = self.sits_df["timestamp"] + 1

        self.xs, self.doys = self.to_numpy_arrays_wo_for()
        self.ys = self.gdf["crop_class"].values
        self.num_classes = int(self.gdf["crop_class"].nunique())

        if normalize:
            self.bands_means = np.array([self.default_means[band] for band in self.bands], dtype=np.float16)
            self.bands_stds = np.array([self.default_stds[band] for band in self.bands], dtype=np.float16)

    def to_numpy_arrays_wo_for(self):
        n_samples = len(self.gdf)
        n_bands = len(self.bands)

        X = np.full((n_samples, self.max_seq_len, n_bands), 0, dtype=np.float16)
        T = np.full((n_samples, self.max_seq_len), 0, dtype=np.int16)

        sits_sorted = self.sits_df.sort_values(["indexnum", "timestamp"]).copy()

        index_map = {idx: i for i, idx in enumerate(self.gdf["indexnum"].values)}
        sits_sorted["parcel_idx"] = sits_sorted["indexnum"].map(index_map)

        sits_sorted["seq_idx"] = sits_sorted.groupby("indexnum").cumcount()

        sits_sorted = sits_sorted[sits_sorted["seq_idx"] < self.max_seq_len]

        band_values = sits_sorted[self.bands].to_numpy()

        time_values = sits_sorted["timestamp"].to_numpy()

        pi = sits_sorted["parcel_idx"].to_numpy()
        si = sits_sorted["seq_idx"].to_numpy()

        X[pi, si, :] = band_values
        T[pi, si] = time_values

        return X, T

    def __len__(self) -> int:
        return self.ys.shape[0]

    def __getitem__(self, idx: int):
        sample = {"doy": self.doys[idx], "x": self.xs[idx], "y": self.ys[idx]}

        if self.transform:
            sample = self.transform(sample)

        if self.normalize:
            sample["x"] = (sample["x"] - self.bands_means) / self.bands_stds

        if self.pytorch:
            for key, value in sample.items():
                if (value.dtype == np.half) or (value.dtype == np.float32) or (value.dtype == np.float64):
                    sample[key] = torch.tensor(value, dtype=torch.float32)
                elif value.dtype == np.bool_:
                    sample[key] = torch.tensor(value, dtype=torch.bool)
                else:
                    sample[key] = torch.tensor(value, dtype=torch.int64)

        return sample


# ------------------ Augmentations ------------------ #


def random_masking(ts, ts_length, seq_len, dimension):
    ts_masking = ts.copy()
    mask = np.zeros((seq_len,), dtype=int)

    for i in range(ts_length):
        prob = random.random()  # noqa: S311
        if prob < 0.15:
            prob /= 0.15
            mask[i] = 1

            if prob < 0.5:
                ts_masking[i, :] += np.random.uniform(low=-0.5, high=0, size=(dimension,))

            else:
                ts_masking[i, :] += np.random.uniform(low=0, high=0.5, size=(dimension,))

    return ts_masking, mask


# ---------------------------- Cloud and Cloud Shadow Augmentation --------------


class RandomCloudAugmentation:
    def __init__(self, probability=0.5):
        self.probability = probability

    def __call__(self, sample):
        if np.random.rand() <= self.probability:
            x = sample["x"]
            doy = sample["doy"]
            x, _ = random_masking(x, x.shape[0] - np.sum(doy == 0), x.shape[0], x.shape[1])
            sample["x"] = x

        return sample


class AddCorruptedSample:
    def __call__(self, sample):
        x = sample["x"]
        doy = sample["doy"]
        corrupted_x, corrupted_mask = random_masking(x, x.shape[0] - np.sum(doy == 0), x.shape[0], x.shape[1])
        sample["corrupted_x"] = corrupted_x
        sample["corrupted_mask"] = corrupted_mask
        return sample


# ---------------------------- Spectral augmentation ----------------------------
class RandomChanSwapping:
    def __init__(self, probability=0.5):
        self.probability = probability

    def __call__(self, sample):
        if np.random.rand() <= self.probability:
            x = sample["x"]
            s_idx = random.sample(range(x.shape[1]), 2)
            idx = list(range(x.shape[1]))
            idx[s_idx[0]] = s_idx[1]
            idx[s_idx[1]] = s_idx[0]
            x = x[:, idx]
            sample["x"] = x

        return sample


class RandomChanRemoval:
    def __init__(self, probability=0.5):
        self.probability = probability

    def __call__(self, sample):
        if np.random.rand() <= self.probability:
            x = sample["x"]
            s_idx = random.sample(range(x.shape[1]), 2)
            idx = list(range(x.shape[1]))
            idx[s_idx[0]] = s_idx[1]
            x = x[:, idx]
            sample["x"] = x

        return sample


class RandomAddNoise:
    def __init__(self, max_noise=0.05, noise_prob=0.5, probability=0.5):
        self.max_noise = max_noise
        self.noise_prob = noise_prob
        self.probability = probability

    def __call__(self, sample):
        if np.random.rand() <= self.probability:
            x = sample["x"]
            t, c = x.shape

            noise_mask = np.random.rand(t, 1) < self.noise_prob
            noise = np.random.uniform(-self.max_noise, self.max_noise, size=(t, c))
            x = x + noise_mask * noise

            sample["x"] = x
        return sample


# ---------------------------- Temporal augmentation ----------------------------
class RandomTempSwapping:
    def __init__(self, max_distance=-1, probability=0.5):
        self.max_distance = max_distance
        self.probability = probability

    def __call__(self, sample):
        if np.random.rand() <= self.probability:
            x = sample["x"]
            length = x.shape[0]
            idx = list(range(length))

            if self.max_distance == -1:
                s_idx = random.sample(range(length), 2)
            else:
                i = random.randrange(length)  # noqa: S311
                min_j = max(0, i - self.max_distance)
                max_j = min(length - 1, i + self.max_distance)
                possible_js = list(range(min_j, max_j + 1))
                possible_js.remove(i)
                j = random.choice(possible_js)  # noqa: S311
                s_idx = [i, j]

            idx[s_idx[0]], idx[s_idx[1]] = idx[s_idx[1]], idx[s_idx[0]]
            x = x[idx]
            sample["x"] = x

        return sample


class RandomTempShift:
    def __init__(self, max_shift=30, probability=0.5):
        self.max_shift = max_shift
        self.probability = probability

    def __call__(self, sample):
        if np.random.rand() <= self.probability:
            x = sample["x"]
            shift = int(np.clip(np.random.randn() * 0.3, -1, 1) * self.max_shift / 5)
            x = np.roll(x, shift, axis=0)
            sample["x"] = x

        return sample


class RandomTempRemoval:
    def __init__(self, probability=0.5):
        self.probability = probability

    def __call__(self, sample):
        if np.random.rand() <= self.probability:
            x = sample["x"]
            doy = sample["doy"]
            mask = np.array([random.random() >= 0.15 for _ in range(x.shape[0])])  # noqa: S311

            x_kept = x[mask]
            doy_kept = doy[mask]

            num_pad = x.shape[0] - x_kept.shape[0]

            x_pad = np.zeros((num_pad,) + x.shape[1:], dtype=x.dtype)
            doy_pad = np.zeros((num_pad,) + doy.shape[1:], dtype=doy.dtype)

            x_new = np.concatenate((x_kept, x_pad), axis=0)
            doy_new = np.concatenate((doy_kept, doy_pad), axis=0)

            sample["x"] = x_new
            sample["doy"] = doy_new

        return sample


class AddMissingMask:
    def __call__(self, sample):
        doy = sample["doy"]
        sample["mask"] = doy == 0
        return sample


class AddNDVIWeights:
    """
    Compute sample weights based on NDVI from input spectral data.

    This function calculates the Normalized Difference Vegetation Index (NDVI) from the input spectral data,
    applies cloud and dark pixel masking, and computes weights based on the NDVI values. The weights are then
    normalized to sum to 1.
    """

    def __call__(self, sample):
        x = sample["x"]
        all_zero_mask = np.all(x == 0, axis=1)

        score = np.ones(x.shape[0])
        score = np.minimum(score, (x[:, [0, 1, 2]].sum(1) - 0.2) / 0.6)  # rgb
        cloud = score * 100 > 20
        dark = x[:, [6, 8, 9]].sum(1) < 0.35  # NIR, SWIR1, SWIR2

        ndvi = (x[:, 6] - x[:, 2]) / (x[:, 6] + x[:, 2] + 1e-6)
        ndvi[cloud] = -1
        ndvi[dark] = -1
        ndvi = ndvi.clip(-1, 1)

        weight = np.exp(ndvi)
        weight /= weight.sum()

        weight[all_zero_mask] = 0

        sample["weight"] = weight
        return sample


class SetSequenceLength:
    def __init__(self, sequence_length):
        self.sequence_length = sequence_length

    def __call__(self, sample):
        current_length = sample["x"].shape[0]

        if current_length < self.sequence_length:
            new_x = np.zeros((self.sequence_length, sample["x"].shape[1]), dtype=np.half)
            new_x[:current_length] = sample["x"]
            sample["x"] = new_x

            new_doy = np.zeros((self.sequence_length), dtype=np.int16)
            new_doy[:current_length] = sample["doy"]
            sample["doy"] = new_doy
        elif current_length > self.sequence_length:
            sample["x"] = sample["x"][: self.sequence_length]
            sample["doy"] = sample["doy"][: self.sequence_length]

        return sample


class ReduceToMonthlyMeans:
    def __call__(self, sample):
        x, doy = sample["x"], sample["doy"]
        num_features = x.shape[1]

        valid_mask = doy > 0
        x_valid = x[valid_mask]
        doy_valid = doy[valid_mask]

        monthly_means = np.zeros((12, num_features), dtype=x.dtype)
        present_mask = np.zeros(12, dtype=bool)

        if len(doy_valid) > 0:
            month_ends = [31, 59, 90, 120, 151, 181, 212, 243, 273, 304, 334, 366]
            months = np.digitize(doy_valid, month_ends, right=True)

            for m in range(12):
                mask = months == m
                if np.any(mask):
                    monthly_means[m, :] = np.mean(x_valid[mask], axis=0)
                    present_mask[m] = True

        out = np.zeros((num_features, 12), dtype=x.dtype)

        if not np.any(present_mask):
            return out

        valid_indices = np.where(present_mask)[0]
        all_indices = np.arange(12)

        for f in range(num_features):
            y_valid = monthly_means[valid_indices, f]
            out[f, :] = np.interp(all_indices, valid_indices, y_valid)

        return {"x": out.T, "y": sample["y"]}


class Pipeline:
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, sample):
        for transform in self.transforms:
            sample = transform(sample)
        return sample
