src/chronos_conference/adapters/ts_plot.py

# ruff: noqa: F403, F405

import plotly.graph_objects as go
from plotly.subplots import make_subplots
import pandas as pd
from chronos_conference.settings import *


def get_plot(df_story: pd.DataFrame, df_pred: pd.DataFrame):
    if not pd.api.types.is_datetime64_any_dtype(df_story[HISTORICAL_DATE_COLUMN]):
        df_story[HISTORICAL_DATE_COLUMN] = pd.to_datetime(
            df_story[HISTORICAL_DATE_COLUMN]
        )

    df_story = df_story.rename(
        columns={
            HISTORICAL_DATE_COLUMN: "datetime",
            HISTORICAL_TARGET_COLUMN: "value",
            HISTORICAL_ITEM_COLUMN: "item_id",
        }
    )

    if not pd.api.types.is_datetime64_any_dtype(df_pred["ds"]):
        df_pred["ds"] = pd.to_datetime(df_pred["ds"])

    item_ids = df_story["item_id"].unique()
    number_columns = len(item_ids)

    subplot_titles = [
        f"{item}<br><span style='font-size:10px;color:gray;'>Unidad: {UNITS_MEASURED[item]} | "
        f"Límite máx: {MAX_SAFETY_LIMITS[item]}</span>"
        for item in item_ids
    ]

    fig = make_subplots(rows=1, cols=number_columns, subplot_titles=subplot_titles)

    for idx, value in enumerate(item_ids):
        df_story_subset = df_story[df_story["item_id"] == value]
        df_pred_subset = df_pred[df_pred["item_id"] == value]

        show_legend = idx == 0

        fig.add_trace(
            go.Scatter(
                x=df_story_subset["datetime"],
                y=df_story_subset["value"],
                mode="lines",
                name="Histórico",
                line=dict(color="blue", width=2),
                showlegend=show_legend,
            ),
            row=1,
            col=idx + 1,
        )

        fig.add_trace(
            go.Scatter(
                x=df_pred_subset["ds"],
                y=df_pred_subset["AWSChronosForecast"],
                mode="lines+markers",
                name="Predicción",
                line=dict(color="orange", dash="dash"),
                marker=dict(symbol="x", size=8, color="orange"),
                showlegend=show_legend,
            ),
            row=1,
            col=idx + 1,
        )

    fig.update_layout(
        showlegend=True,
        height=400,
        width=300 * number_columns,
        title_text="Serie histórica y predicción por ítem",
    )

    return fig