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:
    panel = pd.read_excel('../../data/econometrics-and-panel-data--panel.xlsx', sheet_name='Panel')
except FileNotFoundError:
    panel = pd.read_excel(BASE + 'econometrics-and-panel-data--panel.xlsx', sheet_name='Panel')

pooled = ols('log_wage ~ union + tenure', panel).fit()
print(f"union = {pooled.params['union']:.3f}  CI {np.round(pooled.conf_int().loc['union'].values,3)}  (biased up)")

union = 0.544  CI [0.462 0.626]  (biased up)

fe = ols('log_wage ~ union + tenure + C(worker_id)', panel).fit()
print(f"FE union = {fe.params['union']:.3f}  CI {np.round(fe.conf_int().loc['union'].values,3)}  (true ~0.15)")

FE union = 0.152  CI [0.123 0.181]  (true ~0.15)

ctrl = ols('log_wage ~ union + tenure + ability_index', panel).fit()
print(f"control-ability union = {ctrl.params['union']:.3f}  vs  FE = {fe.params['union']:.3f}  (match)")

control-ability union = 0.157  vs  FE = 0.152  (match)

r = panel[['union','ability_index']].corr().iloc[0,1]
print(f'corr(union, ability) = {r:.2f} (>0), and ability raises wage (>0)  ->  pooled OLS biased UPWARD.')

corr(union, ability) = 0.37 (>0), and ability raises wage (>0)  ->  pooled OLS biased UPWARD.

rng=np.random.default_rng(3); n=2000
u=rng.normal(0,1,n); z=rng.normal(0,1,n)
x=0.8*z+0.9*u+rng.normal(0,0.5,n); y=2.0*x+1.5*u+rng.normal(0,0.5,n)
d=pd.DataFrame({'y':y,'x':x,'z':z})
d['x_hat']=ols('x ~ z', d).fit().fittedvalues
print(f"naive OLS = {ols('y ~ x', d).fit().params['x']:.2f} (biased)")
print(f"2SLS      = {ols('y ~ x_hat', d).fit().params['x_hat']:.2f} (true = 2.00)")

naive OLS = 2.77 (biased)
2SLS      = 1.99 (true = 2.00)