import pandas as pd
import numpy as np
from sqlalchemy import create_engine, text
import joblib
import os
from datetime import datetime

# =========================
# DB
# =========================
user = "root"
password = ""
host = "localhost"
db = "udd_pmi_module"
engine = create_engine(
    f"mysql+pymysql://{user}:{password}@{host}/{db}"
)
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_PATH = os.path.join(BASE_DIR, "model.pkl")
saved = joblib.load(MODEL_PATH)
model = saved["model"]
usage_bins = saved["usage_bins"]

# =========================
# AMBIL SEMUA BARANG
# =========================
items = pd.read_sql(
    """
    SELECT DISTINCT kode
    FROM item_usage_monthly
    ORDER BY kode
    """,
    engine,
)

# =========================
# HAPUS DATA PREDIKSI LAMA
# =========================
with engine.begin() as conn:
    conn.execute(text("""
        DELETE FROM prediction_results
    """))

# =========================
# FEATURE BUILDER
# =========================
def build_features(history, bulan):
    lag = history[-6:]
    ma_6 = np.mean(lag)
    trend_3 = lag[-1] - np.mean(
        [lag[-2], lag[-3], lag[-4]]
    )
    momentum_1 = lag[-1] - lag[-2]
    momentum_2 = lag[-2] - lag[-3]
    rolling_std_6 = np.std(
        lag,
        ddof=1
    )
    max_6 = np.max(lag)
    min_6 = np.min(lag)
    range_6 = max_6 - min_6
    cv_6 = rolling_std_6 / (ma_6 + 1)
    growth_rate = (
        (lag[-1] - lag[-2])
        /
        (lag[-2] + 1)
    )
    usage_level = pd.cut(
        [ma_6],
        bins=usage_bins,
        labels=False,
        include_lowest=True
    )[0]
    if pd.isna(usage_level):
        usage_level = 0
    quarter = ((bulan-1)//3)+1
    return pd.DataFrame([{
        "lag_1":lag[-1],
        "lag_2":lag[-2],
        "lag_3":lag[-3],
        "lag_4":lag[-4],
        "lag_5":lag[-5],
        "lag_6":lag[-6],
        "ma_6":ma_6,
        "trend_3":trend_3,
        "momentum_1":momentum_1,
        "momentum_2":momentum_2,
        "rolling_std_6":rolling_std_6,
        "max_6":max_6,
        "min_6":min_6,
        "range_6":range_6,
        "cv_6":cv_6,
        "growth_rate":growth_rate,
        "usage_level":usage_level,
        "bulan":bulan,
        "quarter":quarter,
        "is_awal_tahun":int(
            bulan in [1,2,3]
        ),
        "is_tengah_tahun":int(
            bulan in [6,7,8]
        ),
        "is_akhir_tahun":int(
            bulan in [10,11,12]
        )
    }])

# =========================
# MAIN
# =========================
results=[]
for kode in items["kode"]:
    hist=pd.read_sql(
        f"""
        SELECT
            tahun,
            bulan,
            total_pemakaian
        FROM item_usage_monthly
        WHERE kode='{kode}'
        ORDER BY tahun,bulan
        """,
        engine
    )
    if hist.empty:
        continue
    hist["total_pemakaian"]=pd.to_numeric(
        hist["total_pemakaian"],
        errors="coerce"
    ).fillna(0)
    if len(hist)<18:
        print(
            f"SKIP {kode}: data < 18"
        )
        continue
    target_hist=hist.tail(12)
    for idx in target_hist.index:
        start_idx=idx-6
        history=hist.iloc[
            start_idx:idx
        ]["total_pemakaian"].tolist()
        if len(history)<6:
            continue
        bulan_target=int(
            hist.loc[idx,"bulan"]
        )
        tahun_target=int(
            hist.loc[idx,"tahun"]
        )
        X=build_features(
            history,
            bulan_target
        )
        pred_log=model.predict(
            X
        )[0]
        pred=np.expm1(
            pred_log
        )
        pred=np.clip(
            round(pred),
            0,
            None
        )
        results.append({
            "kode":kode,
            "bulan_prediksi":bulan_target,
            "tahun_prediksi":tahun_target,
            "jumlah_prediksi":int(pred),
            "tanggal_prediksi":datetime.now().date()
        })

# =========================
# INSERT
# =========================
if results:
    df_insert=pd.DataFrame(results)
    df_insert.to_sql(
        "prediction_results",
        engine,
        if_exists="append",
        index=False
    )
print("SELESAI")