import argparse
import csv
import json
import math
import os
import random
import joblib
import numpy as np
import lightgbm as lgb
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_absolute_error

# =========================================================
# Hyper Parameters
# =========================================================

TRAIN_SIZE = 0
VALID_SIZE = 0

RANDOM_STATE = 42

# LightGBM 특정 하이퍼파라미터
PARAMS = {
    'objective': 'regression',
    'metric': 'mae',
    'verbosity': -1,
    'boosting_type': 'gbdt',
    'random_state': RANDOM_STATE,
    'learning_rate': 0.02, # 더 정밀한 학습을 위해 하향
    'num_leaves': 63,      # 더 복잡한 패턴 학습을 위해 상향
    'feature_fraction': 0.9,
    'bagging_fraction': 0.8,
    'bagging_freq': 5,
    'n_estimators': 3000   # 학습량 대폭 상향
}

CONTINUE_TRAINING = True
ERROR_TOLERANCE = 0.1

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

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

IGNORE_KEYS = {"songno", "difficulty"}

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

def train_model(working_dir: str, data_dir: str):
    random.seed(RANDOM_STATE)

    features_path = os.path.join(working_dir, FEATURES_FILENAME)
    measure_path = os.path.join(data_dir, MEASURE_FILENAME)
    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)

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

    if len(feature_data) == 0:
        raise ValueError("features.json is empty")

    feature_map = {(str(item["songno"]), str(item["difficulty"])): item for item in feature_data}
    feature_names = sorted([k for k in feature_data[0].keys() if k not in IGNORE_KEYS])

    dataset = []
    with open(measure_path, "r", encoding="utf-8") as f:
        reader = csv.reader(f)
        next(reader, None)
        for row in reader:
            if len(row) < 3: continue
            measure, songno, diff = safe_float(row[0]), str(row[1]), str(row[2])
            key = (songno, diff)
            if key in feature_map:
                features = [safe_float(feature_map[key].get(k, 0)) for k in feature_names]
                dataset.append((features, measure))

    random.shuffle(dataset)
    if len(dataset) < (TRAIN_SIZE + VALID_SIZE):
        raise ValueError(f"Not enough dataset ({len(dataset)} < {TRAIN_SIZE + VALID_SIZE})")

    train_dataset = dataset[:TRAIN_SIZE]
    valid_dataset = dataset[TRAIN_SIZE:TRAIN_SIZE + VALID_SIZE]

    X_train = np.array([x for x, _ in train_dataset], dtype=np.float32)
    y_train = np.array([y for _, y in train_dataset], dtype=np.float32)
    X_valid = np.array([x for x, _ in valid_dataset], dtype=np.float32)
    y_valid = np.array([y for _, y in valid_dataset], dtype=np.float32)

    print(f"Train Size: {len(X_train)} | Valid Size: {len(X_valid)} | Features: {len(feature_names)}")

    if CONTINUE_TRAINING and os.path.exists(scaler_path):
        print("Loading existing scaler...")
        scaler = joblib.load(scaler_path)
    else:
        print("Creating new scaler...")
        scaler = StandardScaler()
        scaler.fit(X_train)

    X_train = scaler.transform(X_train)
    X_valid = scaler.transform(X_valid)

    if CONTINUE_TRAINING and os.path.exists(model_path):
        print("Loading existing model for incremental training...")
        model = joblib.load(model_path)
        model.fit(
            X_train, y_train,
            eval_set=[(X_valid, y_valid)],
            init_model=model,
            callbacks=[lgb.early_stopping(stopping_rounds=100)]
        )
    else:
        print("Creating new LightGBM model...")
        model = lgb.LGBMRegressor(**PARAMS)
        model.fit(
            X_train, y_train,
            eval_set=[(X_valid, y_valid)],
            callbacks=[lgb.early_stopping(stopping_rounds=100)]
        )

    pred = model.predict(X_valid)
    mae = mean_absolute_error(y_valid, pred)
    accuracy = np.sum(np.abs(pred - y_valid) <= ERROR_TOLERANCE) / len(y_valid)

    print(f"\nMAE: {mae:.4f} | Accuracy (±{ERROR_TOLERANCE}): {accuracy:.4f}")

    joblib.dump(model, model_path)
    joblib.dump(scaler, scaler_path)
    with open(feature_names_path, "w", encoding="utf-8") as f:
        for name in feature_names: f.write(name + "\n")

    print(f"\nSaved to {working_dir}: {MODEL_FILENAME}, {SCALER_FILENAME}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--workingDir", required=True)
    parser.add_argument("--dataDir", required=True)
    parser.add_argument("--trainSize", required=True, type=int)
    parser.add_argument("--validSize", required=True, type=int)
    args = parser.parse_args()

    TRAIN_SIZE, VALID_SIZE = args.trainSize, args.validSize
    train_model(args.workingDir, args.dataDir)