from sklearn.ensemble import GradientBoostingRegressor
import pandas as pd
import numpy as np
from sklearn.metrics import mean_absolute_error
import joblib
import argparse
import os
import sys

def train():
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset', type=str, default='datas/dataset.csv')
    parser.add_argument('--model_path', type=str, default='model/constant_predictor.joblib')
    parser.add_argument('--batch_size', type=int, default=200)
    parser.add_argument('--iterations', type=int, default=10) # 200개씩 몇 번 반복할지
    args = parser.parse_args()

    if not os.path.exists(args.dataset):
        sys.exit(1)

    df = pd.read_csv(args.dataset)
    
    # 모델 로드 또는 생성
    if os.path.exists(args.model_path):
        print(f"Loading existing model from {args.model_path} for update...")
        model = joblib.load(args.model_path)
        # 기존 모델의 트리 개수를 늘려가며 학습하기 위해 n_estimators 증가
        model.n_estimators += 50 
        model.warm_start = True
    else:
        print("Creating new model...")
        model = GradientBoostingRegressor(
            n_estimators=100,
            learning_rate=0.05,
            max_depth=8,
            warm_start=True,
            random_state=42
        )

    for i in range(args.iterations):
        # 랜덤하게 200개 샘플링
        batch = df.sample(n=min(args.batch_size, len(df)))
        X_batch = batch.drop('target', axis=1)
        y_batch = batch['target']

        model.fit(X_batch, y_batch)
        
        # 전체 데이터에 대한 성능 확인 (학습 경과 관찰용)
        preds = model.predict(df.drop('target', axis=1))
        mae = mean_absolute_error(df['target'], preds)
        
        print(f"Iteration {i+1}/{args.iterations} - Current Model Estimators: {model.n_estimators}, Dataset MAE: {mae:.4f}")
        
        # 매 반복마다 트리 조금씩 추가
        model.n_estimators += 20

    joblib.dump(model, args.model_path)
    print(f"Model updated and saved to {args.model_path}")

if __name__ == "__main__":
    train()