from types import SimpleNamespace

import torch

from runtime import _scale_audio_velocity_for_core


def test_current_flow_av_core_keeps_inner_audio_velocity_unscaled():
    h3 = SimpleNamespace(time_shift_sigma=lambda *args: torch.tensor(0.5))
    audio_out = torch.tensor([1.0, -2.0], dtype=torch.float16)

    result = _scale_audio_velocity_for_core(h3, audio_out, torch.tensor(0.5), 12.0, 3.0)

    assert result is audio_out


def test_legacy_core_scales_inner_audio_velocity_with_schedule_slope():
    calls = []

    def time_shift_slope(sigma_v, shift_v, shift_a):
        calls.append((sigma_v, shift_v, shift_a))
        return torch.tensor(0.25, dtype=torch.float32)

    h3 = SimpleNamespace(time_shift_slope=time_shift_slope)
    sigma_v = torch.tensor(0.5)
    audio_out = torch.tensor([4.0, -8.0], dtype=torch.float16)

    result = _scale_audio_velocity_for_core(h3, audio_out, sigma_v, 12.0, 3.0)

    assert calls == [(sigma_v, 12.0, 3.0)]
    assert result.dtype == torch.float16
    assert torch.equal(result, torch.tensor([1.0, -2.0], dtype=torch.float16))