import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

# ── 데이터 로드 (UCI 직접 또는 S3)
URL = "https://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-red.csv"
df = pd.read_csv(URL, sep=";")
print(f"shape: {df.shape}")
print(df.describe().round(2))
print("결측치:", df.isnull().sum().sum())

# ── 품질 분포 + 상관관계 히트맵
fig, axes = plt.subplots(1, 2, figsize=(12, 4))
df["quality"].value_counts().sort_index().plot(
    kind="bar", ax=axes[0], color="#8B2252", edgecolor="white")
axes[0].set_title("Quality Score Distribution")
sns.heatmap(df.corr(), ax=axes[1], cmap="RdBu_r",
    center=0, annot=False, vmin=-1, vmax=1)
axes[1].set_title("Feature Correlation Heatmap")
plt.tight_layout(); plt.savefig("output/eda_overview.png", dpi=150)
plt.show()

from sklearn.model_selection import train_test_split

# ── low: ≤5 / mid: 6–7 / high: ≥8
def to_label(q):
    if q <= 5:    return "low"
    elif q <= 7:  return "mid"
    else:         return "high"

df["quality_label"] = df["quality"].apply(to_label)
print(df["quality_label"].value_counts())

# ── 피처 / 타겟 분리
FEATURES = [c for c in df.columns if c not in ["quality", "quality_label"]]
X = df[FEATURES]
y = df["quality_label"]

# ── 80/10/10 분리 (stratify로 클래스 비율 유지)
X_train, X_tmp, y_train, y_tmp = train_test_split(
    X, y, test_size=0.2, stratify=y, random_state=42)
X_val, X_test, y_val, y_test = train_test_split(
    X_tmp, y_tmp, test_size=0.5, stratify=y_tmp, random_state=42)

print(f"Train: {len(X_train)} | Val: {len(X_val)} | Test: {len(X_test)}")

from interpret.glassbox import ExplainableBoostingClassifier
import pickle

# ── EBM 학습 (3-class 분류)
ebm = ExplainableBoostingClassifier(
    max_bins=256,
    interactions=3,
    outer_bags=8,
    random_state=42
)
ebm.fit(X_train, y_train)

# ── 검증 성능
from sklearn.metrics import accuracy_score, classification_report
y_pred = ebm.predict(X_val)
print(f"Val Accuracy: {accuracy_score(y_val, y_pred):.3f}")
print(classification_report(y_val, y_pred))

# ── 모델 저장
with open("output/ebm_wine.pkl", "wb") as f:
    pickle.dump(ebm, f)
print("✓ 모델 저장 완료")

from interpret import show

# ── Global 해석 (feature importance + shape functions)
ebm_global = ebm.explain_global(name="Wine EBM")
show(ebm_global)

# ── 특정 피처 shape function 확인
# alcohol, volatile_acidity, sulphates 순서대로 확인하세요

import dice_ml
from dice_ml import Dice

# ── DiCE 데이터 객체 생성
d = dice_ml.Data(
    dataframe=pd.concat([X_train, y_train], axis=1),
    continuous_features=FEATURES,
    outcome_name="quality_label"
)

# ── DiCE 모델 연결 (EBM 사용)
m = dice_ml.Model(model=ebm, backend="sklearn")
exp = Dice(d, m, method="random")

# ── 분석 대상 샘플 선택 (mid-quality 1개)
mid_samples = X_val[y_val == "mid"].head(1)
print(mid_samples)

# ── Actionable 피처: 제조 단계에서 조절 가능한 성분
ACTIONABLE = [
    "alcohol",           # 발효 조절로 변경 가능
    "sulphates",         # 첨가량 조절 가능
    "total sulfur dioxide",  # 이산화황 첨가량 조절
    "free sulfur dioxide",
    "citric acid",       # 첨가 조절 가능
]
# Non-actionable: density, fixed acidity (포도 품종 특성)

cf = exp.generate_counterfactuals(
    mid_samples,
    total_CFs=3,
    desired_class="high",
    features_to_vary=ACTIONABLE
)
cf.visualize_as_dataframe(show_only_changes=True)