"""
04-11  Correlation & Regression Analyzer
Chapters 13-01, 13-02, 13-03

Run from the project folder:   python python/analysis.py
"""
import numpy as np
import pandas as pd
from scipy import stats
import statsmodels.api as sm
from statsmodels.formula.api import ols
from statsmodels.stats.outliers_influence import variance_inflation_factor
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

pd.set_option("display.width", 120)

h = pd.read_csv("data/houses.csv")
n = len(h)

# ── 1. CORRELATION ──────────────────────────────────────────────────────────
print("\n===== 1. CORRELATION MATRIX =====")

num = h[["sqft", "bedrooms", "age_years", "price"]]
print("\nPearson r:")
print(num.corr().round(4).to_string())
print("\nSpearman rho (rank-based, robust to curvature and outliers):")
print(num.corr(method="spearman").round(4).to_string())

print("\nSignificance of each correlation with price:")
for v in ("sqft", "bedrooms", "age_years"):
    res = stats.pearsonr(num[v], num["price"])
    print(f"  {v:<10} r = {res.statistic:+.4f}   r^2 = {res.statistic**2:.4f}   "
          f"p = {res.pvalue:.3e}")

r_sb = h["sqft"].corr(h["bedrooms"])
print("\nWATCH THE PREDICTOR-PREDICTOR CORRELATIONS:")
print(f"  sqft vs bedrooms  r = {r_sb:.4f}"
      + ("   <- high; expect multicollinearity in step 7" if abs(r_sb) > 0.7 else ""))

# ── 2. SIMPLE REGRESSION ────────────────────────────────────────────────────
print("\n===== 2. SIMPLE REGRESSION — price on sqft =====")

m1 = ols("price ~ sqft", data=h).fit()
b0, b1 = m1.params["Intercept"], m1.params["sqft"]
s_e = np.sqrt(m1.mse_resid)

print(f"\nFITTED LINE:  price-hat = {b0:.2f} + {b1:.4f} * sqft")
print("\nSLOPE      each additional square foot is associated with a")
print(f"           ${b1:.2f} increase in predicted price (${100*b1:.0f} per 100 sqft)")
print(f"INTERCEPT  ${b0:.2f} at sqft = 0 -- far outside the observed range")
print(f"           ({h['sqft'].min():.0f} to {h['sqft'].max():.0f} sqft), so it is a "
      "mathematical anchor only")
print(f"r^2        {m1.rsquared:.4f}  -> {100*m1.rsquared:.1f}% of price variation "
      "explained by floor area")
print(f"s_e        ${s_e:.2f}  -> a typical prediction misses by about this much")

# ── 3. INFERENCE ON THE SLOPE ───────────────────────────────────────────────
print("\n===== 3. INFERENCE ON THE SLOPE =====")
print("H0: beta1 = 0   (floor area has no linear predictive value)")
print("H1: beta1 != 0\n")

coef_tab = pd.DataFrame({
    "estimate": m1.params, "std_err": m1.bse,
    "t": m1.tvalues, "p_value": m1.pvalues,
})
print(coef_tab.round(6).to_string())
print(f"\ncritical t(0.025, {int(m1.df_resid)}) = "
      f"{stats.t.ppf(0.975, m1.df_resid):.4f}")

ci = m1.conf_int(alpha=0.05)
ci.columns = ["lower", "upper"]
print("\n95% confidence intervals for the coefficients:")
print(ci.round(4).to_string())
print(f"\n=> each extra square foot is worth between ${ci.loc['sqft','lower']:.2f} "
      f"and ${ci.loc['sqft','upper']:.2f},")
print("   with 95% confidence. That range is what a decision-maker needs.")

print("\nThe three equivalent tests (simple regression only):")
print(f"  slope t          {m1.tvalues['sqft']:.4f}")
print(f"  correlation t    {stats.pearsonr(h['sqft'], h['price']).statistic * np.sqrt((n-2)/(1-h['sqft'].corr(h['price'])**2)):.4f}")
print(f"  sqrt(model F)    {np.sqrt(m1.fvalue):.4f}")

print("\nANOVA decomposition:")
print(sm.stats.anova_lm(m1, typ=1).round(4).to_string())

# ── 4. RESIDUAL DIAGNOSTICS (the LINE assumptions) ──────────────────────────
print("\n===== 4. RESIDUAL DIAGNOSTICS =====")

r = m1.resid
print(f"sum of residuals        {r.sum():.10f}  (must be ~0)")
sw = stats.shapiro(r)
print(f"Shapiro-Wilk on resid   p = {sw.pvalue:.4f}  -> "
      + ("NORMALITY looks fine" if sw.pvalue > 0.05 else "normality is questionable"))
cor_rf = np.corrcoef(np.abs(r), m1.fittedvalues)[0, 1]
print(f"cor(|resid|, fitted)    {cor_rf:+.4f}  -> "
      + ("EQUAL VARIANCE looks fine" if abs(cor_rf) < 0.2
         else "possible heteroscedasticity -- inspect the plot"))
bp = sm.stats.diagnostic.het_breuschpagan(r, m1.model.exog)
print(f"Breusch-Pagan           p = {bp[1]:.4f}")
print(f"Durbin-Watson           {sm.stats.durbin_watson(r):.4f}  (near 2 = independent)")

cooks = m1.get_influence().cooks_distance[0]
infl = np.where(cooks > 4 / n)[0]
print(f"influential points (Cook's D > 4/n): {len(infl)}")
if len(infl):
    print(h.iloc[infl][["house_id", "sqft", "price"]].to_string(index=False))

# ── 5. CONFIDENCE vs PREDICTION INTERVAL ────────────────────────────────────
print("\n===== 5. CONFIDENCE vs PREDICTION INTERVAL at sqft = 2000 =====")

pred = m1.get_prediction(pd.DataFrame({"sqft": [2000]})).summary_frame(alpha=0.05)
fit = pred["mean"].iloc[0]
ci_lo, ci_hi = pred["mean_ci_lower"].iloc[0], pred["mean_ci_upper"].iloc[0]
pi_lo, pi_hi = pred["obs_ci_lower"].iloc[0], pred["obs_ci_upper"].iloc[0]

print(f"point estimate            ${fit:,.2f}")
print(f"95% CONFIDENCE interval   ${ci_lo:,.2f} to ${ci_hi:,.2f}   "
      f"(width ${ci_hi-ci_lo:,.2f})")
print(f"95% PREDICTION interval   ${pi_lo:,.2f} to ${pi_hi:,.2f}   "
      f"(width ${pi_hi-pi_lo:,.2f})")
print(f"\nThe prediction interval is {(pi_hi-pi_lo)/(ci_hi-ci_lo):.1f}x wider, "
      "because it carries BOTH the")
print("uncertainty in the line AND the house-to-house scatter.")
print("Quote the CONFIDENCE interval for 'the average 2000 sqft house';")
print("quote the PREDICTION interval for 'this particular house'.")

# ── 6. MULTIPLE REGRESSION ──────────────────────────────────────────────────
print("\n===== 6. MULTIPLE REGRESSION =====")

m2 = ols("price ~ sqft + bedrooms", data=h).fit()
m3 = ols("price ~ sqft + bedrooms + age_years", data=h).fit()
m4 = ols("price ~ sqft + bedrooms + age_years + C(garage)", data=h).fit()

print("\nFull model (m4):")
print(pd.DataFrame({"estimate": m4.params, "std_err": m4.bse,
                    "t": m4.tvalues, "p_value": m4.pvalues}).round(4).to_string())
print(f"\nR^2 = {m4.rsquared:.4f}    Adjusted R^2 = {m4.rsquared_adj:.4f}    "
      f"s_e = {np.sqrt(m4.mse_resid):.2f}")
print(f"F({int(m4.df_model)}, {int(m4.df_resid)}) = {m4.fvalue:.3f}, "
      f"p = {m4.f_pvalue:.3e}")

print("\nHOW TO READ EACH COEFFICIENT: the change in predicted price per")
print("one-unit rise in that predictor, HOLDING THE OTHERS CONSTANT.")
print("'C(garage)[T.Yes]' is the difference between a house with a garage and one")
print("without, at the same sqft, bedrooms and age. (No = the reference level.)")

# ── 7. MODEL COMPARISON ─────────────────────────────────────────────────────
print("\n===== 7. MODEL COMPARISON =====")

mods = {"m1": m1, "m2": m2, "m3": m3, "m4": m4}
print(pd.DataFrame({
    "model": list(mods),
    "k": [int(m.df_model) for m in mods.values()],
    "r2": [round(m.rsquared, 4) for m in mods.values()],
    "adj_r2": [round(m.rsquared_adj, 4) for m in mods.values()],
    "s_e": [round(np.sqrt(m.mse_resid), 2) for m in mods.values()],
    "AIC": [round(m.aic, 2) for m in mods.values()],
}).to_string(index=False))

print("\nCompare with ADJUSTED R^2 or AIC, never with raw R^2 --")
print("raw R^2 can only rise as predictors are added.")
print("\nNested F-tests:")
print(sm.stats.anova_lm(m1, m2, m3, m4).round(4).to_string())

# ── 8. MULTICOLLINEARITY ────────────────────────────────────────────────────
print("\n===== 8. MULTICOLLINEARITY (VIF) =====")

X = sm.add_constant(h[["sqft", "bedrooms", "age_years"]])
vif = pd.Series([variance_inflation_factor(X.values, i) for i in range(X.shape[1])],
                index=X.columns).drop("const")
print(vif.round(3).to_string())
print("\nVIF < 5 fine   5-10 investigate   > 10 serious")
if (vif > 5).any():
    print("=> " + ", ".join(vif.index[vif > 5])
          + " show inflated variance. A non-significant t here may reflect")
    print("   collinearity, NOT irrelevance -- check before dropping anything.")
else:
    print("=> No serious multicollinearity.")

# ── PLOTS ───────────────────────────────────────────────────────────────────
fig, axes = plt.subplots(2, 3, figsize=(15, 9))
ax = axes.ravel()

ax[0].scatter(h["sqft"], h["price"], color="#5B2A86")
grid = pd.DataFrame({"sqft": np.linspace(h["sqft"].min(), h["sqft"].max(), 100)})
pr = m1.get_prediction(grid).summary_frame(alpha=0.05)
ax[0].plot(grid["sqft"], pr["mean"], color="#0FA3A3", lw=2, label="fit")
ax[0].plot(grid["sqft"], pr["mean_ci_lower"], "--", color="#0FA3A3", label="95% CI")
ax[0].plot(grid["sqft"], pr["mean_ci_upper"], "--", color="#0FA3A3")
ax[0].plot(grid["sqft"], pr["obs_ci_lower"], ":", color="#8A5FBF", label="95% PI")
ax[0].plot(grid["sqft"], pr["obs_ci_upper"], ":", color="#8A5FBF")
ax[0].set(title=f"price vs sqft  (r = {h['sqft'].corr(h['price']):.3f})",
          xlabel="sqft", ylabel="price")
ax[0].legend()

ax[1].scatter(m1.fittedvalues, r, color="#5B2A86")
ax[1].axhline(0, ls="--", color="grey")
ax[1].set(title="Residuals vs fitted", xlabel="fitted", ylabel="residual")

stats.probplot(r, dist="norm", plot=ax[2])
ax[2].set_title("Normal Q-Q of residuals")

ax[3].hist(r, bins=20, color="#0B7A7A", edgecolor="white")
ax[3].set(title="Residual histogram", xlabel="residual")

groups = [g["price"].to_numpy() for _, g in h.groupby("garage")]
ax[4].boxplot(groups, tick_labels=list(h.groupby("garage").groups.keys()))
ax[4].set(title="Price by garage", ylabel="price")

ax[5].scatter(m4.fittedvalues, m4.resid, color="#0B7A7A")
ax[5].axhline(0, ls="--", color="grey")
ax[5].set(title="Full model: residuals vs fitted", xlabel="fitted", ylabel="residual")

plt.tight_layout()
plt.savefig("regression_plots.png", dpi=110)
print("\nWrote regression_plots.png")