예제 소스

main_llm_db_rl_qt_chart_browser.py

"""
DB 캔들(`Market.get_data`)을 ``main_llm_db_rl_viewer.py`` 와 동일하게 불러와
Tk 뷰어와 같은 ``ChartImageGenerator.render_and_slice_by_x`` 로 PNG 차트를 만든 뒤,
Qt 창에서 ``<`` / ``>`` 로 봉 인덱스를 한 칸씩 옮겨 가며 본다.

설정: ``config/llm_rl_train.py``
차트 상수: ``main_llm_db_test_ollama``

실행 (프로젝트 루트):
  pip install PySide6
  python main_llm_db_rl_qt_chart_browser.py

키보드: 왼쪽 화살표 = 이전 봉, 오른쪽 화살표 = 다음 봉
슬라이더: 버튼 아래에서 bar 인덱스를 끌어서 이동 (렌더 중에도 위치 변경 가능, 완료 후 반영)
"""

from __future__ import annotations

import io
import logging
import sys
from typing import TYPE_CHECKING

import main_llm_db_test_ollama as ollama_bt
from ChartImageGenerator import ChartImageConfig, ChartImageGenerator
from Market import Market
from config import llm_rl_train as tr
from lib.rl_data_labels import merge_extreme_labels_from_csv
from PIL import Image

if TYPE_CHECKING:
    import pandas as pd

try:
    from PySide6.QtCore import Qt, QThread, Signal
    from PySide6.QtGui import QImage, QKeySequence, QPixmap, QShortcut
    from PySide6.QtWidgets import (
        QApplication,
        QHBoxLayout,
        QLabel,
        QMainWindow,
        QMessageBox,
        QPushButton,
        QScrollArea,
        QSizePolicy,
        QSlider,
        QVBoxLayout,
        QWidget,
    )
except ImportError as e:
    raise SystemExit(
        "PySide6 가 필요합니다: pip install PySide6\n"
        f"원인: {e}"
    ) from e

logger = logging.getLogger(__name__)


def _load_df():
    market = Market()
    df = market.get_data()
    if df is None or len(df) < tr.MIN_BARS + tr.BAR_STEP:
        raise SystemExit("데이터가 부족합니다.")

    n_full = len(df)
    bar_offset = 0
    cap = getattr(tr, "RL_MAX_BARS", None)
    if cap is not None and n_full > int(cap):
        bar_offset = n_full - int(cap)
        df = df.iloc[-int(cap) :].reset_index(drop=True)
        logger.info(
            "RL_MAX_BARS=%s: tail %s봉 (전체 %s봉, bar_offset=%s)",
            cap,
            len(df),
            n_full,
            bar_offset,
        )

    df = merge_extreme_labels_from_csv(df, tr.LABELS_CSV_PATH, source_bar_offset=bar_offset)
    df = ollama_bt._ensure_pb_mfi10(df)
    return df


def _resize_png_like_viewer(png_bytes: bytes, chart_max_width: int) -> bytes:
    im = Image.open(io.BytesIO(png_bytes)).convert("RGB")
    w, h = im.size
    mw = max(400, int(chart_max_width))
    if w > mw:
        ratio = mw / w
        im = im.resize((int(w * ratio), int(h * ratio)), Image.Resampling.LANCZOS)
    buf = io.BytesIO()
    im.save(buf, format="PNG")
    return buf.getvalue()


def _render_chart_png(df, chart_gen: ChartImageGenerator, bar_idx: int) -> bytes:
    df_hist = df.iloc[: bar_idx + 1].copy()
    chart_idx = len(df_hist)
    base_png, _parts = chart_gen.render_and_slice_by_x(
        df_hist,
        idx=chart_idx,
        max_data=min(ollama_bt.CHART_MAX_DATA, len(df_hist)),
        ticker=ollama_bt.TICKER,
        n_slices=ollama_bt.CHART_N_SLICES,
        overlap_px=ollama_bt.CHART_OVERLAP_PX,
        scale=ollama_bt.CHART_SCALE,
    )
    return base_png


def _png_to_qpixmap(png_bytes: bytes) -> QPixmap:
    img = QImage.fromData(png_bytes, "PNG")
    if img.isNull():
        raise RuntimeError("PNG → QImage 변환 실패")
    return QPixmap.fromImage(img)


class _ChartRenderThread(QThread):
    rendered = Signal(int, bytes, str)

    def __init__(
        self,
        df: pd.DataFrame,
        chart_gen: ChartImageGenerator,
        chart_max_width: int,
        bar: int,
        n: int,
    ) -> None:
        super().__init__()
        self._df = df
        self._chart_gen = chart_gen
        self._chart_max_width = chart_max_width
        self._bar = bar
        self._n = n

    def run(self) -> None:
        raw = _render_chart_png(self._df, self._chart_gen, self._bar)
        disp = _resize_png_like_viewer(raw, self._chart_max_width)
        row = self._df.iloc[self._bar]
        ts = row.name if hasattr(row.name, "strftime") else self._bar
        status = (
            f"bar 인덱스={self._bar} / {self._n - 1} · 시각(인덱스)={ts} · "
            f"종가={float(row['close']):,.0f}"
        )
        self.rendered.emit(self._bar, disp, status)


class ChartBrowserWindow(QMainWindow):
    def __init__(
        self,
        *,
        df: pd.DataFrame,
        chart_gen: ChartImageGenerator,
        chart_max_width: int,
    ) -> None:
        super().__init__()
        self._df = df
        self._chart_gen = chart_gen
        self._chart_max_width = chart_max_width
        self._n = len(df)
        self._min_bar = max(50, int(tr.MIN_BARS))
        if self._n <= self._min_bar:
            raise SystemExit(f"행 수({self._n})가 최소 봉({self._min_bar}) 이하입니다.")
        self._bar = self._min_bar
        self._thread: _ChartRenderThread | None = None
        self._slider_sync = False

        title = getattr(tr, "PNG_STATUS_WINDOW_TITLE", "DB 차트 브라우저")
        self.setWindowTitle(f"{title} (Qt)")
        self.resize(980, 720)

        central = QWidget()
        self.setCentralWidget(central)
        layout = QVBoxLayout(central)

        layout.addWidget(
            QLabel(f"<b>{ollama_bt.TICKER}</b> · DB 차트 (뷰어와 동일 렌더)")
        )

        nav = QHBoxLayout()
        self._btn_prev = QPushButton("< 이전 봉")
        self._btn_next = QPushButton("다음 봉 >")
        self._btn_prev.clicked.connect(self._on_prev)
        self._btn_next.clicked.connect(self._on_next)
        nav.addWidget(self._btn_prev)
        nav.addWidget(self._btn_next)
        nav.addStretch()
        layout.addLayout(nav)

        slider_row = QHBoxLayout()
        self._slider = QSlider(Qt.Orientation.Horizontal)
        self._slider.setMinimum(self._min_bar)
        self._slider.setMaximum(self._n - 1)
        self._slider.setValue(self._bar)
        span = max(1, (self._n - 1) - self._min_bar)
        self._slider.setPageStep(max(1, span // 24))
        self._slider.setSingleStep(1)
        self._slider.setTracking(True)
        self._slider.valueChanged.connect(self._on_slider_value_changed)
        slider_row.addWidget(QLabel("봉"), stretch=0)
        slider_row.addWidget(self._slider, stretch=1)
        self._lbl_slider_pos = QLabel("")
        self._lbl_slider_pos.setMinimumWidth(120)
        self._lbl_slider_pos.setAlignment(Qt.AlignmentFlag.AlignRight | Qt.AlignmentFlag.AlignVCenter)
        self._update_slider_caption()
        slider_row.addWidget(self._lbl_slider_pos, stretch=0)
        layout.addLayout(slider_row)

        self._lbl_status = QLabel("")
        layout.addWidget(self._lbl_status)

        scroll = QScrollArea()
        scroll.setWidgetResizable(True)
        scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAsNeeded)
        scroll.setVerticalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAsNeeded)
        self._lbl_chart = QLabel()
        self._lbl_chart.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self._lbl_chart.setSizePolicy(
            QSizePolicy.Policy.Expanding,
            QSizePolicy.Policy.MinimumExpanding,
        )
        scroll.setWidget(self._lbl_chart)
        layout.addWidget(scroll, stretch=1)

        sc_prev = QShortcut(QKeySequence(Qt.Key.Key_Left), self)
        sc_prev.activated.connect(self._on_prev)
        sc_next = QShortcut(QKeySequence(Qt.Key.Key_Right), self)
        sc_next.activated.connect(self._on_next)

        self._start_render(self._bar)

    def _update_slider_caption(self) -> None:
        self._lbl_slider_pos.setText(f"{self._bar} / {self._n - 1}")

    def _sync_slider_to_bar(self) -> None:
        self._slider_sync = True
        try:
            self._slider.setValue(self._bar)
        finally:
            self._slider_sync = False
        self._update_slider_caption()

    def _set_busy(self, busy: bool) -> None:
        self._btn_prev.setEnabled(not busy)
        self._btn_next.setEnabled(not busy)

    def _on_slider_value_changed(self, v: int) -> None:
        if self._slider_sync:
            return
        self._update_slider_caption()
        if self._thread and self._thread.isRunning():
            return
        v = max(self._min_bar, min(self._n - 1, int(v)))
        if v == self._bar:
            return
        self._bar = v
        self._start_render(self._bar)

    def _on_prev(self) -> None:
        if self._thread and self._thread.isRunning():
            return
        self._bar = max(self._min_bar, self._bar - 1)
        self._sync_slider_to_bar()
        self._start_render(self._bar)

    def _on_next(self) -> None:
        if self._thread and self._thread.isRunning():
            return
        self._bar = min(self._n - 1, self._bar + 1)
        self._sync_slider_to_bar()
        self._start_render(self._bar)

    def _start_render(self, bar: int) -> None:
        self._set_busy(True)
        self._lbl_status.setText(f"렌더 중… (bar={bar})")
        self._thread = _ChartRenderThread(
            self._df,
            self._chart_gen,
            self._chart_max_width,
            bar,
            self._n,
        )
        self._thread.rendered.connect(self._on_rendered)
        self._thread.finished.connect(self._on_thread_finished)
        self._thread.start()

    def _on_rendered(self, bar: int, png_bytes: bytes, status: str) -> None:
        self._bar = bar
        self._sync_slider_to_bar()
        try:
            pix = _png_to_qpixmap(png_bytes)
            self._lbl_chart.setPixmap(pix)
            self._lbl_status.setText(status)
        except Exception as e:
            self._lbl_status.setText(f"표시 오류: {e}")
            QMessageBox.warning(self, "차트", str(e))

    def _on_thread_finished(self) -> None:
        self._set_busy(False)
        v = int(self._slider.value())
        if v != self._bar:
            self._bar = v
            self._update_slider_caption()
            self._start_render(self._bar)


def main() -> None:
    logging.basicConfig(level=logging.INFO)
    print("[qt_chart_browser] Market.get_data() …", flush=True)
    df = _load_df()
    chart_max_width = int(getattr(tr, "CHART_MAX_WIDTH", 900))
    chart_gen = ChartImageGenerator(
        ticker=ollama_bt.TICKER,
        config=ChartImageConfig(show=False),
    )

    app = QApplication(sys.argv)
    win = ChartBrowserWindow(df=df, chart_gen=chart_gen, chart_max_width=chart_max_width)
    win.show()
    print(f"[qt_chart_browser] 준비 완료 — {len(df)}행", flush=True)
    sys.exit(app.exec())


if __name__ == "__main__":
    main()