print(orbit.__version__)
print(pd.__version__)
print(np.__version__)
print(scipy.__version__)

budget_start = pd.to_datetime("2021-01-01")
budget_end = pd.to_datetime("2021-01-02")
optim_channels = mmm.get_spend_cols()[1:3]
optim_channels

temp_mmm = deepcopy(mmm)
attributor = Attributor(model=temp_mmm, start=budget_start, end=budget_end)
t_npm = TimeNetProfitMaximizer(
    ltv_arr=ltv_arr,
    model=temp_mmm,
    attributor=attributor,
    budget_start=budget_start,
    budget_end=budget_end,
    optim_channels=optim_channels,
    spend_scaler=1.0,
    response_scaler=1.0,
)

%%time
x, y, z = time_based_net_profit_response_curve(t_npm=t_npm, model=temp_mmm, n_iters=10)
fig, ax = plt.subplots(1, 1, figsize=(12, 8))
cs = ax.contourf(x, y, z, 20, cmap="RdGy")
fig.colorbar(cs, ax=ax)
ax.set_xlabel("t1")
ax.set_ylabel("t2")
ax.plot(
    [0, t_npm.total_budget],
    [t_npm.total_budget, 0],
    alpha=0.5,
    label="total budget constraint",
    linestyle="--",
)

# run optim steps
t_optim_df = t_npm.optimize(maxiter=500)

# use callback metrics to extract solutions on each step
callback_metrics = t_npm.get_callback_metrics()
xs = np.array(callback_metrics["xs"])
ax.arrow(
    x=xs[0, 0],
    y=xs[0, 1],
    dx=xs[-1, 0] - xs[0, 0],
    dy=xs[-1, 1] - xs[0, 1],
    width=80,
    color="green",
)
fig.legend(loc="lower center")
fig.tight_layout()

temp_mmm = deepcopy(mmm)
ch_npm = ChannelNetProfitMaximizer(
    ltv_arr=ltv_arr,
    model=temp_mmm,
    attributor=attributor,
    optim_channels=optim_channels,
    budget_start=budget_start,
    budget_end=budget_end,
    spend_scaler=1.0,
    response_scaler=1.0,
)

%%time
# we need higher iteration to support higher resolution on channels response surface
x, y, z = ch_based_net_profit_response_curve(ch_npm=ch_npm, model=temp_mmm, n_iters=20)
fig, ax = plt.subplots(1, 1, figsize=(12, 8))
cs = ax.contourf(x, y, z, 20, cmap="RdGy")
fig.colorbar(cs, ax=ax)
ax.set_xlabel("t1")
ax.set_ylabel("t2")
ax.plot(
    [0, ch_npm.total_budget],
    [ch_npm.total_budget, 0],
    alpha=0.5,
    label="total budget constraint",
    linestyle="--",
)

# run optim steps
cn_optim_df = ch_npm.optimize(maxiter=500)

# use callback metrics to extract solutions on each step
callback_metrics = ch_npm.get_callback_metrics()
xs = np.array(callback_metrics["xs"])
ax.arrow(
    x=xs[0, 0],
    y=xs[0, 1],
    dx=xs[-1, 0] - xs[0, 0],
    dy=xs[-1, 1] - xs[0, 1],
    width=80,
    color="green",
)
fig.legend(loc="lower center")
fig.tight_layout()

%%time
attributor = Attributor(model=temp_mmm, start=budget_start, end=budget_end)
ch_npm = ChannelNetProfitMaximizer(
    ltv_arr=ltv_arr,
    model=temp_mmm,
    attributor=attributor,
    optim_channels=optim_channels,
    budget_start=budget_start,
    budget_end=budget_end,
    spend_scaler=1.0,
    response_scaler=1.0,
)

temp_optim_spend_df = ch_npm.optimize(maxiter=500, eps=1e-3)

# expermential plot the revs and costs
fig, ax = plt.subplots(1, 1, figsize=(16, 8))
callback_metrics = ch_npm.get_callback_metrics()

optim_revenues = np.array(callback_metrics["optim_revenues"])
optim_costs = np.array(callback_metrics["optim_costs"])

ax.plot(optim_revenues, label="revs")
ax.plot(optim_costs, label="costs")
ax.plot(optim_revenues - optim_costs, label="net-profit")
ax.set_xlabel("iterations")
ax.set_ylabel("metrics")
fig.legend()
fig.tight_layout();

# total spend suggest
print("Suggested total spend: {:.0f}".format(np.sum(channel_spend_arr)))
# total budget
print("Total Budget: {:.0f}".format(np.sum(init_spend_arr)))

%%time
temp_mmm = deepcopy(mmm)
temp_mmm.raw_df = temp_optim_spend_df
t_npm = TimeNetProfitMaximizer(
    ltv_arr=ltv_arr,
    model=temp_mmm,
    attributor=attributor,
    budget_start=budget_start,
    budget_end=budget_end,
    optim_channels=optim_channels,
    spend_scaler=1.0,
    response_scaler=1.0,
)

optim_spend_df = t_npm.optimize(maxiter=500)

# total spend suggest
print("Suggested total spend: {:.0f}".format(np.sum(optim_spend_matrix)))
# total budget
print("Total Budget: {:.0f}".format(np.sum(init_spend_matrix)))

cost_report = generate_cost_report(
    model=mmm,
    start=budget_start,
    end=budget_end,
    pre_spend_df=df,
    post_spend_df=optim_spend_df,
)
cost_report["ltv"] = ltv_arr
cost_report

init_total_spend = np.sum(init_spend_matrix, 0)
optim_total_spend = np.sum(optim_spend_matrix, 0)
plot_data = np.vstack([init_total_spend, optim_total_spend])
plot_data.shape

fig, ax = plt.subplots(1, 1, figsize=(20, 10))
bottom = 0.0
for idx, label in enumerate(optim_channels):
    ax.bar(x=["Pre", "Post"], height=plot_data[:, idx], bottom=bottom, label=label)
    bottom += plot_data[:, idx]
ax.set_title("Spend Allocation Pre vs. Post Optimization")
ax.legend(
    loc="lower center",
    bbox_to_anchor=(0.5, -0.10),
    ncol=math.ceil(len(optim_channels) / 2),
);

t_based_allocation_comp = t_npm.get_current_state()

fig, ax = plt.subplots(1, 1, figsize=(16, 8))
ax.plot(
    pd.date_range(budget_start, budget_end, freq="1D"),
    t_based_allocation_comp / np.sum(t_based_allocation_comp),
    label="MMM-Time-Based-Budget%",
    alpha=0.5,
)
ax.plot(
    pd.date_range(budget_start, budget_end, freq="1D"),
    organic_attr_arr / np.sum(organic_attr_arr),
    label="MMM-Forecasted-Organic",
    alpha=0.5,
)
fig.legend()
fig.tight_layout()

bounds_and_constraints_df = pd.DataFrame(
    {
        "channels": optim_channels + ["total"],
        "upper": np.array(
            [291918.0, 589910.0, 481658.0, 342542.0, 948156.0, 1327092.0]
        ),
        "lower": np.array([14595.9, 29495.5, 24082.9, 17127.1, 47407.8, 1327092.0]),
    }
)
bounds_and_constraints_df

ch_npm.set_bounds_and_constraints(df=bounds_and_constraints_df)

bounds_and_constraints_df = pd.DataFrame(
    {
        "date": t_npm.df.loc[t_npm.budget_mask, "date"].astype(str).to_list()
        + ["total"],
        "upper": np.concatenate(
            [np.ones(t_npm.n_budget_steps) * 30000, np.ones(1) * 1327092]
        ),
        "lower": np.concatenate([np.zeros(t_npm.n_budget_steps), np.ones(1) * 1327092]),
    }
)
bounds_and_constraints_df

t_npm.set_bounds_and_constraints(bounds_and_constraints_df)