"""FID evaluation helper.

Computes Fréchet Inception Distance between a fixed set of real images
and a batch of generated images. Real images are stored as a tensor on CPU
and moved to device only during evaluation — this avoids re-reading disk
every call while keeping GPU memory free between evaluations.
"""
import torch
from torch.utils.data import DataLoader
from torchmetrics.image.fid import FrechetInceptionDistance


class FIDEvaluator:
    def __init__(self, real_dataset, n_real: int = 10_000, device: str = "cuda"):
        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
        self.n_real = n_real

        # Cache real images as a CPU tensor ([-1, 1] range)
        imgs_list = []
        loader = DataLoader(real_dataset, batch_size=256, shuffle=False,
                            num_workers=4, drop_last=False)
        for batch in loader:
            imgs_list.append(batch.cpu())
            if sum(x.size(0) for x in imgs_list) >= n_real:
                break
        real = torch.cat(imgs_list)[:n_real]
        self._real = real  # stored on CPU, shape (N, 3, H, W) in [-1, 1]

    @torch.no_grad()
    def compute(self, fake_imgs: torch.Tensor) -> float:
        """Compute FID score.

        fake_imgs: float tensor in [-1, 1], shape (N, 3, H, W).
                   N should be at least 2048 for a reliable score.
        """
        fid = FrechetInceptionDistance(feature=2048, normalize=True).to(self.device)

        # Feed real images in batches
        for i in range(0, self._real.size(0), 256):
            batch = (self._real[i:i + 256] * 0.5 + 0.5).clamp(0, 1).to(self.device)
            fid.update(batch, real=True)

        # Feed fake images in batches
        fake = fake_imgs.cpu()
        for i in range(0, fake.size(0), 256):
            batch = (fake[i:i + 256] * 0.5 + 0.5).clamp(0, 1).to(self.device)
            fid.update(batch, real=False)

        return float(fid.compute())
