"""Table VII: how much does the authors' arbitrary row choice matter?

MKTBOX2.DO collapses movies to one row per weekend by sorting on weekend and keeping
the first row. FOURDAY, the month dummies and the CPI factor are taken from that row.
In 4 weekends they differ across the movies opening that weekend, and Stata's sort
order among tied rows is arbitrary. This script tries every combination (80), reports
the range of the key estimates, and finds the one closest to the authors' own log
(original_code/mktbox2.log).

p-values use the t distribution with residual degrees of freedom, as Stata's `reg, robust` does.

Run from the package root:  python3 replication/check_row_choice.py data/movie.dta data/mktbox.txt
Writes replication/output/table7_row_choices.csv
"""
import sys
import itertools
from pathlib import Path
import numpy as np
import pandas as pd
import statsmodels.api as sm

raw = pd.read_stata(sys.argv[1], convert_categoricals=False)
mk = pd.read_csv(sys.argv[2], sep=r"\s+", header=None, names=["x", "mktbox"])
raw["weekend"] = raw.opendate + 6 - raw.openwday
wmap = {}
for w in raw.weekend.unique():
    hit = mk[(mk.x >= w - 3) & (mk.x <= w)]
    if len(hit):
        wmap[w] = hit.mktbox.iloc[0]
m = raw.copy()
m["mktbox"] = m.weekend.map(wmap)
m["mbcpi"] = m.mktbox * m.grosscpi / m.gross
m = m[m.pre != 1].copy()
m["screens"] = m.openscr / 1000
YEARS = [f"year{i}" for i in range(2, 10)]
M11 = [f"month{i}" for i in range(1, 12)]
vars_ = ["dur", "oudur", "uudur", "ou", "uu", "screens", "maltin"]
xv = ["dur", "oudur", "uudur", "ou", "uu", "fourday"] + M11 + YEARS + ["screens", "maltin"]

amb = [w for w, g in m.groupby("weekend") if g.fourday.nunique() > 1 or g.month.nunique() > 1]
print("ambiguous weekends (SAS dates):", amb)
options = [list(m.index[m.weekend == w]) for w in amb]
print("combinations:", int(np.prod([len(o) for o in options])))

# authors' log (original_code/mktbox2.log): FOURDAY, D x TWO UP, R2
LOG = {"sum": (0.3118941, -0.0829835, 0.6772), "mean": (0.2738601, -0.04621, 0.5392)}

records = []
for stat in ("sum", "mean"):
    agg = m.groupby("weekend")[vars_ + ["wk1cpi"]].agg(stat)
    base_first = m.groupby("weekend").head(1).set_index("weekend")
    best = None
    for combo in itertools.product(*options):
        first = base_first.copy()
        for w, idx in zip(amb, combo):
            first.loc[w, ["mbcpi", "fourday"] + YEARS + M11] = m.loc[idx, ["mbcpi", "fourday"] + YEARS + M11].values
        out = first[["mbcpi", "fourday"] + YEARS + M11].join(agg)
        out["ratio"] = out.wk1cpi / out.mbcpi
        out = out[out.mbcpi.notna() & (out.ratio <= 1)]
        out["lresid"] = np.log(out.mbcpi - out.wk1cpi)
        X = sm.add_constant(out[xv].astype(float))
        r = sm.OLS(out.lresid, X).fit(cov_type="HC1", use_t=True)
        got = (r.params["fourday"], r.params["uudur"], r.rsquared)
        dist = sum(abs(a - b) for a, b in zip(got, LOG[stat]))
        records.append(dict(column=stat, kept_titles=" | ".join(m.loc[i, "title"] for i in combo),
                            fourday=r.params["fourday"], d_twoup=r.params["uudur"], d_twoup_se=r.bse["uudur"],
                            d_twoup_p=r.pvalues["uudur"], r2=r.rsquared, n=int(r.nobs), dist_to_log=dist))
        if best is None or dist < best[0]:
            best = (dist, combo, r)
    print(f"\n=== Column '{stat}' ===")
    print(f"Closest to authors' log (distance {best[0]:.5f}); kept: {[m.loc[i, 'title'] for i in best[1]]}")
    r = best[2]
    print(pd.DataFrame({"b": r.params, "se": r.bse, "p": r.pvalues})
          .loc[["dur", "oudur", "uudur", "ou", "uu", "maltin", "screens", "fourday", "const"]].round(4))
    print(f"N {int(r.nobs)}  R2 {r.rsquared:.4f}")

res = pd.DataFrame(records)
OUT = Path(__file__).resolve().parent / "output"
OUT.mkdir(exist_ok=True)
res.to_csv(OUT / "table7_row_choices.csv", index=False)

print("\n=== Range over all row choices ===")
for stat, g in res.groupby("column", sort=False):
    print(f"{stat}: D x TWO UP {g.d_twoup.min():.3f} to {g.d_twoup.max():.3f}; "
          f"p {g.d_twoup_p.min():.3f} to {g.d_twoup_p.max():.3f}; "
          f"share of choices with p < 0.10: {(g.d_twoup_p < 0.10).mean():.2f}; "
          f"FOURDAY {g.fourday.min():.3f} to {g.fourday.max():.3f}; N {sorted(int(x) for x in g.n.unique())}")