from contextlib import AbstractContextManager
from typing import Callable, Iterator

import torch
import torch.distributed as dist

from sarasa.config import Evaluate as EvaluateConfig
from sarasa.metrics import MetricsProcessor
from sarasa.utils import IGNORE_INDEX, world_size


class Evaluator:
    def __init__(
        self,
        config: EvaluateConfig,
        val_loader: Iterator[tuple[dict[str, torch.Tensor], torch.Tensor]],
        amp_context: AbstractContextManager,
        metrics_processor: MetricsProcessor,
        loss_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
        device: torch.device,
    ):
        self.freq = config.freq
        self.val_loader = val_loader
        self.amp_context = amp_context
        self.metrics_processor = metrics_processor
        self.loss_fn = loss_fn
        self.device = device

    def trigger(
        self,
        step: int,
    ) -> bool:
        return step % self.freq == 0

    @torch.no_grad()
    def evaluate(
        self,
        model: torch.nn.Module,
        step: int,
    ) -> None:
        model.eval()

        losses = []
        total_tokens = torch.zeros(1, device=self.device)

        for input_dict, target in self.val_loader:
            self.metrics_processor.ntokens_since_last_log += target.numel()
            input_dict = {
                k: v.to(self.device, non_blocking=(self.device.type == "cuda")) for k, v in input_dict.items()
            }
            target = target.to(self.device, non_blocking=(self.device.type == "cuda"))
            valid_tokens = (target != IGNORE_INDEX).sum()

            with self.amp_context:
                pred = model(**input_dict)
                loss = self.loss_fn(pred, target)

            if world_size() > 1:
                dist.all_reduce(valid_tokens, op=dist.ReduceOp.SUM)
                dist.all_reduce(loss, op=dist.ReduceOp.SUM)

            losses.append(loss)
            total_tokens += valid_tokens

        loss = torch.stack(losses).sum() / total_tokens

        self.metrics_processor.val_log(step=step, val_loss=loss.item())

        model.train()
