'use client'; import { useEffect, useRef, useState, useImperativeHandle, forwardRef } from 'react'; import { createChart, IChartApi, ISeriesApi, CandlestickData, HistogramData, Time, SeriesMarker } from 'lightweight-charts'; import { useTheme } from 'next-themes'; import SpanAnnotationManager from './SpanAnnotationManager'; import { TrendLine } from '@/plugins/trend-line'; import { RectangleDrawingPrimitive, RectanglePoint } from '@/plugins/rectangle-drawing'; import type { PerCandlePrediction, PredictionSpan, ModelInfoResponse, PredictionSummary } from '@/types/predictions'; interface DrawingState { tool: 'line' | 'rectangle'; firstPoint: { time: Time; price: number }; } interface Point { time: Time; price: number; } interface Candle { time: number; open: number; high: number; low: number; close: number; } interface Annotation { id: number; timestamp: number; label_type: string; geometry: { startTime?: number; startPrice?: number; endTime?: number; endPrice?: number; } | null; created_at: number; } type AnnotationType = { id: number; name: string; display_name: string; color: string; category: string; icon: string | null; is_active: boolean; }; interface SpanAnnotation { id: number; chart_id: number; start_time: number; end_time: number; label: string; confidence: number | null; outcome: string | null; notes: string | null; sub_spans: any; color: string; created_at: number; } interface SpanLabelType { id: number; name: string; display_name: string; color: string; hotkey: string | null; is_active: boolean; sort_order: number; created_at: number; } interface CandleChartProps { activeTool: string | null; onAnnotationChange?: () => void; selectedColor: string; selectedLabelId?: number | null; onLabelSelect?: (id: number) => void; activeChartId?: number | null; spanAnnotations?: SpanAnnotation[]; spanLabelTypes?: SpanLabelType[]; selectedSpanId?: number | null; onSpanAnnotationsChange?: () => void; onSelectedSpanChange?: (spanId: number | null) => void; // Prediction rendering props predictionVisible?: boolean; perCandlePredictions?: PerCandlePrediction[]; predictionSpans?: PredictionSpan[]; confidenceThreshold?: number; selectedLabels?: Set; modelInfo?: ModelInfoResponse | null; predictionSummary?: PredictionSummary | null; showOnlyDisagreements?: boolean; onPredictionClick?: (span: PredictionSpan, disagreementType: string | null) => void; onPredictionDismiss?: (span: PredictionSpan, disagreementType: string | null) => void; } export interface CandleChartHandle { refreshData: () => Promise; getVisibleCandles: () => Candle[]; } const CandleChart = forwardRef( ({ activeTool, onAnnotationChange, selectedColor, selectedLabelId, onLabelSelect, activeChartId, spanAnnotations = [], spanLabelTypes = [], selectedSpanId = null, onSpanAnnotationsChange, onSelectedSpanChange, predictionVisible = false, perCandlePredictions = [], predictionSpans = [], confidenceThreshold = 0.5, selectedLabels = new Set(), modelInfo = null, predictionSummary = null, showOnlyDisagreements = false, onPredictionClick, onPredictionDismiss, }, ref) => { const chartContainerRef = useRef(null); const chartRef = useRef(null); const seriesRef = useRef | null>(null); const histogramSeriesRef = useRef | null>(null); const [candles, setCandles] = useState([]); const [annotations, setAnnotations] = useState([]); const [annotationTypes, setAnnotationTypes] = useState([]); const [isEmpty, setIsEmpty] = useState(true); const { theme, resolvedTheme } = useTheme(); const [mounted, setMounted] = useState(false); // Drawing state for line and rectangle tools const [drawingState, setDrawingState] = useState(null); const previewPrimitiveRef = useRef(null); const linePrimitivesRef = useRef>(new Map()); const rectanglePrimitivesRef = useRef>(new Map()); const [selectedRectangleId, setSelectedRectangleId] = useState(null); // Line selection and dragging state const [selectedLineId, setSelectedLineId] = useState(null); const [dragState, setDragState] = useState<{ lineId: number; endpoint: 'p1' | 'p2' } | null>(null); // Track mounted state to avoid hydration mismatch useEffect(() => { setMounted(true); }, []); // Fetch candles from API const fetchCandles = async () => { try { const url = activeChartId ? `/api/candles?chartId=${activeChartId}` : '/api/candles'; const response = await fetch(url); const data = await response.json(); setCandles(data); setIsEmpty(data.length === 0); return data; } catch (error) { console.error('Failed to fetch candles:', error); return []; } }; // Fetch annotations from API const fetchAnnotations = async () => { try { const url = activeChartId ? `/api/annotations?chartId=${activeChartId}` : '/api/annotations'; const response = await fetch(url); const data = await response.json(); setAnnotations(data); return data; } catch (error) { console.error('Failed to fetch annotations:', error); return []; } }; // Fetch annotation types from API const fetchAnnotationTypes = async () => { try { const response = await fetch('/api/annotation-types'); const data = await response.json(); setAnnotationTypes(data.filter((t: AnnotationType) => t.is_active)); return data; } catch (error) { console.error('Failed to fetch annotation types:', error); return []; } }; // Expose methods to parent useImperativeHandle(ref, () => ({ refreshData: async () => { await fetchCandles(); await fetchAnnotations(); await fetchAnnotationTypes(); }, getVisibleCandles: () => { if (!chartRef.current || candles.length === 0) return []; const timeScale = chartRef.current.timeScale(); const visibleRange = timeScale.getVisibleLogicalRange(); if (!visibleRange) return candles; const barsInfo = seriesRef.current?.barsInLogicalRange(visibleRange); if (!barsInfo || !barsInfo.barsBefore && !barsInfo.barsAfter) return candles; // Filter candles by visible time range const from = barsInfo.from as number; const to = barsInfo.to as number; if (typeof from !== 'number' || typeof to !== 'number') return candles; return candles.filter((c) => c.time >= from && c.time <= to); }, })); // Get theme-specific colors const getChartColors = () => { const isDark = resolvedTheme === 'dark'; if (isDark) { return { layout: { background: { color: '#000000' }, textColor: '#00ff41', }, grid: { vertLines: { color: '#003311' }, horzLines: { color: '#003311' }, }, timeScale: { borderColor: '#003311', }, rightPriceScale: { borderColor: '#003311', }, crosshair: { color: '#00ff41', }, }; } else { return { layout: { background: { color: '#ffffff' }, textColor: '#1a1a2e', }, grid: { vertLines: { color: '#d1d5db' }, horzLines: { color: '#d1d5db' }, }, timeScale: { borderColor: '#d1d5db', }, rightPriceScale: { borderColor: '#d1d5db', }, crosshair: { color: '#16a34a', }, }; } }; // Initialize chart useEffect(() => { if (!chartContainerRef.current || isEmpty || !mounted) return; const colors = getChartColors(); const chart = createChart(chartContainerRef.current, { width: chartContainerRef.current.clientWidth, height: chartContainerRef.current.clientHeight, layout: colors.layout, grid: colors.grid, crosshair: { mode: 1, }, timeScale: { timeVisible: true, secondsVisible: false, borderColor: colors.timeScale.borderColor, }, rightPriceScale: { borderColor: colors.rightPriceScale.borderColor, }, }); const candlestickSeries = chart.addCandlestickSeries({ upColor: '#ffffff', downColor: '#444444', borderVisible: true, wickUpColor: '#444444', wickDownColor: '#444444', borderUpColor: '#444444', borderDownColor: '#444444', }); // Add histogram series for prediction overlay (rendered behind candles) const histogramSeries = chart.addHistogramSeries({ color: 'rgba(128, 128, 128, 0.15)', priceFormat: { type: 'volume' }, priceScaleId: 'prediction_overlay', }); histogramSeries.priceScale().applyOptions({ scaleMargins: { top: 0, bottom: 0 }, visible: false, }); chartRef.current = chart; seriesRef.current = candlestickSeries; histogramSeriesRef.current = histogramSeries; // Handle resize const resizeObserver = new ResizeObserver((entries) => { if (entries.length === 0 || !chartContainerRef.current) return; const { width, height } = entries[0].contentRect; chart.applyOptions({ width, height }); }); resizeObserver.observe(chartContainerRef.current); return () => { resizeObserver.disconnect(); histogramSeriesRef.current = null; chart.remove(); }; }, [isEmpty, mounted, resolvedTheme]); // Load candle data into chart useEffect(() => { if (!seriesRef.current || candles.length === 0) return; const chartData: CandlestickData[] = candles .map((candle) => ({ time: candle.time as Time, open: candle.open, high: candle.high, low: candle.low, close: candle.close, })) .sort((a, b) => (a.time as number) - (b.time as number)); seriesRef.current.setData(chartData); }, [candles]); // Update markers from annotations useEffect(() => { if (!seriesRef.current || annotationTypes.length === 0) return; // Get marker type names const markerTypeNames = annotationTypes .filter((t) => t.category === 'marker') .map((t) => t.name); const markerAnnotations = annotations.filter((a) => markerTypeNames.includes(a.label_type) ); const markers = markerAnnotations .map((annotation) => { const annotationType = annotationTypes.find((t) => t.name === annotation.label_type); if (!annotationType) return null; const isSelected = annotation.id === selectedLabelId; // Determine marker shape and position based on icon const isUpArrow = annotationType.icon === 'arrowUp'; return { time: annotation.timestamp as Time, position: isUpArrow ? ('belowBar' as const) : ('aboveBar' as const), color: isSelected ? annotationType.color + 'CC' // Add slight transparency when selected : annotationType.color, shape: isUpArrow ? ('arrowUp' as const) : ('arrowDown' as const), text: annotationType.display_name, size: isSelected ? 2 : 1, }; }) .filter((m) => m !== null) .sort((a, b) => (a!.time as number) - (b!.time as number)); seriesRef.current.setMarkers(markers as any); }, [annotations, selectedLabelId, annotationTypes]); // Render prediction histogram overlay and span markers useEffect(() => { if (!histogramSeriesRef.current || !seriesRef.current || candles.length === 0) return; // If predictions are not visible or no data, clear the histogram and markers if (!predictionVisible || perCandlePredictions.length === 0) { histogramSeriesRef.current.setData([]); histogramSeriesRef.current.setMarkers([]); return; } // Build a label-to-color map from modelInfo // Generate colors for labels since backend no longer provides them const labelColorMap: Record = {}; if (modelInfo?.labels) { const predefinedColors = [ '#3b82f6', // blue '#ef4444', // red '#10b981', // green '#f59e0b', // amber '#8b5cf6', // violet '#ec4899', // pink '#06b6d4', // cyan '#f97316', // orange ]; modelInfo.labels.forEach((label, index) => { labelColorMap[label] = predefinedColors[index % predefinedColors.length]; }); } // Build disagreement lookup map for highlighting // Map: predictionSpan start_time -> disagreement type const disagreementMap = new Map(); const disagreementSpanSet = new Set(); // Set of "start-end" keys for filtering if (predictionSummary?.disagreements) { predictionSummary.disagreements.forEach((d) => { if (d.predictionSpan) { disagreementMap.set(d.predictionSpan.start_time, d.type); disagreementSpanSet.add(`${d.predictionSpan.start_time}-${d.predictionSpan.end_time}`); } }); } // Helper to convert hex to rgba const hexToRgba = (hex: string, alpha: number): string => { hex = hex.replace('#', ''); const r = parseInt(hex.substring(0, 2), 16); const g = parseInt(hex.substring(2, 4), 16); const b = parseInt(hex.substring(4, 6), 16); return `rgba(${r}, ${g}, ${b}, ${alpha})`; }; // Build candle price lookup for histogram values const candleMap = new Map(candles.map((c) => [c.time, c])); // Build a map from prediction time to its span for disagreement lookup const predictionTimeToSpan = new Map(); predictionSpans.forEach((span) => { // Associate all times in the span with the span object for (let t = span.start_time; t <= span.end_time; t += 60) { // Assuming 1-minute candles, adjust if needed predictionTimeToSpan.set(t, span); } }); // Filter and map predictions to histogram data with disagreement highlighting const histogramData: (HistogramData & { color?: string })[] = perCandlePredictions .filter((p) => { // Filter by confidence threshold if (p.confidence < confidenceThreshold) return false; // Filter by selected labels if (!selectedLabels.has(p.label)) return false; // Skip "O" (no-pattern) labels if (p.label === 'O') return false; // If showOnlyDisagreements is enabled, only show predictions that are part of disagreements if (showOnlyDisagreements) { const span = predictionTimeToSpan.get(p.time); if (!span) return false; const spanKey = `${span.start_time}-${span.end_time}`; if (!disagreementSpanSet.has(spanKey)) return false; } return true; }) .map((p) => { const candle = candleMap.get(p.time); // Use candle high as the histogram value so it overlays correctly const value = candle ? candle.high : 0; // Check if this prediction is part of a disagreement const span = predictionTimeToSpan.get(p.time); const disagreementType = span ? disagreementMap.get(span.start_time) : null; let baseColor = labelColorMap[p.label] || '#888888'; let alpha = 0.15; // Apply disagreement-specific colors and styling if (disagreementType === 'missed_by_human') { // Yellow highlight for predictions missed by humans baseColor = '#eab308'; // yellow alpha = 0.25; } else if (disagreementType === 'label_mismatch') { // Orange for label mismatches baseColor = '#f97316'; // orange alpha = 0.25; } return { time: p.time as Time, value, color: hexToRgba(baseColor, alpha), }; }) .sort((a, b) => (a.time as number) - (b.time as number)); histogramSeriesRef.current.setData(histogramData); // Add markers for prediction span labels at span start times const spanMarkers: SeriesMarker