import hashlib
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,
    })

    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,
        class_names=None,
    ):
        if not _has_torch and pytorch:
            raise ImportError("AgriGEELiteDataset requer 'torch' instalado. Instale com 'pip install torch'.")  # noqa: TRY003

        super().__init__()

        self.normalize = normalize
        self.pytorch = pytorch
        self.class_names = class_names

        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

        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")

        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")

        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()

        # Map class names to indices if class_names is provided
        if self.class_names is not None:
            unique_classes = set(self.gdf["crop_class"].unique())
            class_names_set = set(self.class_names)

            # Check if all classes in the dataset are in the provided class_names list
            missing_classes = unique_classes - class_names_set
            if missing_classes:
                raise ValueError(  # noqa: TRY003
                    f"Classes encontradas no dataset que não estão na lista class_names: {missing_classes}"
                )

            # Create mapping from class name to index
            class_to_idx = {class_name: idx for idx, class_name in enumerate(self.class_names)}

            # Map crop_class to indices
            self.ys = self.gdf["crop_class"].map(class_to_idx).values
            self.num_classes = len(self.class_names)
        else:
            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.float16) 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

    def get_evi2_and_doys(self, idx: int, smooth: bool = True):
        x = self.xs[idx]
        doys = self.doys[idx]

        valid_mask = doys != 0

        x = x[valid_mask]
        doys = doys[valid_mask]

        if self.pytorch:
            x = x.numpy()
            doys = doys.numpy()

        nir_idx = self.bands.index("nir")
        red_idx = self.bands.index("red")

        nir = x[:, nir_idx]
        red = x[:, red_idx]

        evi2 = 2.5 * (nir - red) / (nir + 2.4 * red + 1)

        if smooth:
            from scipy.signal import savgol_filter

            if np.sum(valid_mask) >= 5:
                evi2 = savgol_filter(evi2, window_length=5, polyorder=2)

        return doys, evi2

    def __add__(self, other):
        if not isinstance(other, AgriGEELiteDataset):
            raise TypeError(f"Não é possível somar AgriGEELiteDataset com {type(other)}")  # noqa: TRY003

        result = AgriGEELiteDataset.__new__(AgriGEELiteDataset)

        result.normalize = self.normalize
        result.pytorch = self.pytorch
        result.transform = self.transform
        result.bands = self.bands.copy()
        result.timestamp_processing = self.timestamp_processing
        result.class_names = self.class_names

        if self.normalize:
            result.bands_means = self.bands_means.copy()
            result.bands_stds = self.bands_stds.copy()

        max_indexnum_self = self.gdf["indexnum"].max()
        offset = max_indexnum_self + 1

        gdf_other = other.gdf.copy()
        gdf_other["indexnum"] = gdf_other["indexnum"] + offset
        result.gdf = pd.concat([self.gdf, gdf_other], ignore_index=True)

        sits_df_other = other.sits_df.copy()
        sits_df_other["indexnum"] = sits_df_other["indexnum"] + offset
        result.sits_df = pd.concat([self.sits_df, sits_df_other], ignore_index=True)

        result.max_seq_len = max(self.max_seq_len, other.max_seq_len)

        xs_self = self.xs
        doys_self = self.doys
        xs_other = other.xs
        doys_other = other.doys

        if self.max_seq_len < result.max_seq_len:
            pad_len = result.max_seq_len - self.max_seq_len
            xs_self = np.pad(xs_self, ((0, 0), (0, pad_len), (0, 0)), mode="constant", constant_values=0)
            doys_self = np.pad(doys_self, ((0, 0), (0, pad_len)), mode="constant", constant_values=0)

        if other.max_seq_len < result.max_seq_len:
            pad_len = result.max_seq_len - other.max_seq_len
            xs_other = np.pad(xs_other, ((0, 0), (0, pad_len), (0, 0)), mode="constant", constant_values=0)
            doys_other = np.pad(doys_other, ((0, 0), (0, pad_len)), mode="constant", constant_values=0)

        result.xs = np.concatenate([xs_self, xs_other], axis=0)
        result.doys = np.concatenate([doys_self, doys_other], axis=0)
        result.ys = np.concatenate([self.ys, other.ys], axis=0)

        result.sequence_length = result.max_seq_len
        result.num_classes = int(result.gdf["crop_class"].nunique())

        return result

    def get_weights_for_WeightedRandomSampler(self) -> torch.Tensor:
        targets = self.ys
        classes, counts = np.unique(targets, return_counts=True)
        class_weights = 1.0 / counts
        weight_map = dict(zip(classes, class_weights))
        sample_weights = np.array([weight_map[t] for t in targets])
        return torch.from_numpy(sample_weights).double()

    def merge(self, other):
        """
        Merge this dataset with another dataset, combining their timestamps.

        Only parcels with identical geometry, start_date and end_date are merged.
        For each band:
        - Takes data from the other dataset for timestamps that don't exist in this dataset
        - For bands unique to one dataset, interpolates values for the other
        - Keeps means and stds from this dataset, using other dataset only for new bands

        Parameters
        ----------
        other : AgriGEELiteDataset
            The dataset to merge with this one

        Returns
        -------
        AgriGEELiteDataset
            A new merged dataset with only identical parcels
        """
        if not isinstance(other, AgriGEELiteDataset):
            raise TypeError(f"Cannot merge AgriGEELiteDataset with {type(other)}")  # noqa: TRY003

        # Create hash column for both datasets based on geometry, start_date and end_date
        def create_parcel_hash(row):
            """Create a hash from geometry, start_date and end_date."""
            geom_wkb = row.geometry.wkb
            start_str = str(row.get("start_date", ""))
            end_str = str(row.get("end_date", ""))
            combined = geom_wkb + start_str.encode() + end_str.encode()
            return hashlib.sha256(combined).hexdigest()

        gdf_self = self.gdf.copy()
        gdf_other = other.gdf.copy()

        gdf_self["parcel_hash"] = gdf_self.apply(create_parcel_hash, axis=1)
        gdf_other["parcel_hash"] = gdf_other.apply(create_parcel_hash, axis=1)

        # Find common parcels (identical geometry, start_date, end_date)
        common_hashes = set(gdf_self["parcel_hash"]) & set(gdf_other["parcel_hash"])

        if not common_hashes:
            raise ValueError("No identical parcels found between datasets")  # noqa: TRY003

        print(f"Found {len(common_hashes)} identical parcels between datasets")

        # Filter both datasets to keep only common parcels
        gdf_self_filtered = gdf_self[gdf_self["parcel_hash"].isin(common_hashes)].copy()
        gdf_other_filtered = gdf_other[gdf_other["parcel_hash"].isin(common_hashes)].copy()

        # Create mapping from parcel_hash to indexnum for both datasets
        hash_to_indexnum_self = dict(zip(gdf_self_filtered["parcel_hash"], gdf_self_filtered["indexnum"]))
        hash_to_indexnum_other = dict(zip(gdf_other_filtered["parcel_hash"], gdf_other_filtered["indexnum"]))

        # Create result dataset
        result = AgriGEELiteDataset.__new__(AgriGEELiteDataset)
        result.normalize = self.normalize
        result.pytorch = self.pytorch
        result.transform = self.transform
        result.timestamp_processing = self.timestamp_processing
        result.class_names = self.class_names

        # Determine union of bands
        bands_self = set(self.bands)
        bands_other = set(other.bands)
        all_bands = list(dict.fromkeys(self.bands + [b for b in other.bands if b not in bands_self]))
        result.bands = all_bands

        # Setup means and stds for normalization
        if self.normalize:
            result.bands_means = np.zeros(len(all_bands), dtype=np.float16)
            result.bands_stds = np.zeros(len(all_bands), dtype=np.float16)

            for i, band in enumerate(all_bands):
                if band in bands_self:
                    # Use self's stats
                    idx_self = self.bands.index(band)
                    result.bands_means[i] = self.bands_means[idx_self]
                    result.bands_stds[i] = self.bands_stds[idx_self]
                elif band in bands_other:
                    # Use other's stats for bands not in self
                    idx_other = other.bands.index(band)
                    result.bands_means[i] = other.bands_means[idx_other]
                    result.bands_stds[i] = other.bands_stds[idx_other]

        # Create result gdf with filtered parcels and reset indexnum
        result.gdf = gdf_self_filtered.drop(columns=["parcel_hash"]).reset_index(drop=True)
        unique_hashes = sorted(common_hashes)
        new_indexnum_mapping = {hash_val: idx for idx, hash_val in enumerate(unique_hashes)}
        result.gdf["indexnum"] = gdf_self_filtered["parcel_hash"].map(new_indexnum_mapping).values
        result.num_classes = int(result.gdf["crop_class"].nunique())

        # Merge sits_df
        merged_sits_list = []

        for parcel_hash in common_hashes:
            indexnum_self = hash_to_indexnum_self[parcel_hash]
            indexnum_other = hash_to_indexnum_other[parcel_hash]
            new_indexnum = new_indexnum_mapping[parcel_hash]

            sits_self = self.sits_df[self.sits_df["indexnum"] == indexnum_self].copy()
            sits_other = other.sits_df[other.sits_df["indexnum"] == indexnum_other].copy()

            # Get timestamps from both datasets
            timestamps_self = set(sits_self["timestamp"].values)
            timestamps_other = set(sits_other["timestamp"].values)

            # For timestamps unique to other, add them
            new_timestamps = timestamps_other - timestamps_self
            new_rows = sits_other[sits_other["timestamp"].isin(new_timestamps)].copy()

            # For bands in self but not in other, interpolate
            for band in bands_self - bands_other:
                if len(new_rows) > 0:
                    # Get existing values for this band
                    all_times = pd.concat([
                        sits_self[["timestamp", band]],
                        pd.DataFrame({"timestamp": new_rows["timestamp"], band: 0}),
                    ])
                    all_times = all_times.sort_values("timestamp")

                    # Interpolate for new timestamps
                    if len(sits_self) > 0:
                        interpolated = np.interp(
                            new_rows["timestamp"].values, sits_self["timestamp"].values, sits_self[band].values
                        )
                        new_rows[band] = interpolated
                    else:
                        new_rows[band] = 0

            # For bands in other but not in self, interpolate for existing timestamps
            for band in bands_other - bands_self:
                if len(sits_self) > 0:
                    if len(sits_other) > 0:
                        interpolated = np.interp(
                            sits_self["timestamp"].values, sits_other["timestamp"].values, sits_other[band].values
                        )
                        sits_self[band] = interpolated
                    else:
                        sits_self[band] = 0

                # New rows already have this band from other dataset

            # Combine all rows
            combined = pd.concat([sits_self, new_rows], ignore_index=True)

            # Ensure all bands are present
            for band in all_bands:
                if band not in combined.columns:
                    combined[band] = 0

            # Set the new indexnum
            combined["indexnum"] = new_indexnum

            merged_sits_list.append(combined)

        result.sits_df = pd.concat(merged_sits_list, ignore_index=True)
        result.sits_df = result.sits_df.sort_values(["indexnum", "timestamp"]).reset_index(drop=True)

        # Recompute max_seq_len
        result.max_seq_len = result.sits_df.groupby("indexnum").size().max()
        result.sequence_length = result.max_seq_len

        # Rebuild numpy arrays
        result.xs, result.doys = result.to_numpy_arrays_wo_for()
        result.ys = result.gdf["crop_class"].values

        return result


# ------------------ 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
