atm/tests/test_levels.py

"""Tests for src/atm/levels.py."""
from __future__ import annotations

import numpy as np
import pytest

from atm.config import (
    CanaryRegion,
    ColorSpec,
    Config,
    DiscordCfg,
    ROI,
    TelegramCfg,
    YAxisCalib,
)
from atm.levels import Levels, LevelsExtractor, LevelsResult
from atm.vision import pixel_y_to_price

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
# chart_roi starts at (0,0) so test frames can be exactly (H, W)
CHART_ROI = ROI(x=0, y=0, w=600, h=400)
CALIB = YAxisCalib(p1_y=0, p1_price=100.0, p2_y=400, p2_price=80.0)

# light_red RGB=(255,0,0) → BGR=(0,0,255)
# light_green RGB=(0,255,0) → BGR=(0,255,0)
RED_BGR: tuple[int, int, int] = (0, 0, 255)
GREEN_BGR: tuple[int, int, int] = (0, 255, 0)

TOLERANCE = 30.0


def _make_cfg(phaseb_timeout_s: int = 600) -> Config:
    colors = {
        "turquoise":   ColorSpec(rgb=(0, 255, 255),   tolerance=TOLERANCE),
        "yellow":      ColorSpec(rgb=(255, 255, 0),   tolerance=TOLERANCE),
        "dark_green":  ColorSpec(rgb=(0, 100, 0),     tolerance=TOLERANCE),
        "dark_red":    ColorSpec(rgb=(100, 0, 0),     tolerance=TOLERANCE),
        "light_green": ColorSpec(rgb=(0, 255, 0),     tolerance=TOLERANCE),
        "light_red":   ColorSpec(rgb=(255, 0, 0),     tolerance=TOLERANCE),
        "gray":        ColorSpec(rgb=(128, 128, 128), tolerance=TOLERANCE),
    }
    return Config(
        window_title="test",
        dot_roi=ROI(x=0, y=0, w=100, h=50),
        chart_roi=CHART_ROI,
        colors=colors,
        y_axis=CALIB,
        canary=CanaryRegion(
            roi=ROI(x=0, y=0, w=50, h=50),
            baseline_phash="0" * 64,
            drift_threshold=8,
        ),
        discord=DiscordCfg(webhook_url="http://example.com/hook"),
        telegram=TelegramCfg(bot_token="tok", chat_id="123"),
        phaseb_timeout_s=phaseb_timeout_s,
    )


# ---------------------------------------------------------------------------
# Frame helpers
# ---------------------------------------------------------------------------

def _make_chart(*line_specs: tuple[tuple[int, int, int], int]) -> np.ndarray:
    """Return a 400×600 black BGR frame with horizontal lines.

    Each spec is (bgr_color, y_position).  Lines are painted 3px thick
    and span the full width so Hough detection is reliable.
    """
    frame = np.zeros((400, 600, 3), dtype=np.uint8)
    for bgr, y in line_specs:
        y0 = max(0, y - 1)
        y1 = min(400, y + 2)
        frame[y0:y1, :] = bgr
    return frame


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------

def test_three_lines_buy_complete() -> None:
    """BUY with 3 stable lines → complete after 2 calls, prices correct."""
    # red = SL at y=300 (bottom / lowest price for BUY)
    # green = TP1 at y=200, TP2 at y=100
    frame = _make_chart(
        (RED_BGR, 300),
        (GREEN_BGR, 200),
        (GREEN_BGR, 100),
    )
    cfg = _make_cfg()
    ext = LevelsExtractor(cfg, direction="BUY", start_ts=0.0)

    r1 = ext.step(frame, ts=1.0)
    assert r1.status == "partial"  # first call: not yet stable

    r2 = ext.step(frame, ts=2.0)
    assert r2.status == "complete"
    assert r2.levels is not None
    assert r2.levels.partial is False

    expected_tp2 = pixel_y_to_price(100, CALIB)   # 95.0
    expected_tp1 = pixel_y_to_price(200, CALIB)   # 90.0
    expected_sl  = pixel_y_to_price(300, CALIB)   # 85.0

    assert r2.levels.tp2 == pytest.approx(expected_tp2, abs=1.0)
    assert r2.levels.tp1 == pytest.approx(expected_tp1, abs=1.0)
    assert r2.levels.sl  == pytest.approx(expected_sl,  abs=1.0)


def test_two_lines_partial() -> None:
    """2 lines → always partial."""
    frame = _make_chart((RED_BGR, 300), (GREEN_BGR, 100))
    cfg = _make_cfg()
    ext = LevelsExtractor(cfg, direction="BUY", start_ts=0.0)

    result = ext.step(frame, ts=1.0)

    assert result.status == "partial"
    assert result.levels is not None
    assert result.levels.partial is True


def test_zero_lines_waiting() -> None:
    """No lines in chart → waiting."""
    frame = np.zeros((400, 600, 3), dtype=np.uint8)
    cfg = _make_cfg()
    ext = LevelsExtractor(cfg, direction="BUY", start_ts=0.0)

    result = ext.step(frame, ts=1.0)

    assert result.status == "waiting"
    assert result.levels is None


def test_timeout() -> None:
    """Elapsed > phaseb_timeout_s → timeout regardless of lines."""
    frame = np.zeros((400, 600, 3), dtype=np.uint8)
    cfg = _make_cfg(phaseb_timeout_s=600)
    ext = LevelsExtractor(cfg, direction="BUY", start_ts=0.0)

    result = ext.step(frame, ts=700.0)

    assert result.status == "timeout"
    assert result.levels is None
    assert result.elapsed_s == pytest.approx(700.0)


def test_sell_direction_assignment() -> None:
    """SELL: topmost y → SL (highest price), bottom → TP2 (lowest price)."""
    frame = _make_chart(
        (RED_BGR, 300),
        (GREEN_BGR, 200),
        (GREEN_BGR, 100),
    )
    cfg = _make_cfg()
    ext = LevelsExtractor(cfg, direction="SELL", start_ts=0.0)

    ext.step(frame, ts=1.0)           # first call, not yet stable
    r = ext.step(frame, ts=2.0)       # second call, stable → complete

    assert r.status == "complete"
    assert r.levels is not None

    expected_sl  = pixel_y_to_price(100, CALIB)  # 95.0 (topmost = highest price)
    expected_tp1 = pixel_y_to_price(200, CALIB)  # 90.0
    expected_tp2 = pixel_y_to_price(300, CALIB)  # 85.0

    assert r.levels.sl  == pytest.approx(expected_sl,  abs=1.0)
    assert r.levels.tp1 == pytest.approx(expected_tp1, abs=1.0)
    assert r.levels.tp2 == pytest.approx(expected_tp2, abs=1.0)