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


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

FEATURES_FILENAME = "features.json"

MODEL_FILENAME = "model.pkl"
SCALER_FILENAME = "scaler.pkl"
FEATURE_NAMES_FILENAME = "features.txt"


# =========================================================
# safe float
# =========================================================

def safe_float(value):
    if value is None:
        return 0.0

    x = float(value)

    if not math.isfinite(x):
        return 0.0

    return x


# =========================================================
# 예측 함수
# =========================================================

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
    )

    # =====================================================
    # 모델 로드
    # =====================================================

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

    # =====================================================
    # feature 이름 로드
    # =====================================================

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

    # =====================================================
    # features.json 로드
    # =====================================================

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

    # =====================================================
    # target 찾기
    # =====================================================

    targets = []

    for item in data:
        if str(item["songno"]) == str(songno):
            targets.append(item)

    if len(targets) == 0:
        raise ValueError(f"Chart not found: songno={songno}")

    # =====================================================
    # feature vector 생성
    # =====================================================

    results = []

    for target in targets:
        row = []

        for k in feature_names:
            value = target.get(k, 0)
            row.append(safe_float(value))

        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))


# =========================================================
# main
# =========================================================

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
    )