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

from xgboost import XGBRegressor
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_absolute_error

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

TEST_SIZE = 0.2
RANDOM_STATE = 42

N_ESTIMATORS = 500
MAX_DEPTH = 6
LEARNING_RATE = 0.05
SUBSAMPLE = 0.8
COLSAMPLE_BYTREE = 0.8

CONTINUE_TRAINING = True

# 예측 성공으로 간주할 허용 오차
ERROR_TOLERANCE = 0.2

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

FEATURES_FILENAME = "features.json"
MEASURE_FILENAME = "measure.csv"

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

# =========================================================
# 무시할 key
# =========================================================

IGNORE_KEYS = {
    "songno",
    "difficulty"
}


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 train_model(
    working_dir: str,
    data_dir: str
):
    # =====================================================
    # path
    # =====================================================

    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
    )

    # =====================================================
    # features.json
    # =====================================================

    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
    # =====================================================

    feature_map = {}

    for item in feature_data:
        key = (
            str(item["songno"]),
            str(item["difficulty"])
        )

        feature_map[key] = item

    # =====================================================
    # feature names
    # =====================================================

    feature_names = sorted([
        k for k in feature_data[0].keys()
        if k not in IGNORE_KEYS
    ])

    # =====================================================
    # measure.csv
    # =====================================================

    X = []
    y = []

    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 = safe_float(row[0])
            songno = str(row[1])
            diff = str(row[2])

            key = (songno, diff)

            if key not in feature_map:
                print(
                    f"[WARN] feature not found: "
                    f"{songno} {diff}"
                )
                continue

            feature_item = feature_map[key]

            features = [
                safe_float(feature_item.get(k, 0))
                for k in feature_names
            ]

            X.append(features)
            y.append(measure)

    if len(X) == 0:
        raise ValueError("No training data")

    X = np.array(X, dtype=np.float32)
    y = np.array(y, dtype=np.float32)

    print(f"Dataset Size: {len(X)}")
    print(f"Feature Count: {len(feature_names)}")

    # =====================================================
    # split
    # =====================================================

    X_train, X_valid, y_train, y_valid = train_test_split(
        X,
        y,
        test_size=TEST_SIZE,
        random_state=RANDOM_STATE
    )

    # =====================================================
    # scaler
    # =====================================================

    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)

    # =====================================================
    # model
    # =====================================================

    if CONTINUE_TRAINING and os.path.exists(model_path):
        print("Loading existing model...")

        model = joblib.load(model_path)

        previous_booster = model.get_booster()

        model.fit(
            X_train,
            y_train,
            xgb_model=previous_booster
        )

    else:
        print("Creating new model...")

        model = XGBRegressor(
            n_estimators=N_ESTIMATORS,
            max_depth=MAX_DEPTH,
            learning_rate=LEARNING_RATE,
            subsample=SUBSAMPLE,
            colsample_bytree=COLSAMPLE_BYTREE,
            objective="reg:squarederror",
            random_state=RANDOM_STATE
        )

        model.fit(X_train, y_train)

    # =====================================================
    # evaluate
    # =====================================================

    pred = model.predict(X_valid)

    mae = mean_absolute_error(y_valid, pred)

    correct = np.sum(
        np.abs(pred - y_valid) <= ERROR_TOLERANCE
    )

    accuracy = correct / len(y_valid)

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

    # =====================================================
    # feature importance
    # =====================================================

    print("\nFeature Importance:")

    importance = model.feature_importances_

    pairs = list(zip(feature_names, importance))
    pairs.sort(key=lambda x: x[1], reverse=True)

    for name, score in pairs:
        print(f"{name:25} {score:.6f}")

    # =====================================================
    # save
    # =====================================================

    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("\nSaved:")
    print(model_path)
    print(scaler_path)
    print(feature_names_path)


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

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--workingDir",
        required=True
    )

    parser.add_argument(
        "--dataDir",
        required=True
    )

    args = parser.parse_args()

    train_model(
        args.workingDir,
        args.dataDir
    )