import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

from pathlib import Path
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.pipeline import Pipeline
from sklearn.linear_model import RidgeCV
from sklearn.metrics import r2_score, mean_squared_error

pd.set_option("display.max_columns", 100)
pd.set_option("display.float_format", lambda x: f"{x:,.4f}")

DATA_PATH = "data/Sample Media Spend Data.csv"


df = pd.read_csv(DATA_PATH)
print("Shape:", df.shape)
display(df.head())

date_col = "Calendar_Week"
division_col = "Division"
target_col = "Sales"

channel_cols = [
    "Paid_Views",
    "Organic_Views",
    "Google_Impressions",
    "Email_Impressions",
    "Facebook_Impressions",
    "Affiliate_Impressions",
    "Overall_Views",
]

required_cols = [date_col, division_col, target_col] + channel_cols
missing = [c for c in required_cols if c not in df.columns]
if missing:
    raise ValueError(f"Missing expected columns: {missing}")

df[date_col] = pd.to_datetime(df[date_col])
df = df.sort_values([division_col, date_col]).reset_index(drop=True)

print("Date range:", df[date_col].min(), "to", df[date_col].max())
print("Number of divisions:", df[division_col].nunique())

print("\nMissing values:")
display(df[required_cols].isna().sum().to_frame("missing_values"))

print("\nRows per division:")
display(df.groupby(division_col).size().describe().to_frame("rows_per_division"))

weekly = (
    df.groupby(date_col, as_index=False)
      .agg({target_col: "sum", **{c: "sum" for c in channel_cols}})
      .sort_values(date_col)
)

plt.figure(figsize=(12, 5))
plt.plot(weekly[date_col], weekly[target_col])
plt.title("Weekly Sales Over Time")
plt.xlabel("Week")
plt.ylabel("Sales")
plt.xticks(rotation=45)
plt.tight_layout()
plt.show()

normalized = weekly[[date_col] + channel_cols].copy()
for c in channel_cols:
    std = normalized[c].std()
    normalized[c] = 0 if std == 0 else (normalized[c] - normalized[c].mean()) / std

plt.figure(figsize=(12, 6))
for c in channel_cols:
    plt.plot(normalized[date_col], normalized[c], label=c, alpha=0.8)

plt.title("Normalized Channel Activity Over Time")
plt.xlabel("Week")
plt.ylabel("Standardized activity")
plt.legend(loc="upper left", bbox_to_anchor=(1, 1))
plt.xticks(rotation=45)
plt.tight_layout()
plt.show()

df["week_index"] = df.groupby(division_col).cumcount()
df["month"] = df[date_col].dt.month.astype(str)
df["rank_pct"] = df.groupby(division_col)[date_col].rank(pct=True)

train_df = df[df["rank_pct"] <= 0.80].copy()
test_df = df[df["rank_pct"] > 0.80].copy()

print("Train shape:", train_df.shape)
print("Test shape:", test_df.shape)
print("Train date range:", train_df[date_col].min(), "to", train_df[date_col].max())
print("Test date range:", test_df[date_col].min(), "to", test_df[date_col].max())

def fit_ridge_model(train, test, numeric_features, categorical_features, target=target_col):
    """
    Fit a Ridge regression model.

    Numeric features are standardized so channel coefficients are more comparable.
    Categorical features control for division and month.
    """
    all_features = numeric_features + categorical_features

    preprocessor = ColumnTransformer(
        transformers=[
            ("num", StandardScaler(), numeric_features),
            ("cat", OneHotEncoder(handle_unknown="ignore"), categorical_features),
        ]
    )

    model = Pipeline(
        steps=[
            ("preprocessor", preprocessor),
            ("ridge", RidgeCV(alphas=np.logspace(-3, 3, 25))),
        ]
    )

    model.fit(train[all_features], train[target])
    pred = model.predict(test[all_features])

    rmse = mean_squared_error(test[target], pred) ** 0.5
    r2 = r2_score(test[target], pred)

    feature_names = model.named_steps["preprocessor"].get_feature_names_out()
    coefs = pd.Series(model.named_steps["ridge"].coef_, index=feature_names)

    return model, pred, rmse, r2, coefs

categorical_features = [division_col, "month"]
baseline_numeric = channel_cols + ["week_index"]

baseline_model, baseline_pred, baseline_rmse, baseline_r2, baseline_coefs = fit_ridge_model(
    train_df, test_df, baseline_numeric, categorical_features
)

baseline_importance = (
    pd.DataFrame({
        "channel": channel_cols,
        "standardized_coefficient": [
            baseline_coefs.get(f"num__{c}", np.nan) for c in channel_cols
        ],
    })
    .assign(abs_coefficient=lambda x: x["standardized_coefficient"].abs())
    .sort_values("abs_coefficient", ascending=False)
)

print(f"Baseline RMSE: {baseline_rmse:,.0f}")
print(f"Baseline R²: {baseline_r2:,.3f}")
display(baseline_importance)

def add_geometric_adstock(data, cols, group_col, decay):
    """
    Create geometric adstock features within each division.

    This must be computed separately by division.
    Otherwise, activity from one division would incorrectly carry into another.
    """
    out = data.copy()

    for c in cols:
        new_col = f"{c}_adstock"
        out[new_col] = np.nan

        for _, idx in out.groupby(group_col).groups.items():
            carry = 0.0
            values = out.loc[idx, c].fillna(0).values
            transformed = []

            for x in values:
                carry = x + decay * carry
                transformed.append(carry)

            out.loc[idx, new_col] = transformed

    return out

decay = 0.50
df_adstock = add_geometric_adstock(df, channel_cols, division_col, decay=decay)
adstock_cols = [f"{c}_adstock" for c in channel_cols]

train_adstock = df_adstock[df_adstock["rank_pct"] <= 0.80].copy()
test_adstock = df_adstock[df_adstock["rank_pct"] > 0.80].copy()

adstock_model, adstock_pred, adstock_rmse, adstock_r2, adstock_coefs = fit_ridge_model(
    train_adstock, test_adstock, adstock_cols + ["week_index"], categorical_features
)

adstock_importance = (
    pd.DataFrame({
        "channel": channel_cols,
        "standardized_coefficient": [
            adstock_coefs.get(f"num__{c}_adstock", np.nan) for c in channel_cols
        ],
    })
    .assign(abs_coefficient=lambda x: x["standardized_coefficient"].abs())
    .sort_values("abs_coefficient", ascending=False)
)

print(f"Adstock RMSE: {adstock_rmse:,.0f}")
print(f"Adstock R²: {adstock_r2:,.3f}")
display(adstock_importance)

df_sat = df_adstock.copy()

for c in adstock_cols:
    df_sat[f"log_{c}"] = np.log1p(df_sat[c])

sat_cols = [f"log_{c}" for c in adstock_cols]

train_sat = df_sat[df_sat["rank_pct"] <= 0.80].copy()
test_sat = df_sat[df_sat["rank_pct"] > 0.80].copy()

sat_model, sat_pred, sat_rmse, sat_r2, sat_coefs = fit_ridge_model(
    train_sat, test_sat, sat_cols + ["week_index"], categorical_features
)

sat_importance = (
    pd.DataFrame({
        "channel": channel_cols,
        "standardized_coefficient": [
            sat_coefs.get(f"num__log_{c}_adstock", np.nan) for c in channel_cols
        ],
    })
    .assign(abs_coefficient=lambda x: x["standardized_coefficient"].abs())
    .sort_values("abs_coefficient", ascending=False)
)

print(f"Adstock + Saturation RMSE: {sat_rmse:,.0f}")
print(f"Adstock + Saturation R²: {sat_r2:,.3f}")
display(sat_importance)

comparison = pd.DataFrame({
    "model": ["Baseline", "Adstock", "Adstock + Saturation"],
    "rmse": [baseline_rmse, adstock_rmse, sat_rmse],
    "r2": [baseline_r2, adstock_r2, sat_r2],
})

display(comparison)

plt.figure(figsize=(8, 5))
plt.bar(comparison["model"], comparison["r2"])
plt.title("Test R² by MMM Specification")
plt.ylabel("Test R²")
plt.tight_layout()
plt.show()

def rank_table(imp, name):
    out = imp.copy()
    out["rank"] = out["abs_coefficient"].rank(method="first", ascending=False).astype(int)
    out["model"] = name
    return out[["model", "rank", "channel", "standardized_coefficient", "abs_coefficient"]]

ranks = pd.concat([
    rank_table(baseline_importance, "Baseline"),
    rank_table(adstock_importance, "Adstock"),
    rank_table(sat_importance, "Adstock + Saturation"),
], ignore_index=True)

display(ranks.sort_values(["model", "rank"]))

sensitivity_rows = []

for decay_value in [0.10, 0.30, 0.50, 0.70, 0.90]:
    temp = add_geometric_adstock(df, channel_cols, division_col, decay=decay_value)
    temp_cols = [f"{c}_adstock" for c in channel_cols]

    temp_train = temp[temp["rank_pct"] <= 0.80].copy()
    temp_test = temp[temp["rank_pct"] > 0.80].copy()

    _, _, rmse, r2, coefs = fit_ridge_model(
        temp_train, temp_test, temp_cols + ["week_index"], categorical_features
    )

    scores = [(c, abs(coefs.get(f"num__{c}_adstock", 0))) for c in channel_cols]
    top_channel = sorted(scores, key=lambda x: x[1], reverse=True)[0][0]

    sensitivity_rows.append({
        "decay": decay_value,
        "rmse": rmse,
        "r2": r2,
        "top_channel": top_channel,
    })

sensitivity = pd.DataFrame(sensitivity_rows)
display(sensitivity)

plt.figure(figsize=(8, 5))
plt.plot(sensitivity["decay"], sensitivity["r2"], marker="o")
plt.title("Adstock Decay Sensitivity")
plt.xlabel("Decay")
plt.ylabel("Test R²")
plt.tight_layout()
plt.show()