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
import lightgbm as lgb

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

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

MAX_NOTES = 2000
FACTOR_COUNT = 4
INPUT_DIM = MAX_NOTES * FACTOR_COUNT

TRAIN_SIZE = 0
VALID_SIZE = 0
RANDOM_STATE = 42

PARAMS = {
    'objective': 'regression',
    'metric': 'mae',
    'verbosity': -1,
    'boosting_type': 'gbdt',
    'random_state': RANDOM_STATE,
    'learning_rate': 0.02,
    'num_leaves': 63,
    'n_estimators': 2000
}

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 = {(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]
                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, eval_set=[(X_valid, y_valid)], init_model=model, callbacks=[lgb.early_stopping(stopping_rounds=100)])
    else:
        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"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='teal')
    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)