tests/test_inputs.py

import pytest
import torch

from time_anchor.base import BaseTimeAnchorPipeline
from time_anchor.cli.infer_and_explain import (
    build_forecast_rows,
    parse_column_list,
    parse_quantiles,
    variable_impact_to_rows,
)
from time_anchor.explainability import (
    GAFAttribution,
    GAFExplainabilityConfig,
    GeneralizedAttentionFlowExplainer,
)
from time_anchor.utils import left_pad_and_stack_1d


class DummyPipeline(BaseTimeAnchorPipeline):
    forecast_type = None


def test_left_pad_and_stack_1d_preserves_dtype_and_left_pads() -> None:
    batch = left_pad_and_stack_1d(
        [
            torch.tensor([1.0, 2.0], dtype=torch.float64),
            torch.tensor([3.0], dtype=torch.float64),
        ]
    )

    assert batch.dtype == torch.float64
    assert batch.shape == (2, 2)
    assert torch.isnan(batch[1, 0])
    assert batch[1, 1].item() == 3.0


def test_left_pad_and_stack_1d_rejects_empty_input() -> None:
    with pytest.raises(ValueError, match="At least one tensor"):
        left_pad_and_stack_1d([])


def test_pipeline_context_validation_batches_1d_tensors() -> None:
    pipeline = DummyPipeline(inner_model=None)
    context = pipeline._prepare_and_validate_context(torch.tensor([1.0, 2.0]))

    assert context.shape == (1, 2)


def test_pipeline_context_validation_rejects_invalid_rank() -> None:
    pipeline = DummyPipeline(inner_model=None)
    with pytest.raises(ValueError, match="1D or 2D"):
        pipeline._prepare_and_validate_context(torch.zeros(1, 2, 3))


def test_parse_quantiles_rejects_empty_values() -> None:
    with pytest.raises(ValueError, match="At least one quantile"):
        parse_quantiles(" , ")


def test_parse_column_list_trims_empty_values() -> None:
    assert parse_column_list(" Feature1, Feature2,,Feature3 ") == [
        "Feature1",
        "Feature2",
        "Feature3",
    ]


def test_gaf_config_defaults_to_barrier_flow_solver() -> None:
    assert GAFExplainabilityConfig().flow_solver == "barrier"
    assert GAFExplainabilityConfig(flow_solver="simple").flow_solver == "simple"


def test_gaf_config_does_not_accept_precision_override() -> None:
    with pytest.raises(TypeError, match="precision"):
        GAFExplainabilityConfig(precision=4)


def test_build_forecast_rows_uses_quantile_column_names() -> None:
    quantiles = torch.tensor([[[0.1, 0.5, 0.9], [0.2, 0.6, 1.0]]])
    mean = torch.tensor([[0.5, 0.6]])

    rows = build_forecast_rows(quantiles, mean, [0.1, 0.5, 0.9])

    assert rows[0]["horizon"] == 1
    assert rows[0]["mean"] == pytest.approx(0.5)
    assert rows[0]["q10"] == pytest.approx(0.1)
    assert rows[0]["q50"] == pytest.approx(0.5)
    assert rows[0]["q90"] == pytest.approx(0.9)
    assert rows[1]["horizon"] == 2
    assert rows[1]["mean"] == pytest.approx(0.6)
    assert rows[1]["q10"] == pytest.approx(0.2)
    assert rows[1]["q50"] == pytest.approx(0.6)
    assert rows[1]["q90"] == pytest.approx(1.0)


def test_gaf_attribution_json_includes_variable_timestep_attribution() -> None:
    attribution = GAFAttribution(
        token_attribution=[1.0],
        patch_attribution=[1.0],
        timestep_attribution=[{"index": 7, "contribution": 0.1234567}],
        variable_timestep_attribution=[
            {
                "index": 7,
                "time_index": 12,
                "variable_id": 2,
                "variable_name": "explanatory_2",
                "variable_role": "explanatory",
                "contribution": 0.1234567,
            }
        ],
        series_aggregates={"explanatory_2": 0.1234567},
    )

    payload = attribution.to_json()

    assert payload["variable_timestep_attribution"] == [
        {
            "index": 7,
            "time_index": 12,
            "variable_id": 2,
            "variable_name": "explanatory_2",
            "variable_role": "explanatory",
            "contribution": 0.123457,
        }
    ]
    assert payload["series_aggregates"] == {"explanatory_2": 0.123457}


def test_gaf_attribution_json_does_not_accept_precision_override() -> None:
    attribution = GAFAttribution(
        token_attribution=[],
        patch_attribution=[],
        timestep_attribution=[],
        variable_timestep_attribution=[],
        series_aggregates={},
    )

    with pytest.raises(TypeError, match="precision"):
        attribution.to_json(precision=4)


def test_variable_impact_to_rows_exports_variable_columns() -> None:
    rows = variable_impact_to_rows(
        {
            "variable_timestep_attribution": [
                {
                    "index": 3,
                    "time_index": 10,
                    "variable_id": 1,
                    "variable_name": "explanatory_1",
                    "variable_role": "explanatory",
                    "series_id": 4,
                    "segment": 2.0,
                    "contribution": 0.25,
                }
            ]
        }
    )

    assert rows == [
        {
            "index": 3,
            "time_index": 10,
            "variable_id": 1,
            "variable_name": "explanatory_1",
            "variable_role": "explanatory",
            "series_id": 4,
            "segment": 2.0,
            "impact": 0.25,
        }
    ]


def test_variable_timestep_entries_normalize_per_time_index() -> None:
    entries = [
        {"time_index": 0, "variable_id": 0, "variable_name": "target", "contribution": 2.0},
        {
            "time_index": 0,
            "variable_id": 1,
            "variable_name": "explanatory_1",
            "contribution": 1.0,
        },
        {"time_index": 1, "variable_id": 0, "variable_name": "target", "contribution": 4.0},
    ]

    normalized = GeneralizedAttentionFlowExplainer._normalize_variable_timestep_entries(entries)

    sums_by_time = {}
    for entry in normalized:
        sums_by_time.setdefault(entry["time_index"], 0.0)
        sums_by_time[entry["time_index"]] += entry["contribution"]

    assert sums_by_time == {0: pytest.approx(1.0), 1: pytest.approx(1.0)}
    assert normalized[0]["contribution"] == pytest.approx(2 / 3)
    assert normalized[1]["contribution"] == pytest.approx(1 / 3)