from __future__ import annotations

import importlib.util
from pathlib import Path

import pytest
import torch


ROOT = Path(__file__).resolve().parents[2]
BATCH_FILE = ROOT / "nodes" / "core" / "batch.py"

spec = importlib.util.spec_from_file_location(
    "imageops_batch_test_module",
    BATCH_FILE,
)

assert spec is not None
assert spec.loader is not None

batch_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(batch_module)


match_batch = batch_module.match_batch
match_mask_batch = batch_module.match_mask_batch
BatchMismatchError = batch_module.BatchMismatchError
cycle_to_length = batch_module.cycle_to_length


def test_equal_batches_unchanged():
    a = torch.zeros(4, 8, 8, 3)
    b = torch.ones(4, 8, 8, 3)

    aa, bb = match_batch(a, b)

    assert aa.shape == a.shape
    assert bb.shape == b.shape


def test_singleton_broadcast():
    a = torch.zeros(1, 8, 8, 3)
    b = torch.ones(4, 8, 8, 3)

    aa, bb = match_batch(a, b)

    assert aa.shape[0] == 4
    assert bb.shape[0] == 4


def test_strict_mismatch_raises():
    a = torch.zeros(3, 8, 8, 3)
    b = torch.ones(7, 8, 8, 3)

    with pytest.raises(BatchMismatchError):
        match_batch(a, b)


def test_hold_last():
    a = torch.arange(3).view(3, 1)
    b = torch.arange(5).view(5, 1)

    aa, bb = match_batch(
        a,
        b,
        policy="hold_last",
    )

    assert aa[:, 0].tolist() == [0, 1, 2, 2, 2]
    assert bb[:, 0].tolist() == [0, 1, 2, 3, 4]


def test_loop_is_explicit():
    a = torch.arange(3).view(3, 1)
    b = torch.arange(7).view(7, 1)

    aa, _ = match_batch(
        a,
        b,
        policy="loop",
    )

    assert aa[:, 0].tolist() == [
        0, 1, 2,
        0, 1, 2,
        0,
    ]


def test_mask_singleton_broadcast():
    mask = torch.ones(1, 16, 16)

    result = match_mask_batch(
        mask,
        target_batch=4,
    )

    assert result.shape == (4, 16, 16)


def test_mask_mismatch_raises():
    mask = torch.ones(3, 16, 16)

    with pytest.raises(BatchMismatchError):
        match_mask_batch(
            mask,
            target_batch=7,
        )


def test_cycle_to_length_exact_multiple():
    t = torch.arange(3).view(3, 1)

    out = cycle_to_length(t, 9)

    assert out[:, 0].tolist() == [0, 1, 2] * 3


def test_cycle_to_length_truncates_final_repeat():
    t = torch.arange(3).view(3, 1)

    out = cycle_to_length(t, 7)

    assert out[:, 0].tolist() == [0, 1, 2, 0, 1, 2, 0]


def test_cycle_to_length_noop_when_already_target():
    t = torch.arange(4).view(4, 1)

    out = cycle_to_length(t, 4)

    assert out is t