streamlit_app/components/charts.py

"""
Reusable chart components for Streamlit dashboard
"""

import plotly.express as px
import plotly.graph_objects as go
import pandas as pd
import numpy as np
from typing import Optional, List


def create_line_chart(df: pd.DataFrame, x: str, y: str, title: str,
                      color: Optional[str] = None, height: int = 500) -> go.Figure:
    """Create a line chart"""
    fig = px.line(df, x=x, y=y, title=title, color=color, height=height)
    fig.update_layout(
        hovermode='x unified',
        template='plotly_white',
    )
    return fig


def create_bar_chart(df: pd.DataFrame, x: str, y: str, title: str,
                     color: Optional[str] = None, height: int = 500) -> go.Figure:
    """Create a bar chart"""
    fig = px.bar(df, x=x, y=y, title=title, color=color, height=height)
    fig.update_layout(
        template='plotly_white',
        showlegend=True,
    )
    return fig


def create_scatter_plot(df: pd.DataFrame, x: str, y: str, title: str,
                       size: Optional[str] = None, color: Optional[str] = None,
                       height: int = 500) -> go.Figure:
    """Create a scatter plot"""
    fig = px.scatter(df, x=x, y=y, title=title, size=size, color=color, height=height)
    fig.update_layout(
        template='plotly_white',
        hovermode='closest',
    )
    return fig


def create_histogram(df: pd.DataFrame, column: str, title: str,
                    nbins: int = 30, height: int = 500) -> go.Figure:
    """Create a histogram"""
    fig = px.histogram(df, x=column, title=title, nbins=nbins, height=height)
    fig.update_layout(
        template='plotly_white',
        xaxis_title=column,
        yaxis_title='Frequency',
    )
    return fig


def create_box_plot(df: pd.DataFrame, y: str, x: Optional[str] = None,
                   title: str = "Box Plot", height: int = 500) -> go.Figure:
    """Create a box plot"""
    fig = px.box(df, x=x, y=y, title=title, height=height)
    fig.update_layout(template='plotly_white')
    return fig


def create_heatmap(data: np.ndarray, x_labels: List[str], y_labels: List[str],
                   title: str = "Heatmap", height: int = 600) -> go.Figure:
    """Create a heatmap"""
    fig = go.Figure(data=go.Heatmap(
        z=data,
        x=x_labels,
        y=y_labels,
        colorscale='Viridis',
    ))
    fig.update_layout(
        title=title,
        height=height,
        template='plotly_white',
    )
    return fig


def create_pie_chart(df: pd.DataFrame, values: str, names: str,
                    title: str = "Pie Chart", height: int = 500) -> go.Figure:
    """Create a pie chart"""
    fig = px.pie(df, values=values, names=names, title=title, height=height)
    fig.update_layout(template='plotly_white')
    return fig