"""
Multicollinearity check for marketing mix model inputs.

From the Vexpower MMM library: https://vexpower.com/mmm/multicollinearity-in-mmm

Prints, for the predictor columns of a CSV:
  - the correlation matrix
  - the variance inflation factor (VIF) of each predictor
  - the condition number of the standardised predictors, and the
    eigenvector behind the smallest eigenvalue (which columns are tied together)

Usage:
  python vif-check.py                                   # uses mmm-practice-synthetic.csv
  python vif-check.py my_data.csv                       # all numeric columns except the outcome
  python vif-check.py my_data.csv --outcome revenue --drop week promo_code

Requires numpy and pandas.
"""
import argparse
import sys
from pathlib import Path

import numpy as np
import pandas as pd

OUTCOME_GUESSES = {"revenue", "sales", "y", "kpi", "conversions", "orders", "dep_var"}


def vif_table(X: pd.DataFrame) -> pd.DataFrame:
    rows = []
    for col in X.columns:
        y = X[col].to_numpy(float)
        others = X.drop(columns=col).to_numpy(float)
        A = np.column_stack([np.ones(len(y)), others])
        coef, *_ = np.linalg.lstsq(A, y, rcond=None)
        resid = y - A @ coef
        r2 = 1 - resid.var() / y.var()
        vif = np.inf if r2 >= 1 - 1e-12 else 1 / (1 - r2)
        rows.append({"predictor": col, "vif": vif, "r2_on_others": r2, "se_inflation": np.sqrt(vif)})
    return pd.DataFrame(rows).set_index("predictor")


def main():
    default = Path(__file__).with_name("mmm-practice-synthetic.csv")
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("csv", nargs="?", default=str(default))
    p.add_argument("--outcome", default=None, help="outcome column to exclude (guessed if omitted)")
    p.add_argument("--drop", nargs="*", default=[], help="other columns to exclude")
    a = p.parse_args()

    if not Path(a.csv).exists():
        sys.exit(f"Can't find {a.csv}. Download mmm-practice-synthetic.csv from vexpower.com/mmm or pass your own file.")
    df = pd.read_csv(a.csv)
    num = df.select_dtypes("number")
    outcome = a.outcome or next((c for c in num.columns if c.lower() in OUTCOME_GUESSES), None)
    X = num.drop(columns=[c for c in [outcome, *a.drop] if c in num.columns])
    X = X.loc[:, X.std() > 0]
    if X.shape[1] < 2:
        sys.exit("Need at least two non-constant numeric predictor columns.")

    print(f"File: {a.csv}  rows: {len(X)}  outcome excluded: {outcome}")
    print("\nCorrelation matrix")
    print(X.corr().round(2).to_string())

    print("\nVariance inflation factors (5+ watch, 10+ high)")
    t = vif_table(X)
    print(t.round(3).to_string())

    Z = (X - X.mean()) / X.std()
    evals, evecs = np.linalg.eigh(np.corrcoef(Z.to_numpy().T))
    cond = np.sqrt(evals.max() / max(evals.min(), 1e-12))
    print(f"\nCondition number of standardised predictors: {cond:.1f} (on this correlation-matrix scale, 5+ is worth a look and 10+ is serious;\n  Belsley-style condition indices with unit-length columns and an intercept run higher, where 30+ is the usual warning)")
    weakest = pd.Series(evecs[:, 0], index=X.columns)
    print("Columns tied together in the weakest direction (large |loading|):")
    print(weakest.reindex(weakest.abs().sort_values(ascending=False).index).round(2).to_string())


if __name__ == "__main__":
    main()