"""Re-run the key tables of Reinstein & Snyder (JIE 2005) in Python.

Tables I-V. Specifications are copied from the authors' Stata files
original_code/revision4.do (Tables IV, V) and original_code/SELECT2.DO (Table III).
Ports `reg ..., robust cluster(weekend)` and `probit ..., robust cluster(weekend)`
to statsmodels. SCREENS is rescaled to thousands to match the paper's units.

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

pd.set_option("display.width", 250)
pd.set_option("display.max_columns", 20)
d = pd.read_stata(sys.argv[1], convert_categoricals=False)

d["lwk1cpi"] = np.log(d.wk1cpi)
d["lwk1per"] = np.log(d.wk1cpi / d.openscr)
d["lgross"] = np.log(d.grosscpi)
d["screens"] = d.openscr / 1000
d["weekend"] = d.opendate + 6 - d.openwday
d = d[d.pre != 1].copy()
print("N after dropping pre-opening reviews:", len(d))

GENRE = "advent animated children comedy crime document drama fantasy noir horror musical mystery romance scifi thrill war western".split()
PROD = "disney fox sony mgm_ua gram orion univers warner paramnt".split()
YEARS = [f"year{i}" for i in range(2, 10)]
M11 = [f"month{i}" for i in range(1, 12)]
M10 = [f"month{i}" for i in range(1, 11)]
KEY = ["dur", "oudur", "uudur", "ou", "uu"]


def design(df, xvars):
    X = df[xvars].astype(float)
    X = X.loc[:, X.std() > 0]          # Stata drops all-constant columns
    X = sm.add_constant(X)
    # drop exactly collinear columns, as Stata does, keeping earlier ones
    keep = []
    for c in X.columns:
        trial = X[keep + [c]].values
        if np.linalg.matrix_rank(trial) == len(keep) + 1:
            keep.append(c)
    return X[keep]


def ols(df, y, xvars):
    X = design(df, xvars)
    res = sm.OLS(df[y].astype(float), X).fit(
        cov_type="cluster", cov_kwds={"groups": df.weekend.values})
    return res


def probit(df, y, xvars):
    df = df.copy()
    xs = list(xvars)
    # Stata: a dummy that perfectly predicts the outcome is dropped with its observations
    changed = True
    while changed:
        changed = False
        for c in xs:
            v = df[c]
            if set(v.unique()) <= {0, 1} and v.sum() > 0:
                outs = df.loc[v == 1, y].unique()
                if len(outs) == 1:
                    df = df[v != 1]
                    xs.remove(c)
                    changed = True
                    break
    X = design(df, xs)
    res = sm.Probit(df[y].astype(float), X).fit(
        disp=0, cov_type="cluster", cov_kwds={"groups": df.weekend.values})
    return res


def col(res, names):
    out = {}
    for n in names:
        if n in res.params:
            out[n] = f"{res.params[n]:6.2f} ({res.bse[n]:.2f}){stars(res.pvalues[n])}"
        else:
            out[n] = ""
    out["N"] = int(res.nobs)
    return out


def stars(p):
    return "***" if p < .01 else "**" if p < .05 else "*" if p < .10 else ""


# ---------------- Table I ----------------
t1 = pd.DataFrame({
    "TOTREV": d.grosscpi / 1e6, "OPENREV": d.wk1cpi / 1e6, "SISKEL UP": d.sdummy,
    "EBERT UP": d.edummy, "ONE UP": d.ou, "TWO UP": d.uu, "DURING": d.dur,
    "SCREENS": d.screens, "FOURDAY": d.fourday, "MALTIN": d.maltin})
print("\nTABLE I  (paper means: 30.6 6.6 .32 .41 .31 .21 .81 1.49 .07 2.29)")
print(t1.describe().T[["mean", "std", "max", "min"]].round(2))

# ---------------- Table II ----------------
t2 = pd.DataFrame({"lnTOTREV": np.log(d.grosscpi), "lnOPENREV": d.lwk1cpi, "MALTIN": d.maltin,
                   "SISKEL": d.sdummy, "EBERT": d.edummy, "DURING": d.dur})
print("\nTABLE II correlations")
print(t2.corr().round(2))

# ---------------- Table III ----------------
base = GENRE + PROD
sub = d[d.month != 12]
p = {
    "(1)": probit(sub, "dur", ["lwk1cpi", "sdummy", "edummy"] + base + ["maltin"] + YEARS + ["screens", "fourday"] + M10),
    "(2)": probit(sub, "dur", ["lwk1cpi", "ou", "uu"] + base + ["maltin"] + YEARS + ["screens", "fourday"] + M10),
    "(3)": probit(sub, "dur", ["lwk1cpi", "sdummy", "edummy"] + base + YEARS + ["screens", "fourday"] + M10),
    "(4)": probit(sub, "dur", ["lwk1cpi", "ou", "uu"] + base + YEARS + ["screens", "fourday"] + M10),
    "(5)": probit(sub, "dur", YEARS + ["fourday"] + M10),
}
names3 = ["lwk1cpi", "sdummy", "edummy", "ou", "uu", "maltin", "screens", "fourday"]
t3 = pd.DataFrame({k: col(v, names3) for k, v in p.items()})
t3.loc["loglik"] = [round(v.llf) for v in p.values()]
t3.loc["pseudoR2"] = [round(v.prsquared, 2) for v in p.values()]
print("\nTABLE III probit for DURING")
print(t3)

# ---------------- Table IV ----------------
ctrl_all = ["maltin"] + GENRE + PROD + YEARS + ["screens", "fourday"] + M11
ctrl_gen = ["maltin"] + PROD + YEARS + ["screens", "fourday"] + M11
spec = KEY
m4 = {
    "All": ols(d, "lwk1cpi", spec + ctrl_all),
    "Wide": ols(d[d.openscr > 1533], "lwk1cpi", spec + ctrl_all),
    "Narrow": ols(d[d.openscr <= 1533], "lwk1cpi", spec + ctrl_all),
    "Drama": ols(d[d.drama == 1], "lwk1cpi", spec + ctrl_gen),
    "Action": ols(d[d.action == 1], "lwk1cpi", spec + ctrl_gen),
    "Comedy": ols(d[d.comedy == 1], "lwk1cpi", spec + ctrl_gen),
}
names4 = KEY + ["maltin", "screens", "fourday", "const"]
t4 = pd.DataFrame({k: col(v, names4) for k, v in m4.items()})
t4.loc["R2"] = [round(v.rsquared, 2) for v in m4.values()]
print("\nTABLE IV ln(OPENREV)")
print(t4)

# ---------------- Table V ----------------
dr = d[d.drama == 1]
m5 = {
    "(1) lnOPEN": ols(dr, "lwk1cpi", spec + ctrl_gen),
    "(2) no MALTIN": ols(dr, "lwk1cpi", spec + [c for c in ctrl_gen if c != "maltin"]),
    "(3) per screen": ols(dr, "lwk1per", spec + [c for c in ctrl_gen if c != "screens"]),
    "(4) lnTOTREV": ols(dr, "lgross", spec + ctrl_gen),
}
t5 = pd.DataFrame({k: col(v, names4) for k, v in m5.items()})
t5.loc["R2"] = [round(v.rsquared, 2) for v in m5.values()]
print("\nTABLE V dramas")
print(t5)

# ---------------- machine-readable key coefficients ----------------
rows = []
for tab, models in (("IV", m4), ("V", m5)):
    for k, r in models.items():
        for v in KEY:
            rows.append(dict(table=tab, model=k, var=v, b=r.params[v], se=r.bse[v],
                             p=r.pvalues[v], N=int(r.nobs), r2=r.rsquared))
OUT = Path(__file__).resolve().parent / "output"
OUT.mkdir(exist_ok=True)
pd.DataFrame(rows).to_csv(OUT / "keycoefs_tables4_5.csv", index=False)
