"""Tests for Collator (T2.5): JSONL sample dicts → wrapper raw batch."""
import math

import torch

from tsmm.data.collator import Collator


def make_sample(T=8, C=3, with_nan=False, n_events=1, answer="no", instruction="Is there an anomaly?"):
    """Build a dict mimicking the gen_pipeline JSONL schema."""
    series = [[float(i * C + c) for c in range(C)] for i in range(T)]
    if with_nan:
        series[1][0] = float("nan")
    return {
        "series": series,                 # [T, C]
        "attributes": [
            {"name": f"var{c}", "unit": "%", "freq": "1Hz"} for c in range(C)
        ],
        "timestamps": [1700000000 + i for i in range(T)],
        "events": [{"t": 3, "text": "deploy"}] * n_events,
        "instruction": instruction,
        "answer": answer,
        "labels": [0] * T,
    }


class TestCollator:
    def test_basic_shapes_and_types(self):
        col = Collator(max_T=8)
        batch = col([make_sample(T=8, C=3) for _ in range(4)])
        assert batch["series"].shape == (4, 8, 3)
        assert isinstance(batch["attributes"], list) and len(batch["attributes"]) == 4
        assert isinstance(batch["timestamps"], list) and len(batch["timestamps"]) == 4
        assert isinstance(batch["events"], list) and len(batch["events"]) == 4
        assert isinstance(batch["question"], list) and len(batch["question"]) == 4
        assert isinstance(batch["answer"], list) and len(batch["answer"]) == 4

    def test_question_maps_from_instruction(self):
        col = Collator()
        batch = col([make_sample()])
        assert batch["question"][0] == "Is there an anomaly?"

    def test_answer_passthrough(self):
        col = Collator()
        batch = col([make_sample(answer="yes")])
        assert batch["answer"][0] == "yes"

    def test_no_nan_after_collate(self):
        col = Collator(nan_fill=0.0)
        batch = col([make_sample(T=8, C=3, with_nan=True)])
        assert not torch.isnan(batch["series"]).any()
        # the NaN position was filled with 0
        assert batch["series"][0, 1, 0].item() == 0.0

    def test_truncates_long_T(self):
        col = Collator(max_T=16)
        batch = col([make_sample(T=32, C=2)])
        assert batch["series"].shape == (1, 16, 2)

    def test_pads_short_T(self):
        col = Collator(max_T=64)
        batch = col([make_sample(T=16, C=2)])
        assert batch["series"].shape == (1, 64, 2)
        # padded region is zero
        assert batch["series"][0, 16:].abs().sum().item() == 0.0

    def test_pads_variable_C_within_batch(self):
        col = Collator(max_T=16)
        batch = col([make_sample(T=16, C=2), make_sample(T=16, C=5)])
        # both padded to max_C=5
        assert batch["series"].shape == (2, 16, 5)

    def test_attribute_string_is_descriptive(self):
        col = Collator()
        batch = col([make_sample(C=2)])
        s = batch["attributes"][0]
        assert "var0" in s and "var1" in s

    def test_timestamps_string(self):
        col = Collator()
        batch = col([make_sample(T=8)])
        s = batch["timestamps"][0]
        assert isinstance(s, str) and "1700000000" in s

    def test_events_string(self):
        col = Collator()
        batch = col([make_sample(n_events=2)])
        s = batch["events"][0]
        assert isinstance(s, str) and "deploy" in s

    def test_answer_optional(self):
        # answer=None path (inference): collator should omit/None answer
        col = Collator()
        sample = make_sample()
        sample["answer"] = None
        batch = col([sample])
        assert batch.get("answer") is None or batch["answer"][0] is None