20class AIICE:
 21    """
 22    High-level interface for loading Arctic ice data, preparing datasets, and benchmarking models.
 23
 24    This class provides a simple API to:
 25    1. Load historical ice data within a specified date range (see `aiice.loader.Loader`)
 26    2. Convert the data into sliding-window datasets (see `aiice.preprocess.SlidingWindowDataset`)
 27    3. Create a PyTorch DataLoader for batch processing
 28    4. Benchmark any PyTorch model on the OSI-SAF dataset with specified metrics
 29
 30    Args:
 31        pre_history_len (int): Number of past time steps to include in each input sample (X).
 32        forecast_len (int): Number of future time steps to predict (Y) in each sample.
 33        batch_size (int, optional): Batch size for the DataLoader. Defaults to 16.
 34        start (date | str | None, optional): Start date of the data to load. If None, defaults to the earliest available data.
 35        end (date | str | None, optional): End date of the data to load. If None, defaults to the latest available data.
 36        step (int | None, optional): Step in days between data points. Defaults to 1 if not provided.
 37        sea (str, optional): Name of the sea (e.g., "Barents Sea"). Check `Loader.seas` for available ones.
 38        threshold (float | None, optional): Threshold for binarizing the target Y. Values above threshold are set to 1, below or equal set to 0. Defaults to None.
 39        x_binarize (bool, optional): Whether to apply the same threshold binarization to input X. Defaults to False.
 40        threads (int, optional): Number of parallel download threads. You can reduce this value in case of rate limiting HuggingFace API errors. Defaults to 16.
 41        device (str | None, optional): Device to place tensors on ("cpu", "cuda", etc.). If None, uses PyTorch default device.
 42
 43    Example:
 44        >>> aiice = AIICE(pre_history_len=30, forecast_len=7, batch_size=32, start="2022-01-01", end="2022-12-31")
 45        >>> model = MyModel()
 46        >>> results = aiice.bench(model, metrics={"mae", "psnr"})
 47    """
 48
 49    def __init__(
 50        self,
 51        pre_history_len: int,
 52        forecast_len: int,
 53        batch_size: int = 16,
 54        start: date | str | None = None,
 55        end: date | str | None = None,
 56        step: int | None = None,
 57        sea: str | None = None,
 58        threshold: float | None = None,
 59        x_binarize: bool = False,
 60        threads: int = 16,
 61        device: str | None = None,
 62    ):
 63        self._device = device
 64        self._sea = sea
 65
 66        raw_data = Loader().get(
 67            start=start,
 68            end=end,
 69            step=step,
 70            sea=sea,
 71            threads=threads,
 72            tensor_out=True,
 73            idx_out=True,
 74        )
 75
 76        indices = raw_data[0]
 77        matrices = raw_data[1]
 78
 79        dataset = SlidingWindowDataset(
 80            data=matrices,
 81            idx=indices,
 82            pre_history_len=pre_history_len,
 83            forecast_len=forecast_len,
 84            threshold=threshold,
 85            x_binarize=x_binarize,
 86            device=self._device,
 87        )
 88
 89        self._dataloader = DataLoader(
 90            dataset=dataset,
 91            batch_size=batch_size,
 92            collate_fn=self._default_collate_fn,
 93        )
 94
 95    def bench(
 96        self,
 97        model: nn.Module,
 98        metrics: dict[str, MetricFn] | list[str] | None = None,
 99        path: str | None = None,
100        detailed: bool = True,
101        plot_workers: int = 4,
102        fps: int = 2,
103    ) -> dict[str, list[float]]:
104        """
105        Run benchmarking evaluation of a model on the prepared dataset.
106
107        The method iterates over the internal DataLoader, generates model
108        predictions, computes evaluation metrics, and optionally produces
109        visualization GIFs comparing ground truth and predicted forecasts.
110
111        When `path` is provided, visualization generation is executed
112        asynchronously using a thread pool so that plotting does not block
113        model inference.
114
115        Args:
116            model (nn.Module):
117                PyTorch model used to generate predictions. The model is expected
118                to accept inputs `x` with shape `(batch, pre_history_len, ...)`
119                and return predictions compatible with the selected metrics.
120
121            metrics (dict[str, MetricFn] | list[str], optional):
122                Metrics to compute during evaluation. If a list of metric names is
123                provided, the metrics are resolved from the built-in registry.
124                If `None`, default metrics are used.
125                See `aiice.metrics.Evaluator` for details.
126
127            path (str, optional):
128                Directory where forecast visualizations will be saved.
129                If provided, each sample in the dataset will produce a GIF
130                animation showing the forecast horizon, comparing ground truth
131                and model predictions frame by frame.
132
133                The files are named: `<start_forecast_date>_<end_forecast_date>.gif`
134                If `None`, visualization generation is skipped.
135
136            detailed (bool, optional):
137                If True, returns full statistics for each metric like
138                mean, last value, count, min, and max.
139                If False, returns only the mean value per metric.
140
141            plot_workers (int, optional):
142                Number of worker threads used for asynchronous plot generation.
143                Increasing this value can speed up visualization when many samples
144                are processed. Defaults to 4.
145
146            fps (int, optional):
147                Frames per second of the generated GIF animations. Defaults to 2.
148
149        Returns:
150            dict[str, list[float]]:
151                Aggregated metric results returned by the evaluator.
152        """
153        if path is not None:
154            os.makedirs(path, exist_ok=True)
155            executor = ThreadPoolExecutor(max_workers=plot_workers)
156            futures = []
157
158        evaluator = Evaluator(metrics=metrics, accumulate=True)
159
160        model.eval()
161        with torch.no_grad():
162            for batch in tqdm(self._dataloader, desc="Prediction"):
163                dates, x, y = batch
164                x, y = x.to(self._device), y.to(self._device)
165
166                pred = model(x)
167                evaluator.eval(y, pred)
168
169                if path is None:
170                    continue
171
172                futures.append(
173                    executor.submit(
174                        self._save_batch_plot,
175                        sea=self._sea,
176                        path=path,
177                        dates=dates,
178                        y=y.detach().cpu().numpy(),
179                        pred=pred.detach().cpu().numpy(),
180                        fps=fps,
181                    )
182                )
183
184        if path is not None:
185            for f in tqdm(futures, desc="Saving plots"):
186                f.result()
187            executor.shutdown(wait=True)
188
189        return evaluator.report(detailed=detailed)
190
191    @staticmethod
192    def _save_batch_plot(
193        sea: str | None,
194        path: str,
195        dates: list[list[date]],
196        y: np.ndarray,
197        pred: np.ndarray,
198        fps: int,
199    ) -> None:
200        """
201        Generate GIF visualizations for a batch of forecast samples.
202
203        For each sample in the batch, a GIF animation is created showing
204        the temporal evolution of the forecast horizon. Each frame displays
205        a side-by-side comparison between the ground truth ice map and the
206        model prediction for the corresponding forecast date.
207
208        The resulting GIF file is saved to `path` with the name: `<start_forecast_date>_<end_forecast_date>.gif`
209        where the dates correspond to the forecast window of the sample.
210        """
211        matplotlib.use("Agg")
212
213        batch_size, forecast_len = y.shape[:2]
214        for i in range(batch_size):
215
216            start_date = dates[i][-forecast_len].strftime("%d-%m-%Y")
217            end_date = dates[i][-1].strftime("%d-%m-%Y")
218
219            save_path = os.path.join(path, f"{start_date}_{end_date}.gif")
220            fig, axes = plt.subplots(1, 2, figsize=(8, 4))
221
222            im_gt = axes[0].imshow(y[i, 0])
223            axes[0].set_title("Ground Truth")
224            axes[0].axis("off")
225
226            im_pred = axes[1].imshow(pred[i, 0])
227            axes[1].set_title("Prediction")
228            axes[1].axis("off")
229
230            frames = []
231            for j in range(forecast_len):
232
233                im_gt.set_data(y[i, j])
234                im_pred.set_data(pred[i, j])
235
236                forecast_date = dates[i][-forecast_len + j]
237                if sea is None:
238                    fig.suptitle(f"Forecast: {forecast_date.strftime('%d-%m-%Y')}")
239                else:
240                    fig.suptitle(
241                        f"{sea} | Forecast: {forecast_date.strftime('%d-%m-%Y')}"
242                    )
243
244                fig.canvas.draw()
245                frame = np.asarray(fig.canvas.buffer_rgba())[:, :, :3].copy()
246                frames.append(frame)
247
248            plt.close(fig)
249            imageio.mimsave(save_path, frames, duration=1 / fps, loop=0)
250
251    @staticmethod
252    def _default_collate_fn(
253        batch: list[tuple[list[date], torch.Tensor, torch.Tensor]],
254    ) -> tuple[list[list[date]], torch.Tensor, torch.Tensor]:
255        """
256        Collates SlidingWindow dataset samples into a batch
257        input  -> batch of samples
258        output -> batched tensors + list of date sequences
259
260        Example:
261        ```
262        d1 = [date1...date2]
263        x1.shape = (T, H, W)
264        y1.shape = (H, W)
265
266        batch = [
267            (d1, x1, y1),
268            (d2, x2, y2)
269        ]
270
271        Output:
272            dates -> [d1, d2]
273            x     -> torch.Tensor (B, T, H, W)
274            y     -> torch.Tensor (B, H, W)
275        ```
276        """
277        dates, x, y = zip(*batch)
278        return list(dates), torch.stack(x), torch.stack(y)