import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots


# Reusable function to create subplots
def create_lac_count_per_controller_subplots(
    df: pd.DataFrame,
    controller_column: str,
    lac_column: str,
    count_column: str,
    fig_title: str,
):
    # Get unique controller_IDs
    unique_controllers = df[controller_column].unique()

    # Calculate the number of rows needed (4 subplots per row)
    rows_needed = (len(unique_controllers) + 3) // 4  # Round up to ensure enough rows

    # Create subplot structure with a dynamic number of rows and 4 columns per row
    fig = make_subplots(
        rows=rows_needed,
        cols=4,
        shared_xaxes=False,
        subplot_titles=unique_controllers,
    )

    # Add a counter for positioning the subplots
    subplot_position = 1

    # Iterate over each controller_ID
    for controller in unique_controllers:
        # Filter data for each controller_ID (create a small dataframe per controller_ID)
        controller_data = df[df[controller_column] == controller]

        # Determine the row and column for the current subplot
        row = (subplot_position - 1) // 4 + 1
        col = (subplot_position - 1) % 4 + 1

        # Add bar chart to the subplot
        fig.add_trace(
            go.Bar(
                x=controller_data[lac_column],
                y=controller_data[count_column],
                name=controller,
                text=controller_data[count_column],
            ),
            row=row,
            col=col,
        )

        # Move to the next subplot position
        subplot_position += 1

    # Update layout to make it more readable and fit all subplots
    fig.update_layout(
        height=300 * rows_needed,
        title_text=fig_title,
        showlegend=False,
    )

    # Show the plot
    # fig.show()

    return fig


def create_bar_chart(df: pd.DataFrame, title: str = "Chart Title") -> px.bar:
    """
    Create a bar chart using Plotly Express with the first column as x and the second column as y.

    Args:
        df (pd.DataFrame): Input DataFrame

    Returns:
        fig (px.bar): Bar chart figure
    """
    fig = px.bar(
        df,
        x=df.columns[0],
        y=df.columns[1],
        text_auto=True,
        title=title,
        height=300,
        width=600,
    )
    fig.update_xaxes(tickvals=df[df.columns[0]].unique())
    return fig