import pandas as pd
import numpy as np
import statsmodels.api as sm
from statsmodels.stats.outliers_influence import variance_inflation_factor

# ── Load & clean ─────────────────────────────────────────────
df = pd.read_csv('DATA.tsv', sep='\t', on_bad_lines='skip', engine='python')

df_work = pd.DataFrame({
    "COUNTRY":   df["COUNTRY"],
    "PRICE":     df["HH_PRICES"],
    "DC":        df["DC/CAPITA"],
    "GDP":       df["GDP"],
})

for col in ["PRICE", "DC", "GDP"]:
    df_work[col] = pd.to_numeric(
        df_work[col].astype(str).str.replace(',', '.').to_numpy(dtype=np.float64),
        errors='coerce'
    )

#── IQR outlier filter (on PRICE) ─────────────────────────────
# ── Base datasets ─────────────────────────
df_m1 = df_work.dropna(subset=["PRICE", "DC"]).copy()
df_m2 = df_work.dropna(subset=["PRICE", "DC", "GDP"]).copy()

# ── IQR filter ───────────────────────────
def iqr_mask(arr):
    q1, q3 = np.percentile(arr, 25), np.percentile(arr, 75)
    iqr = q3 - q1
    return (arr < q1 - 1.5 * iqr) | (arr > q3 + 1.5 * iqr)


# MODEL 1 mask (PRICE OR DC outlier → remove)
mask_m1 = iqr_mask(df_m1["PRICE"]) | iqr_mask(df_m1["DC"])
df_m1_clean = df_m1[~mask_m1]

# MODEL 2 mask (any outlier → remove)
mask_m2 = (
    iqr_mask(df_m2["PRICE"]) |
    iqr_mask(df_m2["DC"]) |
    iqr_mask(df_m2["GDP"])
)
df_m2_clean = df_m2[~mask_m2]



# ── MODEL 1 ──────────────────────────────
X1 = sm.add_constant(df_m1_clean["DC"])
y1 = df_m1_clean["PRICE"]
m1 = sm.OLS(y1, X1).fit()

# ── MODEL 2 ──────────────────────────────
X2 = sm.add_constant(df_m2_clean[["DC", "GDP"]])
y2 = df_m2_clean["PRICE"]
m2 = sm.OLS(y2, X2).fit()

# ── VIF (Model 2) ────────────────────────────────────────────
vif_df = pd.DataFrame({
    "Variable": ["DC", "GDP"],
    "VIF": [
        variance_inflation_factor(X2.values, 1),
        variance_inflation_factor(X2.values, 2)
    ]
})

# ── Output ───────────────────────────────────────────────────
sep = "─" * 60

print(f"\n{sep}")
print("MODEL 1 — PRICE ~ DC")
print(sep)
print(m1.summary())

print(f"\n{sep}")
print("MODEL 2 — PRICE ~ DC + GDP")
print(sep)
print(m2.summary())

print(f"\n{sep}")
print("VIF — MULTICOLLINEARITY (Model 2)")
print(sep)

for _, row in vif_df.iterrows():
    if row["VIF"] > 10:
        flag = "⚠ HIGH"
    elif row["VIF"] > 5:
        flag = "△ moderate"
    else:
        flag = "✓ OK"
    print(f"{row['Variable']:<10} VIF = {row['VIF']:.3f}  {flag}")