tests/test_tensor_builder_v2.py

from __future__ import annotations

import math

import numpy as np
import pandas as pd
import pytest

from sctm_next_round.tensor_builder_v2 import (
    AUDIENCE_NA_ID,
    AUDIENCE_OTHER_ID,
    CADENCE_TEETH,
    NUMERIC_DROP_FIELDS_V2,
    STATIC_NUMERIC_DATE_COLS,
    STATIC_NUMERIC_DURATION_COLS,
    build_static_numeric,
    compute_observation_features,
    compute_patient_cadence,
    encode_audience,
    encode_drug_flags,
    encode_v1_event_values,
    enumerate_windows,
    fit_cat_vocab_v2,
    fit_dose_winsor_caps,
    split_static_columns,
)


# ---------------------------------------------------------------------------
# window enumeration parity with the v1 inline algorithm
# ---------------------------------------------------------------------------

def v1_reference_windows(n_visits: int, max_seq_len: int, stride: int) -> list[tuple[int, int]]:
    """Literal transcription of tensor_builder.build_split_arrays enumeration."""

    starts = list(range(0, n_visits, max(1, stride)))
    if starts and starts[-1] + max_seq_len < n_visits:
        starts.append(max(0, n_visits - max_seq_len))
    out = []
    for start in starts:
        length = int(min(max_seq_len, max(0, n_visits - start)))
        if length < 1:
            continue
        if n_visits > 1 and length < 2:
            continue
        out.append((start, length))
    return out


class TestWindowEnumeration:
    @pytest.mark.parametrize("n", [1, 2, 63, 64, 65, 127, 128, 129, 200, 475])
    def test_matches_v1_reference(self, n):
        assert enumerate_windows(n, 128, 64) == v1_reference_windows(n, 128, 64)

    def test_singleton_patient_gets_one_window(self):
        assert enumerate_windows(1, 128, 64) == [(0, 1)]

    def test_tail_window_appended(self):
        # 200 visits, stride 64: starts 0,64,128,192 -> the 192 window has length 8 (>=2, kept)
        wins = enumerate_windows(200, 128, 64)
        assert (192, 8) in wins


# ---------------------------------------------------------------------------
# cadence block
# ---------------------------------------------------------------------------

class TestCadence:
    def test_trailing_median_and_deviation(self):
        # gaps: -, 14, 0(same-day), 28, 90
        time_days = np.array([0.0, 14.0, 14.0, 42.0, 132.0], dtype=np.float32)
        cad = compute_patient_cadence(time_days)
        # first visit undefined
        assert cad["cadence_trailing_median"][0] == 0.0
        assert cad["cadence_deviation_log"][0] == 0.0
        # t=1: first positive gap 14 -> trailing median 14, no prior rhythm -> deviation 0
        assert cad["cadence_trailing_median"][1] == 14.0
        assert cad["cadence_deviation_log"][1] == 0.0
        # t=2: same-day 0-gap excluded from the rhythm (red line 3)
        assert cad["cadence_trailing_median"][2] == 14.0
        assert cad["cadence_deviation_log"][2] == 0.0
        assert cad["cadence_tooth_distance"][2] == 0.0
        # t=3: gap 28 vs prior rhythm 14
        assert cad["cadence_trailing_median"][3] == pytest.approx(21.0)
        assert cad["cadence_deviation_log"][3] == pytest.approx(math.log(29.0 / 15.0), abs=1e-6)
        # t=4: gap 90 vs prior rhythm median([14, 28]) = 21
        assert cad["cadence_deviation_log"][4] == pytest.approx(math.log(91.0 / 22.0), abs=1e-6)
        assert cad["cadence_trailing_median"][4] == pytest.approx(28.0)

    def test_trailing_window_is_five(self):
        gaps = [10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0]
        time_days = np.cumsum([0.0] + gaps).astype(np.float32)
        cad = compute_patient_cadence(time_days)
        # at the last visit the trailing window is the last 5 gaps [30,40,50,60,70]
        assert cad["cadence_trailing_median"][-1] == pytest.approx(50.0)

    def test_tooth_distance(self):
        time_days = np.array([0.0, 14.0, 44.0, 134.0, 154.0], dtype=np.float32)
        cad = compute_patient_cadence(time_days)
        # gap 14 -> exactly on the 14 tooth
        assert cad["cadence_tooth_distance"][1] == 0.0
        # gap 30 -> nearest tooth 28: |30-28|/28
        assert cad["cadence_tooth_distance"][2] == pytest.approx(2.0 / 28.0)
        # gap 90 -> on the 90 tooth
        assert cad["cadence_tooth_distance"][3] == 0.0
        # gap 20 -> min(|20-14|/14, |20-28|/28) = 6/14 vs 8/28 -> 8/28
        assert cad["cadence_tooth_distance"][4] == pytest.approx(min(6.0 / 14.0, 8.0 / 28.0, 70.0 / 90.0))
        assert set(CADENCE_TEETH) == {14.0, 28.0, 90.0}


# ---------------------------------------------------------------------------
# observation-feature precompute == formal_model torch loop
# ---------------------------------------------------------------------------

class TestObservationFeatures:
    def test_matches_formal_model_torch_loop(self):
        torch = pytest.importorskip("torch")
        from sctm_next_round.formal_model import SCTMFormalModel

        rng = np.random.default_rng(20260611)
        bsz, seq_len, n_fields = 7, 19, 11
        observed = rng.random((bsz, seq_len, n_fields)) < 0.35
        gaps = rng.integers(0, 60, size=(bsz, seq_len)).astype(np.float64)
        gaps[:, 0] = 0.0
        times = np.cumsum(gaps, axis=1).astype(np.float32)

        batch = {"time_since_start_days": torch.from_numpy(times)}
        # the method touches no instance state — call it unbound on the literal implementation
        torch_since, torch_carry = SCTMFormalModel._time_since_last_observed(
            None, batch, torch.from_numpy(observed)
        )
        np_since, np_carry = compute_observation_features(observed, times)
        np.testing.assert_allclose(np_since, torch_since.numpy(), rtol=0, atol=1e-4)
        np.testing.assert_array_equal(np_carry, torch_carry.numpy())

    def test_window_reset_semantics(self):
        # state resets at the window start: nothing seen before position 0
        observed = np.zeros((1, 4, 1), dtype=bool)
        observed[0, 1, 0] = True
        times = np.array([[0.0, 10.0, 25.0, 60.0]], dtype=np.float32)
        since, carry = compute_observation_features(observed, times)
        np.testing.assert_allclose(since[0, :, 0], [0.0, 0.0, 15.0, 50.0])
        np.testing.assert_array_equal(carry[0, :, 0], [False, True, True, True])

    def test_visit_completeness_formula(self):
        # completeness = mean over (cat observed, numeric mask, ordinal mask)
        missing_ids = np.array([[[0, 3], [3, 3]]])  # cat observed: [T0: 1/2, T1: 0/2]
        numeric_mask = np.array([[[True], [False]]])
        ordinal_mask = np.array([[[True], [True]]])
        observed_all = np.concatenate(
            [missing_ids == 0, numeric_mask, ordinal_mask], axis=2
        )
        completeness = observed_all.mean(axis=2)
        np.testing.assert_allclose(completeness[0], [0.75, 0.25])


# ---------------------------------------------------------------------------
# encoders
# ---------------------------------------------------------------------------

class TestEncoders:
    def test_audience_vocab(self):
        series = pd.Series(["本人", "家属", "居委干部", "社区民警", "民警", "邻居", "其他", "亲戚", None, ""])
        out = encode_audience(series)
        assert out.tolist() == [0, 1, 2, 3, 3, 4, 5, AUDIENCE_OTHER_ID, AUDIENCE_NA_ID, AUDIENCE_NA_ID]
        assert out.dtype == np.int8

    def test_drug_flags(self):
        df = pd.DataFrame(
            {
                "随访药品1名称": [
                    "棕榈酸帕利哌酮注射液",      # LAI by name regex
                    "氯氮平片|阿立哌唑",          # clozapine
                    "盐酸苯海索片",              # anticholinergic
                    "利培酮",                    # none
                    None,
                ],
                "药品1剂量长效药": [None, None, None, "25", None],
            }
        )
        flags = encode_drug_flags(df)
        assert flags.shape == (5, 3)
        assert flags[0].tolist() == [True, False, False]
        assert flags[1].tolist() == [False, True, False]
        assert flags[2].tolist() == [False, False, True]
        # 剂量长效药 nonnull -> LAI even without a regex hit (codex red line 10)
        assert flags[3].tolist() == [True, False, False]
        assert flags[4].tolist() == [False, False, False]

    def test_v1_event_values_vectorized(self):
        df = pd.DataFrame(
            {
                "是否转诊": ["是", "否", ""],
                "是否应急处置": ["否", "是", ""],
                "治疗方式": ["住院", "门诊", ""],
                "随访类型": ["", "住院随访", ""],
                "危险行为等级": ["0级", "1级 ", ""],
                "滋事次数": [0, 2, None],
                "肇事次数": [0, 0, 0],
                "肇祸次数": [0, 0, 0],
                "自杀未遂次数": [0, 0, 0],
                "自伤次数": [0, 0, 0],
                "其他危害他人行为": [0, 0, 0],
            }
        )
        ev = encode_v1_event_values(df)
        np.testing.assert_allclose(ev[0], [1, 0, 1, 0, 0])
        # trailing-space danger token still counts as high danger (matches v1 value_text strip)
        np.testing.assert_allclose(ev[1], [0, 1, 1, 1, 1])
        np.testing.assert_allclose(ev[2], [0, 0, 0, 0, 0])

    def test_cat_vocab_whitespace_normalized(self):
        df = pd.DataFrame(
            {
                "split": ["train"] * 4,
                "危险行为等级": ["1级", "1级 ", "1级　", "2级"],
            }
        )
        vocab = fit_cat_vocab_v2(df, ["危险行为等级"], min_freq=1)
        keys = [k for k in vocab if k.startswith("危险行为等级=")]
        assert "危险行为等级=1级" in keys
        # whitespace variants collapse into ONE token
        assert len([k for k in keys if "1级" in k]) == 1

    def test_dose_winsor_caps_train_only(self):
        train_vals = list(np.linspace(1, 100, 999)) + [100150.0]  # keying error in train
        df = pd.DataFrame(
            {
                "split": ["train"] * 1000 + ["val"],
                "__dose": train_vals + [50.0],
            }
        )
        caps = fit_dose_winsor_caps(df, {"药品1非长效药物晨间服用剂量": "__dose"})
        cap = caps["药品1非长效药物晨间服用剂量"]
        assert cap < 100150.0
        assert cap >= 99.0

    def test_numeric_drop_list_matches_codex_c1(self):
        assert set(NUMERIC_DROP_FIELDS_V2) == {
            "公安滋事总数", "公安肇事总数", "以往住院情况", "本次住院时长天数",
            "易肇事肇祸次数", "其他危害他人行为", "药品1用法长效药",
        }


# ---------------------------------------------------------------------------
# static routing
# ---------------------------------------------------------------------------

class TestStaticRouting:
    def test_split_static_columns(self):
        cols = (
            ["has_static", "A_static_性别", "A_static_性别_missing"]
            + ["A_static_曾肇事肇祸", "A_static_曾肇事肇祸_missing", "A_static_家族史", "A_static_关解锁情况"]
            + STATIC_NUMERIC_DATE_COLS
            + [c + "_missing" for c in STATIC_NUMERIC_DATE_COLS]
            + STATIC_NUMERIC_DURATION_COLS
            + ["A_static_recent_violence_date", "A_static_years_since_recent_violence_at_first_visit"]
        )
        kept, numeric_routed, dropped = split_static_columns(cols)
        assert kept == ["has_static", "A_static_性别", "A_static_性别_missing"]
        assert set(numeric_routed) == set(STATIC_NUMERIC_DATE_COLS + STATIC_NUMERIC_DURATION_COLS)
        assert "A_static_曾肇事肇祸" in dropped
        assert "A_static_recent_violence_date" in dropped

    def test_static_numeric_sentinel_guard_and_zscore(self):
        static = pd.DataFrame(
            {
                "patient_id": ["a", "b", "c", "d"],
                "A_static_birth_date": ["1960-05-01", "1900-01-01", "1975-01-01", ""],
                "A_static_age_at_first_visit_years": [50.0, 40.0, 999.0, np.nan],
            }
        )
        values, mask, cols, scaler = build_static_numeric(static, train_ids={"a", "b"})
        bidx = cols.index("A_static_birth_date__frac_year")
        didx = cols.index("A_static_age_at_first_visit_years")
        # 1900 sentinel masked, empty masked
        assert mask[:, bidx].tolist() == [True, False, True, False]
        # out-of-range duration masked
        assert mask[:, didx].tolist() == [True, True, False, False]
        # masked entries carry 0
        assert values[1, bidx] == 0.0
        assert values[3, didx] == 0.0
        # z-score fitted on train only (patients a,b): age mean 45 std 5
        assert scaler["A_static_age_at_first_visit_years"]["mean"] == pytest.approx(45.0)
        assert values[0, didx] == pytest.approx((50.0 - 45.0) / 5.0)


# ---------------------------------------------------------------------------
# tiny end-to-end through build_split_arrays_v2 (synthetic)
# ---------------------------------------------------------------------------

class TestSplitAssembly:
    def test_label_arrays_window_aligned(self):
        from sctm_next_round.tensor_builder_v2 import build_split_arrays_v2, encode_rows

        n = 5
        rows = pd.DataFrame(
            {
                "A表编号": ["p1"] * 3 + ["p2"] * 2,
                "visit_index": [0, 1, 2, 0, 1],
                "visit_date": pd.to_datetime(
                    ["2020-01-01", "2020-01-15", "2020-02-12", "2021-01-01", "2021-04-01"]
                ),
                "__year__": [2020, 2020, 2020, 2021, 2021],
                "split": ["train"] * n,
                "terminal_or_disposition_label": ["none", "none", "none", "none", "death"],
                "危险行为等级": ["0级", "1级", "0级", "0级", "0级"],
                "本次访视对象": ["本人", "家属", "本人", "居委干部", "本人"],
                "随访药品1名称": ["氯氮平", None, "棕榈酸帕利哌酮", None, None],
                # label-v2 columns
                "terminal_label_v2": ["none", "none", "none", "none", "death"],
                "diseng_state": [0, 1, 0, 0, 0],
                "diseng_episode_start": [False, True, False, False, False],
                "time_to_return_days": [np.nan, 28.0, np.nan, np.nan, np.nan],
                "return_within_365": [-1, 1, -1, -1, -1],
                "post_discharge_gap90": [-1, -1, 0, -1, -1],
                "same_day_flag": [False, False, False, False, False],
                "n_same_day": [1, 1, 1, 1, 1],
                "event_hosp_newdate": [False, True, False, False, False],
                "event_discharge_newdate": [False, False, True, False, False],
            }
        )
        cat_cols = ["危险行为等级", "本次访视对象"]
        ord_cols = ["危险行为等级"]
        num_cols: list[str] = []
        cat_vocab = {"<PAD>": 0, "<UNK>": 1, "危险行为等级=0级": 2, "危险行为等级=1级": 3, "本次访视对象=本人": 4, "本次访视对象=家属": 5, "本次访视对象=居委干部": 6}
        drug_vocab = {"<PAD>": 0, "<UNK>": 1, "氯氮平": 2}
        rank_maps = {"危险行为等级": {"0级": 0, "1级": 1}}
        encoded = encode_rows(
            rows,
            cat_cols=cat_cols,
            num_cols=num_cols,
            ord_cols=ord_cols,
            cat_vocab=cat_vocab,
            drug_vocab=drug_vocab,
            numeric_scaler={},
            parsed_cols={},
            rank_maps=rank_maps,
            cbe_dim=4,
            max_drug_atoms=2,
        )
        arrays = build_split_arrays_v2(
            rows,
            encoded,
            static_index_by_patient={"p1": 0, "p2": 1},
            static_cat_ids=np.array([[7], [8]], dtype=np.int64),
            static_numeric=np.array([[0.5], [-0.5]], dtype=np.float32),
            static_numeric_mask=np.array([[True], [False]]),
            n_cat=len(cat_cols),
            n_num=0,
            n_ord=len(ord_cols),
            cbe_dim=4,
            max_drug_atoms=2,
            max_seq_len=8,
            stride=4,
            progress_every=0,
        )
        assert arrays["valid_mask"].shape[0] == 2  # one window per patient
        w0 = list(arrays["patient_ids"]).index("p1")
        w1 = list(arrays["patient_ids"]).index("p2")
        # window-aligned label arrays
        assert arrays["diseng_state"][w0, :3].tolist() == [0, 1, 0]
        assert arrays["diseng_episode_start"][w0, 1]
        assert arrays["time_to_return_days"][w0, 1] == 28.0
        assert np.isnan(arrays["time_to_return_days"][w0, 0])
        assert np.isnan(arrays["time_to_return_days"][w0, 5])  # padding stays NaN
        assert arrays["return_within_365"][w0, :3].tolist() == [-1, 1, -1]
        assert arrays["post_discharge_gap90"][w0, :3].tolist() == [-1, -1, 0]
        assert arrays["terminal_label_v2"][w1, :2].tolist() == [0, 1]  # none, death
        # event_labels_v2 channel order: referral, emergency, hosp_new, discharge_new, danger, harm
        assert arrays["event_labels_v2"][w0, 1, 2] == 1.0
        assert arrays["event_labels_v2"][w0, 2, 3] == 1.0
        # v2 inputs
        assert arrays["audience_id"][w0, :3].tolist() == [0, 1, 0]
        assert arrays["drug_flags"][w0, 0].tolist() == [False, True, False]   # clozapine
        assert arrays["drug_flags"][w0, 2].tolist() == [True, False, False]   # LAI by regex
        # cadence: p1 gaps 14, 28
        assert arrays["cadence_trailing_median"][w0, :3].tolist() == pytest.approx([0.0, 14.0, 21.0])
        assert arrays["cadence_deviation_log"][w0, 2] == pytest.approx(math.log(29.0 / 15.0), abs=1e-6)
        # static channels broadcast to windows
        assert arrays["static_numeric"][w0, 0] == 0.5
        assert arrays["static_numeric_mask"][w1, 0] == False  # noqa: E712
        # obs features exist with the family-ordered F axis
        assert arrays["obs_time_since_last"].shape == (2, 8, len(cat_cols) + 0 + len(ord_cols))
        assert arrays["visit_completeness"].shape == (2, 8)
        # v1-compatible arrays still present
        for key in ("cat_value_ids", "missing_ids", "ordinal_cbe", "service_state", "delta_t_next_log", "terminal_label"):
            assert key in arrays
        assert arrays["terminal_label"][w1, :2].tolist() == [0, 1]