import pandas as pd
import json, os, sys, joblib
import numpy as np
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.model_selection import train_test_split

def train(script_dir, working_dir, data_dir, train_count, val_count, margin):
    features = pd.read_json(os.path.join(working_dir, 'features.json'))
    measures = pd.read_csv(os.path.join(data_dir, 'measure.csv'))
    
    features['songno'] = features['songno'].astype(int)
    measures['songno'] = measures['songno'].astype(int)
    measures['diff'] = measures['diff'].replace('ura', 'edit')
    
    df = pd.merge(features, measures, on=['songno', 'diff'])
    
    with open(os.path.join(script_dir, 'factor.json'), 'r') as f:
        weights = json.load(f)
    for col in ['physical_density', 'stamina_requirement', 'pattern_complexity', 'rhythmic_complexity', 'reading_gimmick']:
        df[col] = df[col] * weights.get(col, 1.0)

    X_cols = ['physical_density', 'stamina_requirement', 'pattern_complexity', 'rhythmic_complexity', 'reading_gimmick']
    model_path = os.path.join(working_dir, 'model.pkl')
    model = joblib.load(model_path) if os.path.exists(model_path) else GradientBoostingRegressor(n_estimators=200, learning_rate=0.05, max_depth=3)

    iteration = 1
    while True:
        # 3. 데이터 샘플링
        df_sample = df.sample(n=int(train_count) + int(val_count))
        train_df, val_df = train_test_split(df_sample, test_size=int(val_count))

        X_train, y_train = train_df[X_cols], train_df['상수']
        X_val, y_val = val_df[X_cols], val_df['상수']

        # 4. 학습 (최대 10회)
        for attempt in range(1, 11):
            model.fit(X_train, y_train)
            train_err = np.max(np.abs(np.clip(model.predict(X_train), 1.0, 12.0) - y_train))
            print(f"Iteration {iteration} - Attempt {attempt} - Train Error: {train_err:.4f}")
            if train_err <= float(margin): break
            model.set_params(n_estimators=model.n_estimators + 50)

        # 5. 검증
        pred_val = np.clip(model.predict(X_val), 1.0, 12.0)
        val_errors = np.abs(pred_val - y_val)
        
        # 6. 검증 실패 시 재시도
        if np.any(val_errors > float(margin)):
            print(f"Validation failed (max error: {np.max(val_errors):.4f}). Retrying...")
            iteration += 1
            continue
            
        val_result = pd.DataFrame({
            'songno': val_df['songno'],
            'difficulty': val_df['diff'],
            'measure': y_val,
            'predicted_measure': pred_val,
            'error': val_errors
        })
        val_result.to_csv(os.path.join(working_dir, f'validate_result_{iteration}.csv'), index=False)
        val_result.to_csv(os.path.join(working_dir, 'validate_result.csv'), index=False)
        break

    joblib.dump(model, model_path)

if __name__ == "__main__":
    train(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4], sys.argv[5], sys.argv[6])