import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings("ignore")

# 항공 승객 데이터 로드
from statsmodels.datasets import get_rdataset
air = get_rdataset("AirPassengers")
df = air.data.copy()
df.columns = ["date", "passengers"]
df["date"] = pd.to_datetime(df["date"])
df.set_index("date", inplace=True)

print(f"데이터 기간: {df.index.min()} ~ {df.index.max()}")
print(f"데이터 포인트 수: {len(df)}")

# 시계열 시각화
fig, axes = plt.subplots(3, 1, figsize=(14, 10))

# 원본 시계열
axes[0].plot(df["passengers"])
axes[0].set_title("항공 승객 수 (원본)")
axes[0].set_ylabel("승객 수")

# 로그 변환 (분산 안정화)
axes[1].plot(np.log(df["passengers"]))
axes[1].set_title("로그 변환")
axes[1].set_ylabel("log(승객 수)")

# 차분 (추세 제거)
axes[2].plot(np.log(df["passengers"]).diff().dropna())
axes[2].set_title("로그 + 1차 차분")
axes[2].set_ylabel("차분값")

plt.tight_layout()
plt.show()

from statsmodels.tsa.seasonal import seasonal_decompose
from statsmodels.tsa.stattools import adfuller

# 시계열 분해 (추세 + 계절성 + 잔차)
decomposition = seasonal_decompose(df["passengers"], model="multiplicative", period=12)

fig, axes = plt.subplots(4, 1, figsize=(14, 10))
decomposition.observed.plot(ax=axes[0], title="원본")
decomposition.trend.plot(ax=axes[1], title="추세 (Trend)")
decomposition.seasonal.plot(ax=axes[2], title="계절성 (Seasonal)")
decomposition.resid.plot(ax=axes[3], title="잔차 (Residual)")
plt.tight_layout()
plt.show()

# ADF 정상성 검정
def adf_test(series, name=""):
    """ADF(Augmented Dickey-Fuller) 정상성 검정"""
    result = adfuller(series.dropna())
    print(f"=== {name} ADF 검정 ===")
    print(f"검정통계량: {result[0]:.4f}")
    print(f"p-value: {result[1]:.4f}")
    print(f"결론: {'정상 시계열' if result[1] < 0.05 else '비정상 시계열'}")
    print()

adf_test(df["passengers"], "원본")
adf_test(np.log(df["passengers"]).diff().dropna(), "로그 + 1차 차분")

# 시계열은 시간 순서를 유지하며 분할 (셔플 금지!)
train_size = int(len(df) * 0.8)
train = df[:train_size]
test = df[train_size:]

print(f"학습 기간: {train.index.min()} ~ {train.index.max()} ({len(train)}개)")
print(f"테스트 기간: {test.index.min()} ~ {test.index.max()} ({len(test)}개)")

plt.figure(figsize=(14, 5))
plt.plot(train["passengers"], label="학습")
plt.plot(test["passengers"], label="테스트")
plt.axvline(x=test.index[0], color="r", linestyle="--", label="분할점")
plt.title("학습/테스트 분할")
plt.legend()
plt.show()

from statsmodels.tsa.arima.model import ARIMA
from statsmodels.tsa.statespace.sarimax import SARIMAX
from sklearn.metrics import mean_absolute_error, mean_squared_error

# ACF/PACF 확인 (파라미터 결정)
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf

fig, axes = plt.subplots(1, 2, figsize=(14, 4))
log_diff = np.log(train["passengers"]).diff().dropna()
plot_acf(log_diff, ax=axes[0], lags=30)
plot_pacf(log_diff, ax=axes[1], lags=30)
plt.tight_layout()
plt.show()

# SARIMAX 모델 (계절성 포함)
model_sarimax = SARIMAX(
    train["passengers"],
    order=(1, 1, 1),           # (p, d, q)
    seasonal_order=(1, 1, 1, 12),  # (P, D, Q, S) S=12 (월별 계절성)
    enforce_stationarity=False,
    enforce_invertibility=False,
)
result_sarimax = model_sarimax.fit(disp=False)
print(result_sarimax.summary().tables[1])

# 예측
forecast_sarimax = result_sarimax.forecast(steps=len(test))

# 평가
mae_sarimax = mean_absolute_error(test["passengers"], forecast_sarimax)
rmse_sarimax = np.sqrt(mean_squared_error(test["passengers"], forecast_sarimax))
print(f"\nSARIMAX - MAE: {mae_sarimax:.2f}, RMSE: {rmse_sarimax:.2f}")

# 시각화
plt.figure(figsize=(14, 5))
plt.plot(train["passengers"], label="학습")
plt.plot(test["passengers"], label="실제")
plt.plot(test.index, forecast_sarimax, label="SARIMAX 예측", linestyle="--")
plt.title("SARIMAX 예측 결과")
plt.legend()
plt.show()

from lightgbm import LGBMRegressor

# 시계열 특성 생성 함수
def create_time_features(df):
    """시계열 데이터에서 ML 학습용 특성 생성"""
    df = df.copy()
    idx = df.index

    # 날짜 기반 특성
    df["month"] = idx.month
    df["year"] = idx.year
    df["quarter"] = idx.quarter

    # 주기적 인코딩 (월의 순환성 반영)
    df["month_sin"] = np.sin(2 * np.pi * idx.month / 12)
    df["month_cos"] = np.cos(2 * np.pi * idx.month / 12)

    # 래그 특성 (과거 값 활용)
    for lag in [1, 2, 3, 6, 12]:
        df[f"lag_{lag}"] = df["passengers"].shift(lag)

    # 이동 평균
    df["rolling_mean_3"] = df["passengers"].shift(1).rolling(3).mean()
    df["rolling_mean_6"] = df["passengers"].shift(1).rolling(6).mean()
    df["rolling_mean_12"] = df["passengers"].shift(1).rolling(12).mean()

    # 이동 표준편차
    df["rolling_std_6"] = df["passengers"].shift(1).rolling(6).std()

    # 전년 동월 대비 변화율
    df["yoy_change"] = df["passengers"].pct_change(12)

    return df.dropna()

# 특성 생성
df_features = create_time_features(df)

# 학습/테스트 분할 (시간 순서 유지)
train_feat = df_features[:train_size]
test_feat = df_features[train_size:]

# 특성과 타겟 분리
feature_cols = [c for c in df_features.columns if c != "passengers"]
X_train_ml = train_feat[feature_cols]
y_train_ml = train_feat["passengers"]
X_test_ml = test_feat[feature_cols]
y_test_ml = test_feat["passengers"]

# LightGBM 학습
lgbm = LGBMRegressor(
    n_estimators=300,
    max_depth=6,
    learning_rate=0.05,
    num_leaves=31,
    random_state=42,
    verbose=-1,
)
lgbm.fit(X_train_ml, y_train_ml)

# 예측 및 평가
pred_lgbm = lgbm.predict(X_test_ml)
mae_lgbm = mean_absolute_error(y_test_ml, pred_lgbm)
rmse_lgbm = np.sqrt(mean_squared_error(y_test_ml, pred_lgbm))
print(f"LightGBM - MAE: {mae_lgbm:.2f}, RMSE: {rmse_lgbm:.2f}")

# 특성 중요도
feat_imp = pd.Series(lgbm.feature_importances_, index=feature_cols)
feat_imp.sort_values(ascending=True).tail(10).plot(kind="barh", figsize=(10, 5))
plt.title("특성 중요도 (상위 10개)")
plt.xlabel("중요도")
plt.tight_layout()
plt.show()

from sklearn.model_selection import TimeSeriesSplit

# 시계열 교차검증 (시간 순서 보존)
tscv = TimeSeriesSplit(n_splits=5)

cv_scores = []
for fold, (train_idx, val_idx) in enumerate(tscv.split(df_features)):
    train_cv = df_features.iloc[train_idx]
    val_cv = df_features.iloc[val_idx]

    X_tr = train_cv[feature_cols]
    y_tr = train_cv["passengers"]
    X_val = val_cv[feature_cols]
    y_val = val_cv["passengers"]

    lgbm_cv = LGBMRegressor(
        n_estimators=300, max_depth=6,
        learning_rate=0.05, random_state=42, verbose=-1,
    )
    lgbm_cv.fit(X_tr, y_tr)
    pred_cv = lgbm_cv.predict(X_val)

    mae_cv = mean_absolute_error(y_val, pred_cv)
    cv_scores.append(mae_cv)
    print(f"Fold {fold+1}: MAE = {mae_cv:.2f}")

print(f"\n평균 MAE: {np.mean(cv_scores):.2f} (+/- {np.std(cv_scores):.2f})")

# 최종 비교 시각화
plt.figure(figsize=(14, 6))
plt.plot(train["passengers"], label="학습", color="gray", alpha=0.5)
plt.plot(test["passengers"], label="실제", color="black", linewidth=2)
plt.plot(test.index, forecast_sarimax.values,
         label=f"SARIMAX (MAE={mae_sarimax:.1f})", linestyle="--")
plt.plot(test_feat.index, pred_lgbm,
         label=f"LightGBM (MAE={mae_lgbm:.1f})", linestyle="--")
plt.title("SARIMAX vs LightGBM 예측 비교")
plt.xlabel("날짜")
plt.ylabel("승객 수")
plt.legend()
plt.show()

# 결과 요약 테이블
comparison = pd.DataFrame({
    "모델": ["SARIMAX", "LightGBM"],
    "MAE": [mae_sarimax, mae_lgbm],
    "RMSE": [rmse_sarimax, rmse_lgbm],
})
print("\n=== 모델 비교 요약 ===")
print(comparison.to_string(index=False))