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

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

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

TRAIN_SIZE = 0
VALID_SIZE = 0

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.5

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

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"
}


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


# =========================================================
# train
# =========================================================

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

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

    # =====================================================
    # dataset build
    # =====================================================

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

            dataset.append((
                features,
                measure
            ))

    # =====================================================
    # shuffle
    # =====================================================

    random.shuffle(dataset)

    required_size = TRAIN_SIZE + VALID_SIZE

    if len(dataset) < required_size:
        raise ValueError(
            f"Not enough dataset "
            f"({len(dataset)} < {required_size})"
        )

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

    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)}")
    print(f"Valid Size: {len(X_valid)}")
    print(f"Feature Count: {len(feature_names)}")

    # =====================================================
    # 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
    )
    
    parser.add_argument(
        "--trainSize",
        required=True
    )
    
    parser.add_argument(
        "--validSize",
        required=True
    )

    args = parser.parse_args()
    
    TRAIN_SIZE = int(args.trainSize)
    VALID_SIZE = int(args.validSize)

    train_model(
        args.workingDir,
        args.dataDir
    )