import sys
import json
import pandas as pd
import numpy as np
from sqlalchemy import create_engine
import joblib
import os
import time

t0 = time.time()
user = "root"
# password = ""
password = ".IcU&Ic4U,"
host = "localhost"
db = "udd_pmi_module"
engine = create_engine(
    f"mysql+pymysql://{user}:{password}@{host}/{db}",
    pool_pre_ping=True,
    pool_size=5,
    max_overflow=10,
)
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_PATH = os.path.join(BASE_DIR, "model.pkl")
cached_model = None
cached_usage_bins = None
start = time.time()
if cached_model is None:
    saved = joblib.load(MODEL_PATH)
    cached_model = saved["model"]
    cached_usage_bins = saved["usage_bins"]
model = cached_model
usage_bins = cached_usage_bins
t1 = time.time()
if len(sys.argv) < 4:
    print(json.dumps({"error": "input kurang"}))
    sys.exit()
kode_list = sys.argv[1].split(",")
bulan_target = int(sys.argv[2])
tahun_target = int(sys.argv[3])
output = []
for kode in kode_list:
    kode = kode.strip()
    meta_query = f"""
    SELECT
        fluktuatif,
        nama_barang,
        satuan
    FROM items
    WHERE kode='{kode}'
    LIMIT 1
    """
    meta_df = pd.read_sql(meta_query, engine)
    t2 = time.time()
    if meta_df.empty:
        output.append(
            {
                "kode": kode,
                "nama_barang": "",
                "satuan": "",
                "bulan": bulan_target,
                "tahun": tahun_target,
                "prediction": 0,
                "error": "barang tidak ditemukan",
            }
        )
        continue
    fluktuatif = int(meta_df.iloc[0]["fluktuatif"])
    nama_barang = str(meta_df.iloc[0]["nama_barang"])
    satuan = str(meta_df.iloc[0]["satuan"])
    query = f"""
    SELECT
        tahun,
        bulan,
        total_pemakaian
    FROM item_usage_monthly
    WHERE kode='{kode}'
    ORDER BY tahun,bulan
    """
    df = pd.read_sql(query, engine)
    t3 = time.time()
    df["total_pemakaian"] = pd.to_numeric(
        df["total_pemakaian"], errors="coerce"
    ).fillna(0)
    if len(df) < 6:
        output.append(
            {
                "kode": kode,
                "nama_barang": nama_barang,
                "satuan": satuan,
                "bulan": bulan_target,
                "tahun": tahun_target,
                "prediction": 0,
                "error": "data < 6 bulan",
            }
        )
        continue
    if fluktuatif == 0:
        pred_real = int(round(df["total_pemakaian"].tail(3).mean()))
        output.append(
            {
                "kode": kode,
                "nama_barang": nama_barang,
                "satuan": satuan,
                "bulan": bulan_target,
                "tahun": tahun_target,
                "prediction": pred_real,
                "method": "average_3_months",
            }
        )
        continue
    history = list(df["total_pemakaian"])
    # 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_target - 1) // 3) + 1
    # X = 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_target,
    #             "quarter": quarter,
    #             "is_awal_tahun": int(bulan_target in [1, 2, 3]),
    #             "is_tengah_tahun": int(bulan_target in [6, 7, 8]),
    #             "is_akhir_tahun": int(bulan_target in [10, 11, 12]),
    #         }
    #     ]
    # )
    # pred_log = model.predict(X)[0]
    # pred_real = np.expm1(pred_log)
    # if pd.isna(pred_real):
    #     pred_real = 0
    # pred_real = max(0, int(round(pred_real)))
    from datetime import datetime

last_year = int(df.iloc[-1]["tahun"])
last_month = int(df.iloc[-1]["bulan"])

current = datetime(last_year, last_month, 1)
target = datetime(tahun_target, bulan_target, 1)

while current < target:

    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

    if current.month == 12:
        next_month = 1
        next_year = current.year + 1
    else:
        next_month = current.month + 1
        next_year = current.year

    quarter = ((next_month - 1) // 3) + 1

    X = 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": next_month,
                "quarter": quarter,
                "is_awal_tahun": int(next_month in [1, 2, 3]),
                "is_tengah_tahun": int(next_month in [6, 7, 8]),
                "is_akhir_tahun": int(next_month in [10, 11, 12]),
            }
        ]
    )

    pred_log = model.predict(X)[0]
    pred = max(0, int(round(np.expm1(pred_log))))

    history.append(pred)

    current = datetime(next_year, next_month, 1)

pred_real = history[-1]
output.append(
    {
        "kode": kode,
        "nama_barang": nama_barang,
        "satuan": satuan,
        "bulan": bulan_target,
        "tahun": tahun_target,
        "prediction": pred_real,
        "method": "random_forest",
    }
)
# output.append(
#     {
#         "kode": kode,
#         "nama_barang": nama_barang,
#         "satuan": satuan,
#         "bulan": bulan_target,
#         "tahun": tahun_target,
#         "prediction": pred_real,
#         "method": "random_forest",
#     }
# )
end = time.time()
print(
    json.dumps(
        {
            "data": output,
            "debug": {
                "load_model": round(t1 - t0, 2),
                "meta_query": round(t2 - t1, 2),
                "history_query": round(t3 - t2, 2),
                "total": round(end - t0, 2),
            },
        },
        allow_nan=False,
    )
)
