import argparse
import json
import os
import joblib
import numpy as np
from pathlib import Path

# =========================================================
# Configuration (Must match training)
# =========================================================
MAX_NOTES = 2000
FACTOR_COUNT = 4
INPUT_DIM = MAX_NOTES * FACTOR_COUNT

MODEL_FILENAME = "model.pkl"
SCALER_FILENAME = "scaler.pkl"

def safe_float(value):
    try: return float(value)
    except: return 0.0

def predict():
    parser = argparse.ArgumentParser()
    parser.add_argument("--workingDir", required=True)
    parser.add_argument("--songno", required=True)
    parser.add_argument("--factor", required=False) # Input factor JSON file
    args = parser.parse_args()

    model_path = os.path.join(args.workingDir, MODEL_FILENAME)
    scaler_path = os.path.join(args.workingDir, SCALER_FILENAME)

    if not os.path.exists(model_path):
        print(f"Model not found: {model_path}")
        return

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

    with open(args.factor or (Path(args.workingDir) / 'factors.json'), "r", encoding="utf-8") as f:
        data = json.load(f)

    # Filter by songno
    targets = [item for item in data if str(item["songno"]) == str(args.songno)]
    if not targets:
        print(f"No data found for songno: {args.songno}")
        return

    results = []
    for item in targets:
        raw_factors = item["factors"]
        vector = np.zeros(INPUT_DIM, dtype=np.float32)
        
        for i in range(min(len(raw_factors), MAX_NOTES)):
            for j in range(FACTOR_COUNT):
                vector[i * FACTOR_COUNT + j] = safe_float(raw_factors[i][j])
        
        # Scale and Predict
        X = scaler.transform([vector])
        pred = model.predict(X)[0]
        
        results.append({
            "songno": item["songno"],
            "diff": item["difficulty"],
            "predicted": float(pred)
        })

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

if __name__ == "__main__":
    predict()