from datetime import datetime
import pytest
from pathlib import Path
from tempfile import TemporaryDirectory
import numpy as np
from PIL import Image
import torch
import torchvision

from . import check_bad_imports
from mmp.a1 import main, tensors

current_assignment = pytest.mark.skipif(
    not (datetime.now() <= datetime(2025, 10, 22, 23, 59, 59)),
    reason="This is not the current assignment.",
)


@current_assignment
def test_no_abs_import():
    paths = list(Path().glob("mmp/a1/*.py"))
    check_bad_imports(paths)    


@current_assignment
def test_main():
    # for testing: generate random image
    img1 = Image.fromarray((np.random.rand(128, 128, 3) * 255).astype(np.uint8))
    img2 = Image.fromarray((np.random.rand(128, 128, 3) * 255).astype(np.uint8))
    tfm = torchvision.transforms.Resize(320)
    with TemporaryDirectory() as img_dir:
        img_dir = Path(img_dir)
        p1 = str(img_dir / "img1.jpg")
        p2 = str(img_dir / "img2.jpg")
        img1.save(p1)
        img2.save(p2)
        batch = main.build_batch([p1, p2])
        batch_transformed = main.build_batch([p1, p2], transform=tfm)
    assert isinstance(batch, torch.Tensor)
    assert len(batch) == 2

    model = main.get_model()
    assert isinstance(model, torch.nn.Module)

    main.main


@current_assignment
def test_tensors():
    avg = tensors.avg_color(torch.rand(3, 128, 128))
    assert isinstance(avg, torch.Tensor)
    assert len(avg) == 3

    masked = tensors.mask(
        torch.rand(3, 128, 128), torch.rand(3, 128, 128), torch.rand(128, 128), 0.3
    )
    assert isinstance(masked, torch.Tensor)
    assert masked.shape == (3, 128, 128)

    result = tensors.add_matrix_vector(torch.rand(3, 4), torch.rand(3))
    assert isinstance(result, torch.Tensor)
    assert result.shape == (3, 4)