import argparse
import json
import math
import os
import joblib
import numpy as np
import warnings

# 경고 무시 (Feature name 관련 경고 제거)
warnings.filterwarnings("ignore", category=UserWarning)

# =========================================================
# 파일명
# =========================================================

FEATURES_FILENAME = "features.json"
MODEL_FILENAME = "model_lgbm.pkl"
SCALER_FILENAME = "scaler_lgbm.pkl"
FEATURE_NAMES_FILENAME = "features_lgbm.txt"

def safe_float(value):
    if value is None: return 0.0
    x = float(value)
    return x if math.isfinite(x) else 0.0

def predict(working_dir: str, songno: str, feature: str = None):
    features_path = os.path.join(working_dir, FEATURES_FILENAME) if feature is None else feature
    model_path = os.path.join(working_dir, MODEL_FILENAME)
    scaler_path = os.path.join(working_dir, SCALER_FILENAME)
    feature_names_path = os.path.join(working_dir, FEATURE_NAMES_FILENAME)

    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Model not found at {model_path}")

    model = joblib.load(model_path)
    scaler = joblib.load(scaler_path)

    with open(feature_names_path, "r", encoding="utf-8") as f:
        feature_names = [line.strip() for line in f.readlines() if line.strip()]

    with open(features_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    targets = [item for item in data if str(item["songno"]) == str(songno)]
    if len(targets) == 0:
        raise ValueError(f"Chart not found: songno={songno}")

    results = []
    for target in targets:
        row = [safe_float(target.get(k, 0)) for k in feature_names]
        X = np.array([row], dtype=np.float32)
        X = scaler.transform(X)
        pred = model.predict(X)[0]

        results.append({
            "songno": str(songno),
            "diff": target.get("difficulty", "unknown"),
            "predicted": round(float(pred), 4)
        })

    print(json.dumps(results, indent=2, ensure_ascii=False))

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--workingDir", required=True)
    parser.add_argument("--feature", required=False)
    parser.add_argument("--songno", required=True)
    args = parser.parse_args()

    predict(args.workingDir, args.songno, args.feature)