"""Claim 4: on Mushrooms (22 categorical variables, hypergrid ~10^14) the exact
functional ANOVA decomposition achieves near-perfect reconstruction (R^2 ~ 1,
MSE ~ 1e-15) in ~0.3 s. The prior revision did not rerun this benchmark.

Functional ANOVA on a categorical hypergrid:
    f(x) = mu + sum_j f_j(x_j) + sum_{j<k} f_{jk}(x_j,x_k) + ...
Each order is spanned by centred indicator contrasts, so the decomposition is a
linear model in an explicit basis and "exact reconstruction" is the statement
that the observed function lies in the span of the retained orders.
"""
import json, time, itertools, numpy as np, pandas as pd

RES = {}


def build(df, target):
    cols = [c for c in df.columns if c != target]
    codes = {c: pd.Categorical(df[c]).codes for c in cols}
    nlev = {c: int(codes[c].max())+1 for c in cols}
    grid = 1.0
    for c in cols: grid *= nlev[c]
    y = (pd.Categorical(df[target]).codes).astype(float)
    return cols, codes, nlev, grid, y


def design(cols, codes, nlev, n, order, maxcols=None):
    """Centred indicator contrasts: drop one level per variable for identifiability."""
    blocks = [np.ones((n, 1))]
    for c in cols:
        k = nlev[c]
        if k < 2: continue
        M = np.zeros((n, k-1))
        M[np.arange(n), np.clip(codes[c], 0, k-2)] = 1.0
        M[codes[c] == k-1] = -1.0/(k-1) if False else 0.0
        M[codes[c] == k-1, :] = -1.0
        blocks.append(M)
    if order >= 2:
        for a, b in itertools.combinations(cols, 2):
            ka, kb = nlev[a], nlev[b]
            if ka < 2 or kb < 2: continue
            A = blocks[1+cols.index(a)]; B = blocks[1+cols.index(b)]
            inter = (A[:, :, None]*B[:, None, :]).reshape(n, -1)
            blocks.append(inter)
            if maxcols and sum(x.shape[1] for x in blocks) > maxcols: break
    return np.concatenate(blocks, axis=1)


def run():
    df = pd.read_pickle("mushrooms.pkl")
    target = df.columns[-1]
    cols, codes, nlev, grid, y = build(df, target)
    n = len(df)
    print("  n=%d, variables=%d, hypergrid=%.3e" % (n, len(cols), grid), flush=True)
    rows = []
    for order, label in ((1, "main effects"), (2, "main + pairwise")):
        t0 = time.perf_counter()
        D = design(cols, codes, nlev, n, order)
        beta, *_ = np.linalg.lstsq(D, y, rcond=None)
        pred = D @ beta
        el = time.perf_counter()-t0
        mse = float(np.mean((y-pred)**2))
        r2 = float(1-np.sum((y-pred)**2)/np.sum((y-y.mean())**2))
        rows.append({"order": order, "label": label, "basis_elements": int(D.shape[1]),
                     "R2": round(r2, 10), "MSE": mse, "seconds": round(el, 4)})
        print("  %-16s basis=%-6d  R^2=%.10f  MSE=%.3e  %.3fs"
              % (label, D.shape[1], r2, mse, el), flush=True)
    RES["claim4_mushrooms"] = {
        "n_rows": n, "n_variables": len(cols),
        "levels_per_variable": {c: nlev[c] for c in cols},
        "hypergrid_size": grid, "rows": rows,
        "paper_reported": {"R2": "~1", "MSE": "~1e-15", "seconds": 0.3},
        "best_R2": max(r["R2"] for r in rows), "best_MSE": min(r["MSE"] for r in rows)}
    json.dump(RES, open("anova_results.json", "w"), indent=1)


if __name__ == "__main__":
    run(); print("DONE")