import pandas as pd
import numpy as np
import joblib
import os
import matplotlib.pyplot as plt
from sqlalchemy import create_engine
from xgboost import XGBRegressor
from sklearn.tree import DecisionTreeRegressor
from sklearn.metrics import mean_absolute_error
from sklearn.metrics import mean_squared_error
from sklearn.metrics import r2_score
from sklearn.metrics import mean_absolute_percentage_error
from sklearn.ensemble import RandomForestRegressor
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.linear_model import LinearRegression

user = "root"
password = ""
host = "localhost"
db = "udd_pmi_module"
engine = create_engine(f"mysql+pymysql://{user}:{password}@{host}/{db}")

# =========================
# FOLDER OUTPUT VISUALISASI
# =========================
output_folder = "machine_learning"
os.makedirs(output_folder, exist_ok=True)

query = """
SELECT 
    m.kode,
    m.tahun,
    m.bulan,
    m.total_pemakaian
FROM item_usage_monthly m
JOIN items i ON m.kode = i.kode
WHERE i.kategori = 0
ORDER BY m.kode, m.tahun, m.bulan;
"""

# =========================
# LOAD DATA
# =========================
df = pd.read_sql(query, engine)
df = df.sort_values(by=["kode", "tahun", "bulan"]).reset_index(drop=True)

# =========================
# VALIDASI DATA
# =========================
df = df.drop_duplicates()
df = df[df["bulan"].between(1, 12)]
df = df[df["tahun"] > 2000]
df = df[df["total_pemakaian"].fillna(0) >= 0]
df = df.dropna(subset=["kode", "tahun", "bulan"])

# =========================
# HANDLE MISSING VALUE (MENGISI NILAI KOSONG MENJADI 0)
# =========================
df["total_pemakaian"] = df["total_pemakaian"].fillna(0)

# =========================
# HANDLE OUTLIER (MEMBATASI PEMAKAIAN MINIMUM 0)
# =========================
df["total_pemakaian"] = df["total_pemakaian"].clip(lower=0)

# =========================
# TYPE CASTING & ROUNDING
# =========================
df["total_pemakaian"] = df["total_pemakaian"].round(0).astype(int)

# =========================
# FILTER BARANG STABIL
# =========================
# std_per_barang = df.groupby("kode")["total_pemakaian"].std()
# mean_per_barang = df.groupby("kode")["total_pemakaian"].mean()
# cv = (std_per_barang / mean_per_barang).fillna(0)
# barang_stabil = cv[cv < 1.0].index
# df = df[df["kode"].isin(barang_stabil)]

# =========================
# LAG FEATURES
# =========================
df["lag_1"] = df.groupby("kode")["total_pemakaian"].shift(1)
df["lag_2"] = df.groupby("kode")["total_pemakaian"].shift(2)
df["lag_3"] = df.groupby("kode")["total_pemakaian"].shift(3)
df["lag_4"] = df.groupby("kode")["total_pemakaian"].shift(4)
df["lag_5"] = df.groupby("kode")["total_pemakaian"].shift(5)
df["lag_6"] = df.groupby("kode")["total_pemakaian"].shift(6)
df["quarter"] = ((df["bulan"] - 1) // 3) + 1
df["is_awal_tahun"] = df["bulan"].isin([1, 2, 3]).astype(int)
df["is_tengah_tahun"] = df["bulan"].isin([6, 7, 8]).astype(int)
df["is_akhir_tahun"] = df["bulan"].isin([10, 11, 12]).astype(int)

# =========================
# DROP NAN
# =========================
df = df[df["lag_6"].notna()]

# =========================
# MOVING AVERAGE
# =========================
df["ma_6"] = df[["lag_1", "lag_2", "lag_3", "lag_4", "lag_5", "lag_6"]].mean(axis=1)

# =========================
# TREND FEATURE
# =========================
df["trend_3"] = df["lag_1"] - df[["lag_2", "lag_3", "lag_4"]].mean(axis=1)

# =========================
# MOMENTUM FEATURE
# =========================
df["momentum_1"] = df["lag_1"] - df["lag_2"]
df["momentum_2"] = df["lag_2"] - df["lag_3"]
df["rolling_std_6"] = df[["lag_1", "lag_2", "lag_3", "lag_4", "lag_5", "lag_6"]].std(
    axis=1
)
df["max_6"] = df[["lag_1", "lag_2", "lag_3", "lag_4", "lag_5", "lag_6"]].max(axis=1)
df["min_6"] = df[["lag_1", "lag_2", "lag_3", "lag_4", "lag_5", "lag_6"]].min(axis=1)

# =========================
# FITUR TAMBAHAN
# =========================
df["range_6"] = df["max_6"] - df["min_6"]
df["cv_6"] = df["rolling_std_6"] / (df["ma_6"] + 1)
df["growth_rate"] = (df["lag_1"] - df["lag_2"]) / (df["lag_2"] + 1)
usage_level, usage_bins = pd.qcut(
    df["ma_6"], q=4, labels=False, retbins=True, duplicates="drop"
)
df["usage_level"] = usage_level

# =========================
# SPLIT TRAINING & TEST (TIME-BASED)
# =========================
df["date"] = pd.to_datetime(
    df["tahun"].astype(str) + "-" + df["bulan"].astype(str) + "-01"
)
df = df.sort_values(["date", "kode"])
split_date = df["date"].quantile(0.8)
train_df = df[df["date"] <= split_date]
test_df = df[df["date"] > split_date]

# =========================
# PREDIKSI
# =========================
features = [
    "lag_1",
    "lag_2",
    "lag_3",
    "lag_4",
    "lag_5",
    "lag_6",
    "ma_6",
    "trend_3",
    "momentum_1",
    "momentum_2",
    "rolling_std_6",
    "max_6",
    "min_6",
    "range_6",
    "cv_6",
    "growth_rate",
    "usage_level",
    "bulan",
    "quarter",
    "is_awal_tahun",
    "is_tengah_tahun",
    "is_akhir_tahun",
]
X_train = train_df[features]
y_train = np.log1p(train_df["total_pemakaian"])
X_test = test_df[features]
y_test = test_df["total_pemakaian"]

# =========================
# MODEL RANDOM FOREST
# =========================
rf_model = RandomForestRegressor(
    n_estimators=300,
    max_depth=15,
    min_samples_split=5,
    min_samples_leaf=2,
    max_features="sqrt",
    bootstrap=True,
    random_state=42,
    n_jobs=-1,
)
rf_model.fit(X_train, y_train)
rf_pred_log = rf_model.predict(X_test)
rf_pred = np.expm1(rf_pred_log)
rf_pred_real = np.round(rf_pred).astype(int)

# =========================
# MODEL LINEAR REGRESSION
# =========================
lr_model = LinearRegression()
lr_model.fit(X_train, y_train)
lr_pred_log = lr_model.predict(X_test)
lr_pred = np.expm1(lr_pred_log)
lr_pred_real = np.round(lr_pred).astype(int)

# =========================
# MODEL DECISION TREE
# =========================
dt_model = DecisionTreeRegressor(
    max_depth=10, min_samples_split=12, min_samples_leaf=6, random_state=42
)
dt_model.fit(X_train, y_train)
dt_pred_log = dt_model.predict(X_test)
dt_pred = np.expm1(dt_pred_log)
dt_pred_real = np.round(dt_pred).astype(int)

# =========================
# CLIP NILAI MINIMUM
# =========================
rf_pred_real = np.clip(rf_pred_real, 0, None)
lr_pred_real = np.clip(lr_pred_real, 0, None)
dt_pred_real = np.clip(dt_pred_real, 0, None)

# =========================
# SAVE LOAD MODEL
# =========================
# joblib.dump(
#     {"model": rf_model, "usage_bins": usage_bins},
#     f"{output_folder}/model.pkl",
#     compress=0,
# )

# =========================
# FUNCTION EVALUASI
# =========================
def smape(y_true, y_pred):
    y_true = np.array(y_true)
    y_pred = np.array(y_pred)
    return (
        np.mean((np.abs(y_true - y_pred) / (np.abs(y_true) + np.abs(y_pred) + 1))) * 100
    )

def evaluate_model(name, y_true, y_pred):
    mae = mean_absolute_error(y_true, y_pred)
    r2 = r2_score(y_true, y_pred)
    mape = mean_absolute_percentage_error(y_true.replace(0, 1), y_pred) * 100
    smape_value = smape(y_true, y_pred)
    print(f"\n=== {name} ===")
    print("MAE   :", round(mae, 2))
    print("R2    :", round(r2, 2))
    print("MAPE  :", round(mape, 2), "%")
    print("SMAPE :", round(smape_value, 2), "%")

# =========================
# EVALUASI SEMUA MODEL
# =========================
evaluate_model("RANDOM FOREST", y_test, rf_pred_real)
evaluate_model("LINEAR REGRESSION", y_test, lr_pred_real)
evaluate_model("DECISION TREE", y_test, dt_pred_real)

# =========================
# HASIL PREDIKSI RANDOM FOREST
# =========================
limit = min(10, len(rf_pred_real))
hasil_tabel = pd.DataFrame(
    {
        "Data": [f"Data {i+1}" for i in range(limit)],
        "Actual": y_test.iloc[:limit].values,
        "RandomForest": rf_pred_real[:limit],
        "LinearRegression": lr_pred_real[:limit],
        "DecisionTree": dt_pred_real[:limit],
    }
)

# =========================
# FUNCTION METRIK VISUAL
# =========================
def get_metrics(y_true, y_pred):
    mae = mean_absolute_error(y_true, y_pred)
    r2 = r2_score(y_true, y_pred)
    mape = mean_absolute_percentage_error(y_true.replace(0, 1), y_pred) * 100
    return mae, mape, r2

# =========================
# AMBIL HASIL METRIK
# =========================
rf_mae, rf_mape, rf_r2 = get_metrics(y_test, rf_pred_real)
lr_mae, lr_mape, lr_r2 = get_metrics(y_test, lr_pred_real)
dt_mae, dt_mape, dt_r2 = get_metrics(y_test, dt_pred_real)
models = ["Random Forest", "Linear Regression", "Decision Tree"]
mae_values = [rf_mae, lr_mae, dt_mae]
mape_values = [rf_mape, lr_mape, dt_mape]
r2_values = [rf_r2, lr_r2, dt_r2]

# =========================
# FUNCTION METRIK VISUAL
# =========================
def get_metrics(y_true, y_pred):
    mae = mean_absolute_error(y_true, y_pred)
    r2 = r2_score(y_true, y_pred)
    mape = mean_absolute_percentage_error(y_true.replace(0, 1), y_pred) * 100
    smape_value = smape(y_true, y_pred)
    return mae, mape, r2, smape_value

# =========================
# AMBIL HASIL METRIK
# =========================
rf_mae, rf_mape, rf_r2, rf_smape = get_metrics(y_test, rf_pred_real)
lr_mae, lr_mape, lr_r2, lr_smape = get_metrics(y_test, lr_pred_real)
dt_mae, dt_mape, dt_r2, dt_smape = get_metrics(y_test, dt_pred_real)
models = ["Random Forest", "Linear Regression", "Decision Tree"]
mae_values = [rf_mae, lr_mae, dt_mae]
mape_values = [rf_mape, lr_mape, dt_mape]
r2_values = [rf_r2, lr_r2, dt_r2]
smape_values = [rf_smape, lr_smape, dt_smape]

# =========================
# VISUALISASI KOMPARASI
# =========================
fig, axes = plt.subplots(2, 2, figsize=(16, 10))
fig.suptitle(
    "Evaluasi dan Komparasi Model Prediksi",
    fontsize=22,
    fontweight="bold",
    y=0.98,
)

# =========================
# STYLE UMUM
# =========================
for ax in axes.flat:
    ax.set_facecolor("#f8fafc")
    for spine in ax.spines.values():
        spine.set_edgecolor("#d1d5db")
        spine.set_linewidth(1)
    ax.tick_params(axis="x", labelsize=11, rotation=0, pad=6)
    ax.tick_params(axis="y", labelsize=11)

# =========================
# R2
# =========================
r2_visual = [max(0, value) for value in r2_values]
bars = axes[0, 0].bar(models, r2_visual, width=0.6)
axes[0, 0].set_title("Perbandingan R²", fontsize=15, pad=12)
axes[0, 0].set_ylabel("R² Score", fontsize=12)
axes[0, 0].set_ylim(0, max(r2_visual) + 0.1)
for i, bar in enumerate(bars):
    visual_height = bar.get_height()
    real_value = r2_values[i]
    axes[0, 0].text(
        bar.get_x() + bar.get_width() / 2,
        visual_height + 0.02,
        round(real_value, 2),
        ha="center",
        va="bottom",
        fontsize=10,
    )

# =========================
# MAE
# =========================
bars = axes[0, 1].bar(models, mae_values, width=0.6)
axes[0, 1].set_title("Perbandingan MAE", fontsize=15, pad=12)
axes[0, 1].set_ylabel("Nilai MAE", fontsize=12)
for bar in bars:
    yval = bar.get_height()
    axes[0, 1].text(
        bar.get_x() + bar.get_width() / 2,
        yval + (max(mae_values) * 0.01),
        round(yval, 2),
        ha="center",
        va="bottom",
        fontsize=10,
    )

# =========================
# MAPE
# =========================
bars = axes[1, 0].bar(models, mape_values, width=0.6)
axes[1, 0].set_title("Perbandingan MAPE", fontsize=15, pad=12)
axes[1, 0].set_ylabel("Persentase (%)", fontsize=12)
for bar in bars:
    yval = bar.get_height()
    axes[1, 0].text(
        bar.get_x() + bar.get_width() / 2,
        yval + (max(mape_values) * 0.02),
        round(yval, 2),
        ha="center",
        va="bottom",
        fontsize=10,
    )

# =========================
# SMAPE
# =========================
bars = axes[1, 1].bar(models, smape_values, width=0.6)
axes[1, 1].set_title("Perbandingan SMAPE", fontsize=15, pad=12)
axes[1, 1].set_ylabel("Persentase (%)", fontsize=12)
for bar in bars:
    yval = bar.get_height()
    axes[1, 1].text(
        bar.get_x() + bar.get_width() / 2,
        yval + (max(smape_values) * 0.02),
        round(yval, 2),
        ha="center",
        va="bottom",
        fontsize=10,
    )

# =========================
# RAPATKAN LAYOUT
# =========================
plt.tight_layout(rect=[0.02, 0.02, 1, 0.95])

# =========================
# SIMPAN GAMBAR
# =========================
# plt.savefig(
#     f"{output_folder}/evaluasi_komparasi_model.png",
#     dpi=300,
#     bbox_inches="tight"
# )
# plt.close()
# print("\nVisualisasi berhasil dibuat")
# print(f"Lokasi file: {output_folder}/evaluasi_komparasi_model.png")

# =========================
# OUTPUT
# =========================
print("\n=== PERBANDINGAN HASIL PREDIKSI ===\n")
print(hasil_tabel.to_string(index=False))
