"""Model F — rolling-origin forecasts of next-week posted log price and promotion indicator (preregistration §2 F).

Features use only information up to week t (strict lags). Origins t = 6..10 (train weeks <= t, test week t+1).
Models: seasonal naive (last value), banner×category mean reversion (shrink toward 8-week listing mode), ridge, LightGBM.
Promotion target: next week p <= 0.9 x trailing 8-week mode. Metrics: MAE/RMSE (log price), Brier/AUC (promotion).
"""
import json, numpy as np, pandas as pd, lightgbm as lgb
from sklearn.linear_model import Ridge, LogisticRegression
from sklearn.metrics import roc_auc_score, brier_score_loss
from src.common import path, ROOT, get_logger, ledger, seed_everything, SEED
log = get_logger("modelF"); P = path("processed"); E = ROOT / "experiments"

def rolling_mode_series(s, window=8):
    out = []
    v = s.tolist()
    for i in range(len(v)):
        w = pd.Series(v[max(0, i - window + 1): i + 1]).value_counts(); out.append(max(w[w == w.max()].index))
    return pd.Series(out, index=s.index)

def main():
    seed_everything(); out = {}
    pw = pd.read_parquet(P / "panel_weekly.parquet"); sf = pw[pw.channel == "online_storefront"].sort_values(["uid", "week"]).copy()
    sf = sf[sf.groupby("uid").week.transform("nunique") >= 6]
    sf["t"] = sf.week.rank(method="dense").astype(int)
    g = sf.groupby("uid")
    sf["ref"] = g.p_mode.transform(rolling_mode_series); sf["promo"] = (sf.p_mode <= 0.9 * sf.ref).astype(int)
    for k in (1, 2, 3): sf[f"lp_l{k}"] = g.log_p.shift(k)
    sf["lp_ref"] = np.log(sf.ref); sf["promo_l1"] = g.promo.shift(1); sf["promo_l2"] = g.promo.shift(2)
    sf["n_changes_l0"] = sf.n_changes; sf["rel_range_l0"] = sf.rel_range; sf["multi_l0"] = sf.multi_price_week.astype(int)
    sf["promo_rate_hist"] = g.promo.transform(lambda s: s.shift(1).expanding().mean())
    sf["y_lp"] = g.log_p.shift(-1); sf["y_promo"] = g.promo.shift(-1); sf["t_next"] = g.t.shift(-1)
    sf = sf[(sf.t_next == sf.t + 1)]
    sf["banner_c"] = sf.banner.astype("category").cat.codes; sf["cat_c"] = sf.category.astype("category").cat.codes
    feats = ["log_p", "lp_l1", "lp_l2", "lp_l3", "lp_ref", "promo", "promo_l1", "promo_l2", "n_changes_l0", "rel_range_l0", "multi_l0", "promo_rate_hist", "banner_c", "cat_c", "t"]
    res = []
    for t in range(6, int(sf.t.max())):
        tr = sf[(sf.t <= t - 1)].dropna(subset=feats + ["y_lp", "y_promo"]); te = sf[sf.t == t].dropna(subset=feats + ["y_lp", "y_promo"])
        if len(te) == 0 or len(tr) == 0: continue
        # price models
        preds = {"naive_last": te.log_p.values, "mean_reversion_ref": 0.5 * te.log_p.values + 0.5 * te.lp_ref.values}
        bc = tr.groupby(["banner", "category"]).apply(lambda d: (d.y_lp - d.log_p).mean()).rename("drift").reset_index()
        preds["banner_cat_drift"] = te.log_p.values + te.merge(bc, on=["banner", "category"], how="left").drift.fillna(0).values
        X = tr[feats].values; Xt = te[feats].values
        rd = Ridge(alpha=1.0).fit(X, tr.y_lp.values); preds["ridge"] = rd.predict(Xt)
        gb = lgb.LGBMRegressor(n_estimators=400, learning_rate=0.05, num_leaves=63, min_child_samples=50, subsample=0.8, colsample_bytree=0.8, random_state=SEED, verbose=-1).fit(X, tr.y_lp.values)
        preds["lightgbm"] = gb.predict(Xt)
        for m, p in preds.items():
            e = p - te.y_lp.values; res.append(dict(origin_t=t, target="log_price", model=m, n=len(te), mae=float(np.mean(np.abs(e))), rmse=float(np.sqrt(np.mean(e ** 2)))))
        # promotion models
        pp = {"persist": te.promo.values.astype(float), "base_rate": np.full(len(te), tr.y_promo.mean()), "hist_rate": te.promo_rate_hist.fillna(tr.y_promo.mean()).values}
        lr = LogisticRegression(max_iter=2000).fit(X, tr.y_promo.values); pp["logit"] = lr.predict_proba(Xt)[:, 1]
        gc = lgb.LGBMClassifier(n_estimators=400, learning_rate=0.05, num_leaves=63, min_child_samples=50, subsample=0.8, colsample_bytree=0.8, random_state=SEED, verbose=-1).fit(X, tr.y_promo.values)
        pp["lightgbm"] = gc.predict_proba(Xt)[:, 1]
        for m, p in pp.items():
            y = te.y_promo.values
            res.append(dict(origin_t=t, target="promotion", model=m, n=len(te), brier=float(brier_score_loss(y, np.clip(p, 0, 1))), auc=float(roc_auc_score(y, p)) if len(set(y)) > 1 else np.nan, base_rate=float(y.mean())))
    r = pd.DataFrame(res); r.to_csv(P / "modelF_rolling_results.csv", index=False)
    summ = r.groupby(["target", "model"]).agg(origins=("origin_t", "nunique"), n=("n", "sum"), mae=("mae", "mean"), rmse=("rmse", "mean"), brier=("brier", "mean"), auc=("auc", "mean")).reset_index()
    summ.to_csv(P / "modelF_summary.csv", index=False); out["summary"] = summ.to_dict(orient="records")
    (E / "modelF_summary.json").write_text(json.dumps(out, indent=2, default=float))
    log.info("\n%s", summ.round(4).to_string(index=False))
    ledger("4", "model_f", "rolling-origin forecasts done")

if __name__ == "__main__":
    main()