import gradio as gr
import pandas as pd
import numpy as np
from sklearn.linear_model import LinearRegression
import matplotlib.pyplot as plt

# To avoid backend display issues on HF
plt.switch_backend("Agg")


def generate_fake_data():
    np.random.seed(42)
    weeks = 104

    data = pd.DataFrame({
        "week": range(1, weeks + 1),
        "tv_spend": np.random.gamma(30, 15, weeks),
        "social_spend": np.random.gamma(20, 20, weeks),
        "search_spend": np.random.gamma(25, 10, weeks),
        "price": np.random.normal(10, 0.4, weeks),
    })

    # true coefficients
    true_coef = {"tv": 0.8, "social": 1.2, "search": 0.6, "price": -20}

    data["sales"] = (
        true_coef["tv"] * data["tv_spend"] +
        true_coef["social"] * data["social_spend"] +
        true_coef["search"] * data["search_spend"] +
        true_coef["price"] * data["price"] +
        np.random.normal(0, 50, weeks)
    ).round(2)

    return data


def run_mmm(data, incremental_increase):
    X = data[["tv_spend", "social_spend", "search_spend", "price"]]
    y = data["sales"]

    model = LinearRegression()
    model.fit(X, y)

    coef = pd.Series(model.coef_, index=X.columns).round(4)
    intercept = round(model.intercept_, 2)

    # Incrementality
    social_coef = coef["social_spend"]
    incremental_sales = round(social_coef * incremental_increase, 2)

    # Plot actual vs predicted
    data["pred_sales"] = model.predict(X)

    plt.figure(figsize=(8, 4))
    plt.plot(data["sales"], label="Actual Sales")
    plt.plot(data["pred_sales"], label="Predicted Sales")
    plt.legend()
    plt.title("Actual vs Predicted Sales (MMM)")
    plt.xlabel("Week")
    plt.ylabel("Sales")

    plot_path = "plot.png"
    plt.tight_layout()
    plt.savefig(plot_path)
    plt.close()

    return (
        coef.to_frame("Coefficient").reset_index().rename(columns={"index": "Variable"}),
        f"Incremental sales from +${incremental_increase} social spend: {incremental_sales}",
        plot_path
    )


def ui_run(incremental_increase):
    data = generate_fake_data()
    coef_table, inc_msg, plot_path = run_mmm(data, incremental_increase)
    return coef_table, inc_msg, plot_path


### -------- Gradio Interface -------- ###
demo = gr.Interface(
    fn=ui_run,
    inputs=[
        gr.Number(label="Increase in Social Spend (In Dollars)", value=10000)
    ],
    outputs=[
        gr.Dataframe(label="MMM Coefficients"),
        gr.Textbox(label="Incrementality Result"),
        gr.Image(label="Actual vs Predicted Sales")
    ],
    title="Marketing Mix Model (MMM) + Incrementality",
    description="A simple MMM that models TV, Social, Search, and Price to predict Sales. Enter an additional Social Ad spend to estimate the incremental uplift.",
    allow_flagging="never"
)

if __name__ == "__main__":
    demo.launch(share=True)