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

MAX_NOTES = 2000 # 분석할 최대 노트 수
FACTOR_COUNT = 4 # [type, bpm, scroll, delta]
INPUT_DIM = MAX_NOTES * FACTOR_COUNT

TRAIN_SIZE = 0
VALID_SIZE = 0
RANDOM_STATE = 42

N_ESTIMATORS = 500
MAX_DEPTH = 6
LEARNING_RATE = 0.05
CONTINUE_TRAINING = True
ERROR_TOLERANCE = 0.1

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

FACTORS_FILENAME = "factors.json"
MEASURE_FILENAME = "measure.csv"
MODEL_FILENAME = "model.pkl"
SCALER_FILENAME = "scaler.pkl"

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)

    factors_path = os.path.join(working_dir, FACTORS_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)

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

    # feature_map build: key -> list of factors
    feature_map = {(str(item["songno"]), str(item["difficulty"])): item["factors"] for item in factor_data}

    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:
                raw_factors = feature_map[key]
                # 고정 길이 벡터로 변환 (Padding or Truncating)
                vector = np.zeros(INPUT_DIM, dtype=np.float32)
                for i in range(min(len(raw_factors), MAX_NOTES)):
                    for j in range(FACTOR_COUNT):
                        vector[i * FACTOR_COUNT + j] = safe_float(raw_factors[i][j])
                
                dataset.append((vector, measure, songno, diff))

    random.shuffle(dataset)
    if len(dataset) < (TRAIN_SIZE + VALID_SIZE):
        raise ValueError(f"Dataset size {len(dataset)} < required {TRAIN_SIZE + VALID_SIZE}")

    train_ds = dataset[:TRAIN_SIZE]
    valid_ds = dataset[TRAIN_SIZE:TRAIN_SIZE + VALID_SIZE]

    X_train = np.array([x for x, _, _, _ in train_ds])
    y_train = np.array([y for _, y, _, _ in train_ds])
    X_valid = np.array([x for x, _, _, _ in valid_ds])
    y_valid = np.array([y for _, y, _, _ in valid_ds])
    valid_info = [(s, d) for _, _, s, d in valid_ds]

    if CONTINUE_TRAINING and os.path.exists(scaler_path):
        scaler = joblib.load(scaler_path)
    else:
        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):
        model = joblib.load(model_path)
        model.fit(X_train, y_train, xgb_model=model.get_booster())
    else:
        model = XGBRegressor(n_estimators=N_ESTIMATORS, max_depth=MAX_DEPTH, learning_rate=LEARNING_RATE, objective="reg:squarederror", random_state=RANDOM_STATE)
        model.fit(X_train, y_train)

    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"MAE: {mae:.4f} | Accuracy: {accuracy:.4f}")

    # Results save
    validate_details = []
    for i in range(len(y_valid)):
        validate_details.append({"songno": valid_info[i][0], "diff": valid_info[i][1], "actual": float(y_valid[i]), "predicted": float(pred[i]), "error": float(y_valid[i] - pred[i])})
    
    validate_details.sort(key=lambda x: abs(x["error"]), reverse=True)
    with open(os.path.join(working_dir, "validate.json"), "w", encoding="utf-8") as f:
        json.dump({"summary": {"mae": float(mae), "accuracy": float(accuracy)}, "details": validate_details}, f, indent=2)

    # Plot
    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', color='crimson')
    plt.axhline(0.2, color='green', linestyle='--')
    plt.ylim(0, 4)
    plt.savefig(os.path.join(working_dir, "validate.png"))
    
    joblib.dump(model, model_path)
    joblib.dump(scaler, scaler_path)

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)