src/ui/chart_widget.py

import pyqtgraph as pg
from PyQt6.QtWidgets import QWidget, QVBoxLayout
from PyQt6.QtCore import pyqtSlot
import numpy as np
import pandas as pd
from datetime import datetime
import pytz

class ChartWidget(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.layout = QVBoxLayout()
        self.setLayout(self.layout)
        
        # Initialize PyQtGraph layout
        self.glw = pg.GraphicsLayoutWidget()
        self.layout.addWidget(self.glw)
        
        # Col 0: Profile (Count vs Price)
        self.profile_plot = self.glw.addPlot(row=0, col=0)
        self.profile_plot.setMaximumWidth(200)
        self.profile_plot.hideAxis('bottom')
        self.profile_plot.showAxis('top')
        self.profile_plot.setLabel('top', 'Volume')
        self.profile_plot.setClipToView(True)
        # self.profile_plot.setDownsampling(auto=True, mode='peak')

        # Col 1: Price (Time vs Price)
        self.price_plot = self.glw.addPlot(row=0, col=1)
        self.price_plot.setLabel('bottom', 'Time')
        self.price_plot.setLabel('right', 'Price')
        self.price_plot.showAxis('right')
        self.price_plot.hideAxis('left')
        self.price_plot.showAxis('right')
        self.price_plot.hideAxis('left')
        self.price_plot.setClipToView(False) # Disable for debugging
        # self.price_plot.setDownsampling(auto=True, mode='peak') # Disable for debugging
        
        # Link Y-axes
        # self.profile_plot.setYLink(self.price_plot) # temporarily unlink to rule out profile plot issues
        
        # Initialize Chart Items
        self.bid_curve = self.price_plot.plot(pen=pg.mkPen('b', width=1), name="Bid")
        self.ask_curve = self.price_plot.plot(pen=pg.mkPen('r', width=1), name="Ask")
        
        # Profile Histogram Item
        self.profile_bars = pg.BarGraphItem(x0=0, y=0, width=0, height=0.01, brush='c')
        self.profile_plot.addItem(self.profile_bars)
        
        # Developing Lines (PlotCurveItems)
        self.curve_vah = self.price_plot.plot(pen=pg.mkPen('g', width=2), name="VAH")
        self.curve_val = self.price_plot.plot(pen=pg.mkPen('m', width=2), name="VAL")
        self.curve_poc = self.price_plot.plot(pen=pg.mkPen('y', width=2), name="POC")

        # Date axis formatter
        self.date_axis = self.price_plot.getAxis('bottom')
        # self.date_axis.setTickSpacing(3600, 1800) # Grid every hour - Caused MemoryError
        self.price_plot.showGrid(x=True, y=True, alpha=0.3)

        # Data storage
        self.times = np.array([])
        self.bids = np.array([])
        self.asks = np.array([])
        
        # Storage for developing levels
        self.level_times = np.array([])
        self.level_vah = np.array([])
        self.level_val = np.array([])
        self.level_poc = np.array([])

    def clear(self):
        self.times = np.array([])
        self.bids = np.array([])
        self.asks = np.array([])
        self.level_times = np.array([])
        self.level_vah = np.array([])
        self.level_val = np.array([])
        self.level_poc = np.array([])
        
        self.bid_curve.setData([], [])
        self.ask_curve.setData([], [])
        self.profile_bars.setOpts(x0=[], y=[], width=[], height=[])
        self.curve_vah.setData([], [])
        self.curve_val.setData([], [])
        self.curve_poc.setData([], [])

    def update_ticks(self, df):
        """
        Updates the tick chart by appending new data.
        df: DataFrame with 'datetime' (ns timestamp) and 'bid', 'ask'.
        """
        if df.empty:
            return

        # Convert timestamps for pyqtgraph (seconds since epoch)
        # df['datetime'] is numpy datetime64[ns]
        new_times = df['datetime'].values.astype(np.float64) / 1e9 
        new_bids = df['bid'].values
        new_asks = df['ask'].values if 'ask' in df.columns else np.zeros_like(new_bids)
        
        if len(self.times) == 0:
            self.times = new_times
            self.bids = new_bids
            self.asks = new_asks
        else:
            self.times = np.concatenate((self.times, new_times))
            self.bids = np.concatenate((self.bids, new_bids))
            self.asks = np.concatenate((self.asks, new_asks))
        
        # Debug Log
        if len(self.times) > 0:
            t_min, t_max = self.times[0], self.times[-1]
            b_min, b_max = np.min(self.bids), np.max(self.bids)
            print(f"DEBUG: Chart Ticks: {len(self.times)} pts.")
            print(f"DEBUG: Time Range: {t_min:.1f} -> {t_max:.1f} ({datetime.fromtimestamp(t_min)} -> {datetime.fromtimestamp(t_max)})")
            print(f"DEBUG: Price Range: {b_min:.4f} -> {b_max:.4f}")
        
        # Update curves
        self.bid_curve.setData(self.times, self.bids)
        if len(self.asks) > 0:
            self.ask_curve.setData(self.times, self.asks)
        
        # Force range on first large update
        if len(self.times) > 0 and len(self.times) == len(new_times): 
             self.price_plot.setXRange(self.times[0], self.times[-1], padding=0.02)
             self.price_plot.setYRange(np.min(self.bids), np.max(self.bids), padding=0.02)

    def update_profile(self, counts_dict, unit_size=0.01):
        """
        Updates the side profile histogram.
        counts: dict {price: count}
        """
        if not counts_dict:
            return

        prices = np.array(list(counts_dict.keys()))
        counts = np.array(list(counts_dict.values()))
        
        # Horizontal bars: x0=0, y=prices, width=counts, height=unit_size
        self.profile_bars.setOpts(x0=np.zeros(len(prices)), y=prices, width=counts, height=unit_size, brush=(0, 255, 255, 100))

    def update_levels(self, new_times, new_vah, new_val, new_poc):
        """
        Updates the developing VAH/VAL/POC lines.
        Expects arrays or scalars.
        """
        try:
            # Ensure inputs are 1D arrays
            nt = np.atleast_1d(np.array(new_times, dtype=np.float64))
            nv = np.atleast_1d(np.array(new_vah, dtype=np.float64))
            nl = np.atleast_1d(np.array(new_val, dtype=np.float64))
            yp = np.atleast_1d(np.array(new_poc, dtype=np.float64))
            
            if len(nt) == 0:
                # print("Chart Update Levels: Empty new_times")
                return

            # Append logic
            if len(self.level_times) == 0:
                self.level_times = nt
                self.level_vah = nv
                self.level_val = nl
                self.level_poc = yp
            else:
                self.level_times = np.concatenate((self.level_times, nt))
                self.level_vah = np.concatenate((self.level_vah, nv))
                self.level_val = np.concatenate((self.level_val, nl))
                self.level_poc = np.concatenate((self.level_poc, yp))
            
            if len(nt) > 1:
                print(f"DEBUG: Chart Levels Loaded: {len(nt)} pts. POC Range: {self.level_poc[0]:.2f} -> {self.level_poc[-1]:.2f}")
            
            # Update plots
            self.curve_vah.setData(self.level_times, self.level_vah)
            self.curve_val.setData(self.level_times, self.level_val)
            self.curve_poc.setData(self.level_times, self.level_poc)
        except Exception as e:
            print(f"Error updating levels: {e}")