ui/plotly_charts.py

"""
Plotly chart builders for CSH2 Web Dashboard — Dark HUD / Brutalist Industrial Theme
Plotly can't read CSS variables, so we mirror the :root tokens as Python constants.
These MUST stay in sync with ui/styles.py :root block.
"""
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import pandas as pd
import numpy as np
from typing import List, Dict, Optional, Tuple


# ── Token mirror (sync with :root in styles.py) ──
_BG_BASE      = '#0C0A08'
_BG_PANEL     = '#12100E'
_BG_ELEVATED  = '#1A1714'
_BORDER       = '#2A2520'
_TEXT_PRIMARY  = '#E8E0D4'
_TEXT_SECONDARY= '#B0A898'
_TEXT_MUTED    = '#6B6358'
_ACCENT       = '#D4A04A'
_FONT_MONO    = 'JetBrains Mono, monospace'
_FONT_BODY    = 'DM Sans, sans-serif'

# Chart color palette (8 series + 2 alternates)
COLORS = [
    '#D4A04A',  # Amber (primary)
    '#5B9A6E',  # Green (nominal)
    '#5BA3B5',  # Teal (informational)
    '#CC3030',  # Red (alarm)
    '#E8C47A',  # Light amber
    '#7DB88E',  # Light green
    '#7BBFCC',  # Light teal
    '#E06060',  # Light red
    '#C4A35A',  # Gold
    '#A0856A',  # Warm brown
]

# Downsample threshold for charts (browser performance)
MAX_CHART_POINTS = 5000

# Reusable HUD layout config for all Plotly charts
HUD_LAYOUT = dict(
    template='plotly_dark',
    paper_bgcolor=_BG_BASE,
    plot_bgcolor=_BG_BASE,
    font=dict(family=_FONT_BODY, color=_TEXT_SECONDARY, size=11),
    title_font=dict(family=_FONT_MONO, color=_ACCENT, size=14),
    hovermode='x unified',
    hoverlabel=dict(
        bgcolor=_BG_PANEL,
        bordercolor=_ACCENT,
        font=dict(family=_FONT_MONO, color=_TEXT_PRIMARY, size=11),
    ),
    legend=dict(
        orientation='h',
        yanchor='bottom',
        y=1.02,
        xanchor='right',
        x=1,
        font=dict(family=_FONT_BODY, size=10, color=_TEXT_SECONDARY),
        bgcolor=f'rgba(12, 10, 8, 0.8)',
        bordercolor=_BORDER,
        borderwidth=1,
    ),
)

# Reusable HUD axis config
HUD_AXIS = dict(
    gridcolor=_BG_ELEVATED,
    zerolinecolor=_BORDER,
    tickfont=dict(family=_FONT_MONO, color=_TEXT_MUTED, size=10),
    title_font=dict(family=_FONT_BODY, color=_TEXT_SECONDARY, size=11),
)


def _downsample_df(
    df: pd.DataFrame,
    tags: List[str],
    max_points: int = MAX_CHART_POINTS,
) -> Tuple[pd.DataFrame, bool]:
    """
    Downsample a DataFrame for charting using min/max bucketing.
    Preserves peaks and valleys while reducing point count.

    Returns (downsampled_df, was_downsampled).
    """
    if len(df) <= max_points:
        return df, False

    n_buckets = max_points // 2  # each bucket yields ~2 rows (min + max)
    bucket_size = max(len(df) // n_buckets, 2)

    df = df.reset_index(drop=True)
    df['_bucket'] = df.index // bucket_size

    result_indices = {0, len(df) - 1}  # always keep first and last

    # Find the primary tag to drive bucket selection
    primary_tag = None
    for tag in tags:
        if tag in df.columns:
            primary_tag = tag
            break

    if primary_tag is None:
        # Fallback: uniform sampling
        step = max(len(df) // max_points, 1)
        return df.iloc[::step].drop(columns=['_bucket'], errors='ignore'), True

    for _, group in df.groupby('_bucket'):
        if len(group) == 0:
            continue
        col = group[primary_tag].dropna()
        if len(col) > 0:
            result_indices.add(col.idxmin())
            result_indices.add(col.idxmax())

    result = df.loc[sorted(result_indices)].drop(columns=['_bucket'])
    return result, True


def create_timeseries_chart(
    df_pivot: pd.DataFrame,
    tags: List[str],
    title: str = "Sensor Data Over Time",
    height: int = 500,
    time_range: Optional[Tuple] = None,
) -> go.Figure:
    """Create an interactive time series chart from pivoted data"""
    fig = go.Figure()

    # Downsample for browser performance
    df_plot, downsampled = _downsample_df(df_pivot, tags)

    for i, tag in enumerate(tags):
        if tag in df_plot.columns:
            fig.add_trace(go.Scatter(
                x=df_plot['timestamp'],
                y=df_plot[tag],
                mode='lines',
                name=tag,
                line=dict(color=COLORS[i % len(COLORS)], width=1.5),
                hovertemplate=f'{tag}: %{{y:.2f}}<br>%{{x}}<extra></extra>',
            ))

    ds_note = f" (downsampled {len(df_pivot):,} \u2192 {len(df_plot):,} pts)" if downsampled else ""

    # Compute explicit x-axis range
    xaxis_range = None
    if time_range is not None:
        _x_min, _x_max = pd.Timestamp(time_range[0]), pd.Timestamp(time_range[1])
        duration = (_x_max - _x_min).total_seconds()
        margin = pd.Timedelta(seconds=max(duration * 0.02, 1))
        xaxis_range = [_x_min - margin, _x_max + margin]
    elif 'timestamp' in df_plot.columns and len(df_plot) > 0:
        _x_min = df_plot['timestamp'].min()
        _x_max = df_plot['timestamp'].max()
        duration = (_x_max - _x_min).total_seconds()
        margin = pd.Timedelta(seconds=max(duration * 0.02, 1))
        xaxis_range = [_x_min - margin, _x_max + margin]

    xaxis_cfg = dict(
        title='Time (ET)',
        rangeslider=dict(visible=True, thickness=0.05, bgcolor='#12100E'),
        type='date',
        **HUD_AXIS,
    )
    if xaxis_range is not None:
        xaxis_cfg['range'] = xaxis_range

    fig.update_layout(
        **HUD_LAYOUT,
        title=dict(text=title + ds_note, font=HUD_LAYOUT['title_font']),
        xaxis=xaxis_cfg,
        yaxis=dict(title='Value', **HUD_AXIS),
        height=height,
        margin=dict(l=60, r=20, t=60, b=40),
    )

    return fig


def create_multi_axis_chart(
    df: pd.DataFrame,
    pressure_tags: List[str],
    temp_tags: List[str],
    other_tags: List[str] = None,
    title: str = "Cycle Detail",
    height: int = 600,
    plateaus: Dict[str, List[Dict]] = None,
    time_range: Optional[Tuple] = None,
) -> go.Figure:
    """
    Create multi-axis chart for cycle detail view.
    Left Y: Pressure, Right Y: Temperature, Subplot: Other sensors
    """
    has_other = other_tags and any(t in df.columns for t in other_tags)
    rows = 2 if has_other else 1

    # Downsample for browser performance
    all_chart_tags = pressure_tags + temp_tags + (other_tags or [])
    df_plot, _ = _downsample_df(df, all_chart_tags)

    fig = make_subplots(
        rows=rows, cols=1,
        shared_xaxes=True,
        vertical_spacing=0.12,
        specs=[[{"secondary_y": True}]] + ([[{"secondary_y": False}]] if has_other else []),
    )

    color_idx = 0

    # Pressure traces (left Y)
    for tag in pressure_tags:
        if tag in df_plot.columns:
            fig.add_trace(go.Scatter(
                x=df_plot['timestamp'], y=df_plot[tag],
                mode='lines', name=tag,
                line=dict(color=COLORS[color_idx % len(COLORS)], width=2),
                hovertemplate=f'{tag}: %{{y:.1f}} bar<extra></extra>',
            ), row=1, col=1, secondary_y=False)
            color_idx += 1

    # Temperature traces (right Y)
    for tag in temp_tags:
        if tag in df_plot.columns:
            fig.add_trace(go.Scatter(
                x=df_plot['timestamp'], y=df_plot[tag],
                mode='lines', name=tag,
                line=dict(color=COLORS[color_idx % len(COLORS)], width=2, dash='dash'),
                hovertemplate=f'{tag}: %{{y:.1f}} K<extra></extra>',
            ), row=1, col=1, secondary_y=True)
            color_idx += 1

    # Other sensor traces (subplot 2)
    if has_other:
        for tag in other_tags:
            if tag in df_plot.columns:
                fig.add_trace(go.Scatter(
                    x=df_plot['timestamp'], y=df_plot[tag],
                    mode='lines', name=tag,
                    line=dict(color=COLORS[color_idx % len(COLORS)], width=1.5),
                    hovertemplate=f'{tag}: %{{y:.2f}}<extra></extra>',
                ), row=2, col=1)
                color_idx += 1

    # Compute x-axis bounds for clamping vrects and setting explicit range
    if time_range is not None:
        _x_min, _x_max = pd.Timestamp(time_range[0]), pd.Timestamp(time_range[1])
    elif 'timestamp' in df_plot.columns and len(df_plot) > 0:
        _x_min = df_plot['timestamp'].min()
        _x_max = df_plot['timestamp'].max()
    else:
        _x_min = _x_max = None

    # Add plateau highlights (clamped to data range)
    if plateaus:
        for tag, periods in plateaus.items():
            color = 'rgba(212, 160, 74, 0.1)' if 'PT' in tag or 'pressure' in tag.lower() else 'rgba(91, 154, 110, 0.1)'
            for p in periods:
                vx0, vx1 = pd.Timestamp(p['start']), pd.Timestamp(p['end'])
                # Normalize timezone awareness to match _x_min/_x_max
                if _x_min is not None:
                    if _x_min.tzinfo is not None and vx0.tzinfo is None:
                        vx0 = vx0.tz_localize(_x_min.tzinfo)
                        vx1 = vx1.tz_localize(_x_min.tzinfo)
                    elif _x_min.tzinfo is None and vx0.tzinfo is not None:
                        vx0 = vx0.tz_localize(None)
                        vx1 = vx1.tz_localize(None)
                # Clamp vrect bounds to prevent x-axis expansion
                if _x_min is not None:
                    vx0 = max(vx0, _x_min)
                if _x_max is not None:
                    vx1 = min(vx1, _x_max)
                if _x_min is not None and vx1 < _x_min:
                    continue
                if _x_max is not None and vx0 > _x_max:
                    continue
                fig.add_vrect(
                    x0=vx0, x1=vx1,
                    fillcolor=color,
                    layer='below',
                    line_width=0,
                    row=1, col=1,
                    annotation_text=f"{tag} plateau",
                    annotation_position="top left",
                    annotation_font_size=9,
                    annotation_font_color='#6B6358',
                )

    fig.update_layout(
        **HUD_LAYOUT,
        title=dict(text=title, font=HUD_LAYOUT['title_font']),
        height=height,
        margin=dict(l=60, r=60, t=120, b=40),
    )
    # Override legend position separately to avoid duplicate 'legend' kwarg
    # (HUD_LAYOUT already contains 'legend', so passing it again raises TypeError)
    fig.update_layout(legend_y=1.10)

    fig.update_yaxes(title_text="Pressure (bar)", row=1, col=1, secondary_y=False, **HUD_AXIS)
    fig.update_yaxes(title_text="Temperature (K)", row=1, col=1, secondary_y=True, **HUD_AXIS)
    if has_other:
        fig.update_yaxes(title_text="Motor & Flow", row=2, col=1, **HUD_AXIS)

    # Set explicit x-axis range to prevent Plotly auto-expansion
    if _x_min is not None and _x_max is not None:
        duration = (_x_max - _x_min).total_seconds()
        margin = pd.Timedelta(seconds=max(duration * 0.02, 1))
        fig.update_xaxes(range=[_x_min - margin, _x_max + margin], **HUD_AXIS)
    else:
        fig.update_xaxes(**HUD_AXIS)

    return fig


def create_cycle_timeline(cycles: List[Dict], month: str, height: int = 300) -> go.Figure:
    """Create a Gantt-style timeline of testing cycles in a month"""
    fig = go.Figure()

    for i, c in enumerate(cycles):
        color = COLORS[i % len(COLORS)]
        # Use Eastern display times for hover, fall back to start_time
        _start_et = c.get('start_time_et') or c['start_time']
        _end_et = c.get('end_time_et') or c['end_time']
        fig.add_trace(go.Bar(
            x=[c['duration_minutes']],
            y=[f"Cycle {c['cycle_id']}"],
            orientation='h',
            marker=dict(color=color, line=dict(width=1, color='#2A2520')),
            text=f"{c['duration_minutes']:.0f} min | Peak (PT130): {c.get('peak_pressure', 0):.0f} bar",
            textposition='inside',
            textfont=dict(color='#0C0A08', size=10),
            hovertemplate=(
                f"Cycle {c['cycle_id']}<br>"
                f"Start: {_start_et.strftime('%b %d %H:%M')} ET<br>"
                f"End: {_end_et.strftime('%b %d %H:%M')} ET<br>"
                f"Duration: {c['duration_minutes']:.1f} min<br>"
                f"Peak Pressure (PT130): {c.get('peak_pressure', 0):.0f} bar"
                "<extra></extra>"
            ),
            showlegend=False,
        ))

    fig.update_layout(
        **HUD_LAYOUT,
        title=dict(text=f"TESTING CYCLES — {month}", font=HUD_LAYOUT['title_font']),
        xaxis=dict(title='Duration (minutes)', **HUD_AXIS),
        yaxis=dict(autorange='reversed', **HUD_AXIS),
        height=height,
        margin=dict(l=80, r=20, t=50, b=40),
    )

    return fig


def create_comparison_chart(
    df_a: pd.DataFrame,
    df_b: pd.DataFrame,
    tag: str,
    label_a: str = "Cycle A",
    label_b: str = "Cycle B",
    height: int = 400,
) -> go.Figure:
    """Overlay two cycles on normalized time axis for comparison"""
    fig = go.Figure()

    if tag in df_a.columns:
        # Normalize time to minutes from start
        t0_a = df_a['timestamp'].min()
        minutes_a = (df_a['timestamp'] - t0_a).dt.total_seconds() / 60
        fig.add_trace(go.Scatter(
            x=minutes_a, y=df_a[tag],
            mode='lines', name=f"{label_a} — {tag}",
            line=dict(color=COLORS[0], width=2),
        ))

    if tag in df_b.columns:
        t0_b = df_b['timestamp'].min()
        minutes_b = (df_b['timestamp'] - t0_b).dt.total_seconds() / 60
        fig.add_trace(go.Scatter(
            x=minutes_b, y=df_b[tag],
            mode='lines', name=f"{label_b} — {tag}",
            line=dict(color=COLORS[1], width=2),
        ))

    fig.update_layout(
        **HUD_LAYOUT,
        title=dict(text=f"{tag} — CYCLE COMPARISON", font=HUD_LAYOUT['title_font']),
        xaxis=dict(title='Minutes from Cycle Start', **HUD_AXIS),
        yaxis=dict(title=tag, **HUD_AXIS),
        height=height,
        margin=dict(l=60, r=20, t=50, b=40),
    )

    return fig