"""
Tests for small but training-critical pure functions that had no coverage:
  - train.get_lr          : warmup + cosine-decay learning-rate schedule
  - data.data_utils._is_batch_valid / synchronized_dataloader_step (single-GPU path)
"""
import unittest

from train import get_lr
from data.data_utils import _is_batch_valid, synchronized_dataloader_step

import torch


class TestGetLr(unittest.TestCase):
    MAX_LR = 1.0
    MAX_STEPS = 1000          # warmup_steps = 0.03 * 1000 = 30
    WARMUP = 30
    MIN_LR = 0.1              # 0.1 * MAX_LR

    def test_warmup_starts_near_zero(self):
        # it=0 -> max_lr * 1/warmup_steps, small and positive (not zero, not max).
        lr0 = get_lr(0, self.MAX_LR, self.MAX_STEPS)
        self.assertAlmostEqual(lr0, self.MAX_LR / self.WARMUP, places=6)
        self.assertGreater(lr0, 0.0)

    def test_warmup_reaches_peak_at_end(self):
        lr = get_lr(self.WARMUP - 1, self.MAX_LR, self.MAX_STEPS)
        self.assertAlmostEqual(lr, self.MAX_LR, places=6)

    def test_warmup_is_monotonic_increasing(self):
        lrs = [get_lr(it, self.MAX_LR, self.MAX_STEPS) for it in range(self.WARMUP)]
        for a, b in zip(lrs, lrs[1:]):
            self.assertLess(a, b)

    def test_floor_after_max_steps(self):
        self.assertAlmostEqual(get_lr(self.MAX_STEPS + 1, self.MAX_LR, self.MAX_STEPS),
                               self.MIN_LR, places=6)

    def test_decay_is_monotonic_non_increasing(self):
        lrs = [get_lr(it, self.MAX_LR, self.MAX_STEPS)
               for it in range(self.WARMUP, self.MAX_STEPS + 1)]
        for a, b in zip(lrs, lrs[1:]):
            self.assertLessEqual(b, a + 1e-9)

    def test_cosine_midpoint_is_halfway(self):
        # At decay_ratio = 0.5 the cosine coeff is 0.5 -> lr = min + 0.5*(max-min).
        mid_it = int(self.WARMUP + (self.MAX_STEPS - self.WARMUP) / 2)
        expected = self.MIN_LR + 0.5 * (self.MAX_LR - self.MIN_LR)
        self.assertAlmostEqual(get_lr(mid_it, self.MAX_LR, self.MAX_STEPS), expected, places=2)

    def test_never_below_floor_after_warmup(self):
        # The floor (0.1 * max_lr) only applies once warmup is done; during warmup the
        # schedule intentionally ramps up from ~0, so it is below the floor early on.
        for it in range(self.WARMUP, self.MAX_STEPS + 1, 17):
            self.assertGreaterEqual(get_lr(it, self.MAX_LR, self.MAX_STEPS), self.MIN_LR - 1e-9)


def _valid_batch():
    return {
        "input_ids": torch.zeros(2, 3, dtype=torch.long),  # len() == batch dim == 2
        "labels": torch.zeros(2, 3, dtype=torch.long),
        "attention_mask": torch.ones(2, 3, dtype=torch.long),
        "images": [[torch.zeros(3, 4)], [torch.zeros(3, 4)]],
        "model_patch_positions": [[torch.zeros(3, 2)], [torch.zeros(3, 2)]],
    }


class TestIsBatchValid(unittest.TestCase):
    def test_none_or_empty_is_invalid(self):
        self.assertFalse(_is_batch_valid(None))
        self.assertFalse(_is_batch_valid({}))

    def test_empty_input_ids_is_invalid(self):
        b = _valid_batch(); b["input_ids"] = []
        self.assertFalse(_is_batch_valid(b))

    def test_empty_images_is_invalid(self):
        b = _valid_batch(); b["images"] = []
        self.assertFalse(_is_batch_valid(b))

    def test_images_with_no_actual_image_is_invalid(self):
        # images present as nested lists but all empty -> would deadlock DDP, must be rejected.
        b = _valid_batch(); b["images"] = [[], []]
        self.assertFalse(_is_batch_valid(b))

    def test_well_formed_batch_is_valid(self):
        self.assertTrue(_is_batch_valid(_valid_batch()))


class TestSynchronizedDataloaderStepSingleGPU(unittest.TestCase):
    def test_filters_invalid_batches_when_not_distributed(self):
        good1 = _valid_batch()
        bad = _valid_batch(); bad["images"] = [[], []]
        good2 = _valid_batch()
        loader = [good1, bad, good2]
        out = list(synchronized_dataloader_step(loader, is_dist=False))
        self.assertEqual(len(out), 2)
        self.assertIs(out[0], good1)
        self.assertIs(out[1], good2)

    def test_empty_loader_yields_nothing(self):
        self.assertEqual(list(synchronized_dataloader_step([], is_dist=False)), [])


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