#!/usr/bin/env python3
"""Tests for active_learning.py (numpy + Pillow only). Run: python3 test_active_learning.py"""
import csv
import shutil
import tempfile
from pathlib import Path

import numpy as np
from PIL import Image

import active_learning as al
import flywheel as fw


def _img(path, seed):
    rng = np.random.default_rng(seed)
    Image.fromarray(rng.integers(0, 255, (32, 32, 3), dtype=np.uint8)).save(path)


def test_novelty_score():
    value, distance = al.novelty_score("00ff00ff00ff00ff", [])
    assert value == 1.0 and distance is None
    same = al.novelty_score("00ff00ff00ff00ff", ["00ff00ff00ff00ff"])
    assert same[0] == 0.0 and same[1] == 0
    different = al.novelty_score("0000000000000000", ["ffffffffffffffff"])
    assert different[0] == 1.0 and different[1] == 64


def test_domain_score_prefers_rare_values():
    counts = {"environment": {"indoor": 10, "outdoor": 1}}
    common, _ = al.domain_score({"environment": "indoor"}, counts)
    rare, reasons = al.domain_score({"environment": "outdoor"}, counts)
    assert rare > common
    assert any("rare environment=outdoor" in reason for reason in reasons)


def test_ranking_puts_novel_rare_above_duplicate():
    with tempfile.TemporaryDirectory() as folder:
        root = Path(folder)
        data = root / "data"
        (data / "train" / "images").mkdir(parents=True)
        labelled = data / "train" / "images" / "existing.jpg"
        _img(labelled, 1)
        with (data / "metadata.csv").open("w", newline="") as stream:
            writer = csv.DictWriter(stream, fieldnames=["split", "image", "wall_id", "session_id", "environment"])
            writer.writeheader()
            for index in range(10):
                writer.writerow({"split": "train", "image": f"i{index}.jpg", "wall_id": f"w{index}",
                                 "session_id": "s", "environment": "indoor"})
            writer.writerow({"split": "train", "image": "existing.jpg", "wall_id": "w", "session_id": "s",
                             "environment": "outdoor"})

        candidates = root / "candidates"
        candidates.mkdir()
        shutil.copy(labelled, candidates / "duplicate.jpg")
        _img(candidates / "novel.jpg", 99)

        index = fw.build_index(data)
        counts = al.labelled_counts(data)
        rows = {"novel.jpg": {"environment": "outdoor"}}
        ranked = al.score_candidates([candidates / "duplicate.jpg", candidates / "novel.jpg"],
                                     index, counts, rows=rows)
        by_name = {Path(item["image"]).name: item for item in ranked}
        assert by_name["novel.jpg"]["score"] > by_name["duplicate.jpg"]["score"]
        assert by_name["duplicate.jpg"]["novelty"] == 0.0
        assert any("rare environment=outdoor" in r for r in by_name["novel.jpg"]["reasons"])


def test_uncertainty_enters_the_score():
    with tempfile.TemporaryDirectory() as folder:
        root = Path(folder)
        data = root / "data"
        (data / "train" / "images").mkdir(parents=True)
        _img(data / "train" / "images" / "existing.jpg", 5)
        candidates = root / "candidates"
        candidates.mkdir()
        _img(candidates / "a.jpg", 2)
        _img(candidates / "b.jpg", 3)
        index = fw.build_index(data)
        counts = al.labelled_counts(data)
        uncertain = {str(candidates / "b.jpg"): 0.9}
        ranked = al.score_candidates([candidates / "a.jpg", candidates / "b.jpg"], index, counts,
                                     uncertainties=uncertain)
        by_name = {Path(item["image"]).name: item for item in ranked}
        assert by_name["b.jpg"]["uncertainty"] == 0.9
        assert any("high model uncertainty" in r for r in by_name["b.jpg"]["reasons"])


def main():
    tests = [value for name, value in sorted(globals().items()) if name.startswith("test_")]
    for test in tests:
        test()
        print(f"ok  {test.__name__}")
    print(f"{len(tests)} tests passed")


if __name__ == "__main__":
    main()