#!/usr/bin/env python3
# (c) Facebook, Inc. and its affiliates. Confidential and proprietary.

from typing import Any

import torch
from flsim.interfaces.batch_metrics import IFLBatchMetrics


class FLBatchMetrics(IFLBatchMetrics):
    def __init__(
        self,
        *,
        loss: torch.Tensor,
        num_examples: int,
        predictions: torch.Tensor,
        targets: torch.Tensor,
        model_inputs: Any,
    ) -> None:
        self._loss = loss
        self._num_examples = num_examples
        self._predictions = predictions
        self._targets = targets
        self._model_inputs = model_inputs

    @property
    def loss(self) -> torch.Tensor:
        return self._loss

    @property
    def num_examples(self) -> int:
        return self._num_examples

    @property
    def predictions(self) -> torch.Tensor:
        return self._predictions

    @property
    def targets(self) -> torch.Tensor:
        return self._targets

    @property
    def model_inputs(self) -> Any:
        return self._model_inputs
