import altair as alt
import gradio as gr
import pandas as pd
import numpy as np


def plot_usage_volume(system1, users1, timeline1, system2, users2, timeline2):
    # Create empty dataframe to hold usage data
    data = pd.DataFrame(columns=["System", "Month", "Usage"])

    # Generate usage data for each system based on typical rollout behaviors
    for i, system in enumerate([system1, system2]):
        # Calculate the monthly usage volume based on the estimated number of users and timeline
        usage = np.concatenate([
            np.linspace(0, [users1, users2][i], int([timeline1, timeline2][i] * 0.25)),
            np.linspace([users1, users2][i] * 0.5, [users1, users2][i], int([timeline1, timeline2][i] * 0.5)),
            np.linspace([users1, users2][i], [users1, users2][i] * 0.75, int([timeline1, timeline2][i] * 0.25))
        ])

        # Add the usage data to the dataframe
        months = ["Month {}".format(j) for j in range(1, len(usage) + 1)]
        system_data = pd.DataFrame({"System": [system] * len(usage), "Month": months, "Usage": usage})
        data = pd.concat([data, system_data], axis=0)

    # Create a multiline plot of the usage volume for each system
    line = alt.Chart(data).mark_line().encode(
        x=alt.X('Month', sort=list(data['Month'].unique())),
        y='Usage:Q',
        color='System:N'
    ).properties(title="System Rollout Usage Volume Plot")

    # Create an aggregated line for the total usage across both systems
    aggregated_data = data.groupby(['Month']).sum().reset_index()
    aggregated_data['System'] = 'Total'
    data = pd.concat([data, aggregated_data], axis=0)

    aggregated_line = alt.Chart(aggregated_data).mark_line(color='black').encode(
        x=alt.X('Month', sort=list(data['Month'].unique())),
        y='Usage:Q'
    )

    # Combine the multiline and aggregated line plots
    chart = alt.layer(line, aggregated_line)

    # Return the plot as a Gradio output
    return chart

# Define the input components
inputs = [
    gr.inputs.Textbox(label="System 1 Name", default="System1"),
    gr.inputs.Number(label="System 1 Users", default=1000),
    gr.inputs.Number(label="System 1 Timeline (months)", default=12),
    gr.inputs.Textbox(label="System 2 Name", default="System2"),
    gr.inputs.Number(label="System 2 Users", default=500),
    gr.inputs.Number(label="System 2 Timeline (months)", default=18)
]

# Define the output component
output = gr.Plot(label="Usage Volume Plot")

# Create the interface
iface = gr.Interface(fn=plot_usage_volume, inputs=inputs, outputs=output, title="System Rollout Usage Volume Plot")

# Launch the interface
iface.launch()