src/plotting/visualization.py

import pandas as pd
import plotly.graph_objects as go


NEWS_COUNT_COLUMN = 0

def plot_time_series(file_path):
    df = pd.read_csv(file_path)
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=df.iloc[:, 0], y=df.iloc[:, 1], mode='lines', name='Time Series'))
    fig.update_layout(title='Disease Mention Time Series', xaxis_title='Date', yaxis_title='Count')

    return fig

def plot_anomalies(df, anomaly_col='new_label'):
    print(df)
    fig = go.Figure()
    
    fig.add_trace(go.Scatter(
        x=df.index, 
        y=df.iloc[:, NEWS_COUNT_COLUMN], 
        mode='lines', 
        name='Time Series',
        line=dict(color='blue')
    ))
    
    anomalies = df[df[anomaly_col] == 1]
    
    fig.add_trace(go.Scatter(
        x=anomalies.index,
        y=anomalies.iloc[:, NEWS_COUNT_COLUMN],
        mode='markers',
        name='Anomalies',
        marker=dict(color='red', size=10, symbol='circle')
    ))
    
    fig.update_layout(
        title='Disease Mention Time Series with Detected Anomalies',
        xaxis_title='Date',
        yaxis_title='Count',
        showlegend=True
    )
    
    return fig