from __future__ import annotations

import random
import unittest

import numpy as np
import torch
from PIL import Image

from inpainting.masks import MASK_KINDS, apply_mask, mask_area_fraction, random_mask
from inpainting.objective import min_snr_weights, spatial_loss_weights, weighted_mse


class MaskTests(unittest.TestCase):
    def test_every_mask_family_is_binary_and_nonempty(self) -> None:
        for index, kind in enumerate(MASK_KINDS):
            mask = random_mask(
                (256, 256),
                random.Random(1000 + index),
                min_area=0.03,
                max_area=0.70,
                kind=kind,
            )
            values = set(np.unique(np.asarray(mask)).tolist())
            self.assertTrue(values <= {0, 255}, (kind, values))
            area = mask_area_fraction(mask)
            self.assertGreater(area, 0.005, kind)
            self.assertLess(area, 0.90, kind)

    def test_apply_mask_only_blacks_white_region(self) -> None:
        image = Image.new("RGB", (8, 8), (20, 40, 60))
        mask = Image.new("L", (8, 8), 0)
        for y in range(2, 6):
            for x in range(3, 7):
                mask.putpixel((x, y), 255)
        result = np.asarray(apply_mask(image, mask))
        self.assertTrue(np.all(result[2:6, 3:7] == 0))
        self.assertTrue(np.all(result[0, 0] == (20, 40, 60)))


class ObjectiveTests(unittest.TestCase):
    def test_spatial_weights_prioritize_mask_and_boundary(self) -> None:
        mask = torch.zeros(1, 1, 16, 16)
        mask[:, :, 5:11, 5:11] = 1
        weights = spatial_loss_weights(
            mask,
            context_weight=0.25,
            masked_weight=2.5,
            boundary_weight=2.0,
            boundary_radius=1,
        )
        self.assertAlmostEqual(weights.mean().item(), 1.0, places=5)
        self.assertGreater(weights[0, 0, 5, 5], weights[0, 0, 8, 8])
        self.assertGreater(weights[0, 0, 8, 8], weights[0, 0, 0, 0])

    def test_min_snr_and_weighted_mse_are_finite(self) -> None:
        alphas = torch.linspace(0.999, 0.001, 1000)
        timesteps = torch.tensor([0, 250, 999])
        sample_weights = min_snr_weights(
            alphas,
            timesteps,
            gamma=5.0,
            prediction_type="epsilon",
        )
        prediction = torch.ones(3, 4, 8, 8)
        target = torch.zeros_like(prediction)
        spatial = torch.ones(3, 1, 8, 8)
        loss = weighted_mse(
            prediction,
            target,
            spatial_weights=spatial,
            sample_weights=sample_weights,
        )
        self.assertTrue(torch.isfinite(loss))
        self.assertGreater(loss.item(), 0)


if __name__ == "__main__":
    unittest.main()