import pandas as pd
import matplotlib.pyplot as plt
import karpiu
from karpiu.models import MMM

pd.set_option("display.float_format", lambda x: "%.3f" % x)

print(karpiu.__version__)

df = pd.read_csv(RAW_DATA_FILE, parse_dates=["date"])
df.head(5)

scalability_df = pd.read_csv(SCALABILITY_FILE)
scalability_df.head(5)

mmm = MMM(
    kpi_col="sales",
    date_col="date",
    spend_cols=paid_channels,
    scalability_df=scalability_df,
    event_cols=event_cols,
    seed=2022,
    seasonality=[7, 365.25],
    fs_orders=[2, 3],
    events_sigma_prior=0.3,
)

best_params = {
    "damped_factor": 0.949,
    "level_sm_input": 0.00245,
}
mmm.set_hyper_params(best_params)

mmm.fit(df, num_warmup=1000, num_sample=1000)

ATTR_START = "2019-03-01"
ATTR_END = "2019-03-31"
attributor = AttributorGamma(model=mmm, start=ATTR_START, end=ATTR_END)
activities_attr_df, spend_attr_df, spend_df, cost_df = attributor.make_attribution()

ax = plot_attribution_with_time(
    model=mmm,
    attr_df=activities_attr_df,
    figsize=(10, 5.5),
    colors=ColorConstants.RAINBOW_SIX,
    show=False,
    dt_col="date",
    include_organic=False,
)
ax.set_title("Attribution by Activities Date", fontdict={"fontsize": 12})
ax.set_xlabel("Date", fontdict={"fontsize": 12})
ax.set_ylabel("Sales", fontdict={"fontsize": 12})

fig = ax.figure
fig

ax = plot_attribution_waterfall(
    model=mmm,
    attr_df=spend_attr_df,
    figsize=(8, 4.5),
    include_organic=False,
    colors=ColorConstants.RAINBOW_SIX,
    show=False,
)
ax.set_title("Sales Decomposition", fontdict={"fontsize": 12})
fig = ax.figure
fig

cc = CostCurves(
    model=mmm,
    curve_type="individual",
    n_steps=50,
    spend_start=ATTR_START,
    spend_end=ATTR_END,
)
cc.generate_cost_curves()

cc.plot(spend_scaler=10, outcome_scaler=10, include_organic=False);

mmm.get_regression_summary()

plot_predicted_data(
    training_actual_df=df[-365:],
    predicted_df=pred_df[-(720):],
    date_col="date",
    actual_col="sales",
);

plot_predicted_components(
    predicted_df=pred_df[-720:],
    date_col="date",
    plot_components=[
        "trend",
        "paid",
        "events",
        "s-7 seasonality",
        "s-365.25 seasonality",
    ],
);