dcrm/plotting.py

import plotly.graph_objects as go
from plotly.subplots import make_subplots

def create_dcrm_plot(df, zones):
    # Create figure with secondary y-axis
    fig = make_subplots(specs=[[{"secondary_y": True}]])

    # Add Traces
    # Ensure column names match what is in the DF (dcrm_llm_app.py uses 'Time (ms)', 'Current', 'Resistance', 'Travel')
    # The user's code uses 'Time_ms'. I need to be careful here.
    # process_uploaded_image returns DF with 'Time (ms)', 'Current', 'Resistance', 'Travel'.
    # I should adapt the plotting code to use the existing column names OR rename columns in the DF.
    # The user's code expects 'Time_ms'.
    # I will handle this mapping inside the function to be safe.
    
    time_col = 'Time (ms)' if 'Time (ms)' in df.columns else 'Time_ms'
    
    fig.add_trace(go.Scatter(x=df[time_col], y=df['Current'], name="Current (A)", line=dict(color='#2980b9', width=2)), secondary_y=False)
    fig.add_trace(go.Scatter(x=df[time_col], y=df['Resistance'], name="Resistance (uOhm)", line=dict(color='#27ae60', width=2)), secondary_y=False)
    fig.add_trace(go.Scatter(x=df[time_col], y=df['Travel'], name="Travel (mm)", line=dict(color='#c0392b', width=2)), secondary_y=True)

    # Zone Colors
    zone_colors = {
        "zone_1_pre_contact": "rgba(52, 152, 219, 0.1)",
        "zone_2_arcing_engagement": "rgba(231, 76, 60, 0.1)",
        "zone_3_main_conduction": "rgba(46, 204, 113, 0.1)",
        "zone_4_parting": "rgba(155, 89, 182, 0.1)",
        "zone_5_final_open": "rgba(149, 165, 166, 0.1)"
    }

    # Add Zone Rectangles
    # The user's code expects 'zones' to be a dict of zone details.
    # The result_json has "zones" key.
    zones_dict = zones.get("zones", {}) if "zones" in zones else zones

    for zone_name, details in zones_dict.items():
        start = details.get("start_ms")
        end = details.get("end_ms")
        color = zone_colors.get(zone_name, "rgba(0,0,0,0)")
        
        if start is not None and end is not None:
            fig.add_vrect(
                x0=start, x1=end,
                fillcolor=color, opacity=1,
                layer="below", line_width=0,
                annotation_text=zone_name.split('_')[1].upper(),
                annotation_position="top left",
                annotation_font_color="#7f8c8d"
            )

    fig.update_layout(
        title_text="<b>Main Signals & Zones</b>",
        height=500,
        hovermode="x unified",
        plot_bgcolor="white",
        paper_bgcolor="white",
        font=dict(family="Segoe UI, sans-serif"),
        legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
        margin=dict(l=20, r=20, t=60, b=20)
    )
    fig.update_xaxes(showgrid=True, gridwidth=1, gridcolor='#f0f0f0')
    fig.update_yaxes(title_text="Current / Resistance", secondary_y=False, showgrid=True, gridwidth=1, gridcolor='#f0f0f0')
    fig.update_yaxes(title_text="Travel", secondary_y=True, showgrid=False)

    return fig

def create_velocity_plot(df):
    time_col = 'Time (ms)' if 'Time (ms)' in df.columns else 'Time_ms'
    
    # Calculate Velocity (Derivative of Travel)
    # V = d(Travel) / d(Time)
    # Units: mm/ms = m/s
    df['Velocity'] = df['Travel'].diff() / df[time_col].diff()
    
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=df[time_col], y=df['Velocity'], name="Velocity (m/s)", line=dict(color='#e67e22', width=2), fill='tozeroy'))
    
    fig.update_layout(
        title_text="<b>Contact Velocity Profile</b>",
        height=300,
        hovermode="x unified",
        plot_bgcolor="white",
        paper_bgcolor="white",
        font=dict(family="Segoe UI, sans-serif"),
        margin=dict(l=20, r=20, t=40, b=20)
    )
    fig.update_xaxes(showgrid=True, gridwidth=1, gridcolor='#f0f0f0')
    fig.update_yaxes(title_text="Velocity (m/s)", showgrid=True, gridwidth=1, gridcolor='#f0f0f0')
    return fig

def create_resistance_zoom_plot(df):
    time_col = 'Time (ms)' if 'Time (ms)' in df.columns else 'Time_ms'
    
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=df[time_col], y=df['Resistance'], name="Resistance", line=dict(color='#27ae60', width=2)))
    
    fig.update_layout(
        title_text="<b>Detailed Resistance (Log Scale)</b>",
        height=300,
        hovermode="x unified",
        plot_bgcolor="white",
        paper_bgcolor="white",
        font=dict(family="Segoe UI, sans-serif"),
        yaxis_type="log", # Log scale to see details
        margin=dict(l=20, r=20, t=40, b=20)
    )
    fig.update_xaxes(showgrid=True, gridwidth=1, gridcolor='#f0f0f0')
    fig.update_yaxes(title_text="Resistance (uOhm)", showgrid=True, gridwidth=1, gridcolor='#f0f0f0')
    return fig