"""Re-check of the diamond-price regression (github.com/alexly0329-cpu/diamond-price-regression).

Reproduces every figure quoted on the portfolio case study:
  * R-squared for each specification, centered and (for the no-intercept model) uncentered
  * coefficients with and without the intercept
  * the depth_pct reconciliation check
  * a 75/25 hold-out comparison (seed 42)

Run:  pip install pandas numpy statsmodels scikit-learn
      python diamond-regression-recheck.py
"""
import numpy as np
import pandas as pd
import statsmodels.formula.api as smf
from sklearn.metrics import mean_absolute_percentage_error as mape
from sklearn.metrics import r2_score
from sklearn.model_selection import train_test_split

URL = "https://raw.githubusercontent.com/mwaskom/seaborn-data/master/diamonds.csv"

df = pd.read_csv(URL).rename(
    columns={"x": "length_mm", "y": "width_mm", "z": "depth_mm", "depth": "depth_pct"}
)

# --- same transformations and ordinal encodings as the notebook -------------------
df["log_price"] = np.log1p(df["price"])
df["log_carat"] = np.log1p(df["carat"])
df["log_length_mm"] = np.log1p(df["length_mm"])
orders = {
    "cut": ["Fair", "Good", "Very Good", "Premium", "Ideal"],
    "color": ["J", "I", "H", "G", "F", "E", "D"],
    "clarity": ["I1", "SI2", "SI1", "VS2", "VS1", "VVS2", "VVS1", "IF"],
}
for col, order in orders.items():
    df[f"{col}_code"] = df[col].astype(pd.CategoricalDtype(order, ordered=True)).cat.codes

QUALITY = "log_length_mm + cut_code + color_code + clarity_code"
SPECS = {
    "log_price ~ carat": "log_price ~ carat",
    "log_price ~ log_carat": "log_price ~ log_carat",
    "log_price ~ log_length_mm": "log_price ~ log_length_mm",
    "final, no intercept (notebook)": f"log_price ~ {QUALITY} - 1",
    "final, with intercept": f"log_price ~ {QUALITY}",
}

print(f"observations: {len(df):,}\n")
print(f"{'specification':34s} {'R2 (statsmodels)':>17s} {'R2 centered':>12s} {'cond. no.':>10s}")
for name, formula in SPECS.items():
    fit = smf.ols(formula, df).fit()
    resid = df["log_price"] - fit.fittedvalues
    centered = 1 - (resid**2).sum() / ((df["log_price"] - df["log_price"].mean()) ** 2).sum()
    print(f"{name:34s} {fit.rsquared:17.3f} {centered:12.3f} {fit.condition_number:10.1f}")

print("\ncoefficients")
for name in ("final, no intercept (notebook)", "final, with intercept"):
    params = smf.ols(SPECS[name], df).fit().params.round(4)
    print(f"  {name}: {params.to_dict()}")

# --- depth_pct reconciliation ----------------------------------------------------
zero_dim = (df["length_mm"] == 0) | (df["width_mm"] == 0) | (df["depth_mm"] == 0)
implied = (2 * df["depth_mm"] / (df["length_mm"] + df["width_mm"]) * 100).where(~zero_dim)
match = np.isclose(implied, df["depth_pct"], atol=0.1)
gap = (implied - df["depth_pct"]).abs()[~match & ~zero_dim]
print("\ndepth_pct = 2 * depth_mm / (length_mm + width_mm) * 100, tolerance 0.1 pt")
print(f"  reconciles: {match.sum():,} rows ({match.mean():.1%})")
print(f"  does not:   {(~match).sum():,} rows ({(~match).mean():.1%}), of which {zero_dim.sum()} have a zero dimension")
print(f"  gap on the rest: median {gap.median():.2f} pt, 90th pct {gap.quantile(0.9):.2f} pt")

# --- hold-out ---------------------------------------------------------------------
train, test = train_test_split(df, test_size=0.25, random_state=42)
print("\nhold-out (75/25, seed 42)")
for name in ("final, no intercept (notebook)", "final, with intercept"):
    pred = smf.ols(SPECS[name], train).fit().predict(test)
    print(
        f"  {name:32s} R2 (log price) {r2_score(test['log_price'], pred):.3f}"
        f" | MAPE on dollar price {mape(test['price'], np.expm1(pred)):.1%}"
    )