import numpy as np, pandas as pd, matplotlib.pyplot as plt
from scipy import stats
import seaborn as sns   # seaborn = high-level statistical plots (heatmaps, regplots, pairplots)
import statsmodels.api as sm
from statsmodels.formula.api import ols
from statsmodels.stats.outliers_influence import variance_inflation_factor
from statsmodels.nonparametric.smoothers_lowess import lowess
from sklearn.linear_model import LinearRegression, Ridge, Lasso, LogisticRegression
from sklearn.preprocessing import StandardScaler, PolynomialFeatures
from sklearn.pipeline import make_pipeline
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.metrics import r2_score, mean_squared_error
ORG="#ea580c"; DEEP="#c2410c"; LIGHT="#fdba74"; INK="#1a2138"; GRID="#e6e9f2"; GREEN="#059669"; RED="#ef4444"; AMBER="#d97706"; BLUE="#2563eb"; PUR="#9333ea"; GREY="#94a3b8"; SLATE="#475569"
plt.rcParams.update({"figure.facecolor":"white","axes.facecolor":"white","figure.dpi":110,"font.size":11,
   "axes.edgecolor":GRID,"axes.grid":True,"grid.color":GRID,"axes.axisbelow":True,"axes.spines.top":False,
   "axes.spines.right":False,"axes.titlesize":12,"axes.titleweight":"bold","legend.frameon":False})
sns.set_style("whitegrid")
BASE="https://raw.githubusercontent.com/johnfisher-ai/Statistics-Data-Science-AI-Visual-Book/main/data/"

try:
    wages = pd.read_excel('../../data/the-ols-framework--wages.xlsx', sheet_name='Workers')
except FileNotFoundError:
    wages = pd.read_excel(BASE + 'the-ols-framework--wages.xlsx', sheet_name='Workers')
X = np.column_stack([np.ones(len(wages)), wages['education'], wages['experience']])
y = wages['wage'].values
model = ols('wage ~ education + experience', wages).fit()

beta = np.linalg.solve(X.T@X, X.T@y)
print('matrix    :', np.round(beta,4))
print('statsmodels:', np.round(model.params.values,4))
print('match:', np.allclose(beta, model.params.values))

matrix    : [2.415  2.3758 0.3193]
statsmodels: [2.415  2.3758 0.3193]
match: True

e = y - X@model.params.values
print('X.T @ e =', np.round(X.T@e, 6))
SST=((y-y.mean())**2).sum(); SSE=(e**2).sum(); SSR=((model.fittedvalues-y.mean())**2).sum()
print(f'SST {SST:.1f} = SSR {SSR:.1f} + SSE {SSE:.1f} -> {np.isclose(SST,SSR+SSE)}')

X.T @ e = [-0. -0. -0.]
SST 49392.2 = SSR 28170.8 + SSE 21221.4 -> True

H = X @ np.linalg.inv(X.T@X) @ X.T
print('y-hat = Hy ?', np.allclose(H@y, model.fittedvalues))
print(f'trace(H) = {np.trace(H):.3f}  (= {X.shape[1]} parameters)')
print('leverage (diagonal of H), first 5:', np.round(np.diag(H)[:5], 4))

y-hat = Hy ? True
trace(H) = 3.000  (= 3 parameters)
leverage (diagonal of H), first 5: [0.0144 0.0111 0.0046 0.004  0.0062]

rng = np.random.default_rng(1); n=40; b_ols=[]; b_two=[]
for _ in range(3000):
    x=np.linspace(0,10,n); yy=2+3*x+rng.normal(0,4,n)
    b_ols.append(np.polyfit(x,yy,1)[0]); b_two.append((yy[-1]-yy[0])/(x[-1]-x[0]))
print(f'OLS sd       = {np.std(b_ols):.3f}')
print(f'two-point sd = {np.std(b_two):.3f}')
print('Both unbiased (means ~3); OLS has the smaller variance, as Gauss-Markov guarantees.')

OLS sd       = 0.210
two-point sd = 0.561
Both unbiased (means ~3); OLS has the smaller variance, as Gauss-Markov guarantees.

from statsmodels.stats.diagnostic import het_breuschpagan
print(f'Breusch-Pagan p = {het_breuschpagan(model.resid, model.model.exog)[1]:.4f} (heteroscedastic)')
robust = ols('wage ~ education + experience', wages).fit(cov_type='HC3')
print(pd.DataFrame({'coef':model.params.round(3),'SE_classical':model.bse.round(3),'SE_HC3':robust.bse.round(3)}))
print('Report the HC3 standard errors when the equal-variance assumption fails.')

Breusch-Pagan p = 0.0000 (heteroscedastic)
             coef  SE_classical  SE_HC3
Intercept   2.415         2.060   1.764
education   2.376         0.127   0.134
experience  0.319         0.050   0.051
Report the HC3 standard errors when the equal-variance assumption fails.