import argparse
import csv
import json
import math
import os
import random
import joblib
import numpy as np
import lightgbm as lgb
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
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.pkl"
SCALER_FILENAME = "scaler.pkl"
FEATURE_NAMES_FILENAME = "features.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, songno, diff))

    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)
    valid_info = [(s, d) for _, _, s, d in valid_dataset]

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

    # =====================================================
    # 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_lightgbm.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')
        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}")

    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)