import argparse
import csv
import json
import math
import os
import random
import joblib
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

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

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

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,
                songno,
                diff
            ))

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

    valid_info = [
        (s, d) for _, _, s, d in valid_dataset
    ]

    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 validate.json
    # =====================================================

    validate_details = []
    for i in range(len(y_valid)):
        actual = float(y_valid[i])
        predicted = float(pred[i])
        songno, diff = valid_info[i]
        
        validate_details.append({
            "songno": songno,
            "diff": diff,
            "actual": actual,
            "predicted": predicted,
            "error": actual - predicted
        })

    # 에러 절댓값 기준 정렬
    validate_details.sort(key=lambda x: abs(x["error"]), reverse=True)

    validate_result = {
        "summary": {
            "total_compared": len(y_valid),
            "average_absolute_error": float(mae),
            "accuracy": float(accuracy),
            "timestamp": "now",
            "script_used": "train/train_xgboost.py"
        },
        "details": validate_details
    }

    validate_path = os.path.join(working_dir, "validate.json")
    with open(validate_path, "w", encoding="utf-8") as f:
        json.dump(validate_result, f, indent=2, ensure_ascii=False)

    print(f"Validation result saved: {validate_path}")

    # =====================================================
    # save validate.png
    # =====================================================
    try:
        plt.switch_backend('Agg') # GUI 없는 환경 대응
        df_plot = pd.DataFrame(validate_details)
        df_plot['abs_error'] = df_plot['error'].abs()
        df_plot = df_plot.sort_values('abs_error', ascending=False).reset_index(drop=True)

        plt.figure(figsize=(12, 6))
        sns.scatterplot(data=df_plot, x=df_plot.index, y='abs_error', alpha=0.6, s=20, color='darkorange')
        
        plt.axhline(0.2, color='green', linestyle='--', linewidth=0.8, alpha=0.5, label='Target (0.2)')
        plt.axhline(0.5, color='blue', linestyle='--', linewidth=0.8, alpha=0.5)
        plt.axhline(1.0, color='red', linestyle='--', linewidth=0.8, alpha=0.5)
        
        plt.ylim(0, max(3.5, df_plot['abs_error'].max() + 0.5))
        plt.title(f'Validation Absolute Error - {os.path.basename(working_dir)}', fontsize=14)
        plt.xlabel('Samples (Sorted by Error Magnitude)', fontsize=12)
        plt.ylabel('Absolute Error', fontsize=12)
        plt.grid(True, axis='y', alpha=0.3)
        
        plot_path = os.path.join(working_dir, "validate.png")
        plt.savefig(plot_path)
        plt.close()
        print(f"Validation plot saved: {plot_path}")
    except Exception as e:
        print(f"[WARN] Failed to create validation plot: {e}")

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