# 이 코드는 다음과 같은 최적화를 포함하고 있습니다:
# 메모리 사용량 최소화: 필요한 열만 유지하고 나머지는 처리 후 즉시 삭제합니다.
# 디스크 공간 관리: 최대 3개의 배치 파일만 유지하며, 공간이 부족할 경우 이전 파일을 삭제합니다.
# 점진적 학습: 각 배치마다 모델을 학습하고 성능을 평가하여 최상의 모델을 유지합니다.
# 효율적인 데이터 처리: 멀티프로세싱을 사용하여 데이터 처리 속도를 향상시킵니다.
# 메모리 및 디스크 사용량 모니터링: 임계값을 초과할 경우 적절한 조치를 취합니다.
# 이 코드는 10GB RAM과 30GB 디스크 공간 제한 내에서 작동하도록 설계되었습니다. 그러나 실제 데이터 크기에 따라 CHUNK_SIZE, BATCH_SIZE, MAX_BATCHES 등의 값을 조정해야 할 수 있습니다.
# 또한, 모델 학습이 완료된 후 predict_price 함수를 사용하여 새로운 입력에 대한 가격을 예측할 수 있습니다. 이 함수는 학습된 모델, 스케일러, 특성 이름, 그리고 CPU와 GPU 이름 목록을 사용하여 가격을 예측합니다.

import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import xgboost as xgb
from sklearn.metrics import mean_squared_error, r2_score
import joblib
from db_config import get_db_connection
import gc
import time
from datetime import datetime
import psutil
import os
import multiprocessing
from fuzzywuzzy import process
import shutil

# 전역 변수 설정
N_CORES = multiprocessing.cpu_count()
MEMORY_THRESHOLD = 80  # 메모리 사용량 임계값 (%)
DISK_THRESHOLD = 5  # 디스크 여유 공간 임계값 (GB)
CHUNK_SIZE = 100000
BATCH_SIZE = 1000000  # 한 번에 처리할 행 수
MAX_BATCHES = 3  # 한 번에 유지할 최대 배치 파일 수

def normalize_cpu_gpu(model):
    model_lower = model.lower()
    if 'intel' in model_lower:
        return 'Intel'
    elif 'amd' in model_lower:
        return 'AMD'
    elif 'nvidia' in model_lower or 'geforce' in model_lower:
        return 'NVIDIA'
    else:
        return 'Other'

def process_chunk(chunk):
    chunk['cpu_brand'] = chunk['cpu_name'].apply(normalize_cpu_gpu)
    chunk['gpu_brand'] = chunk['gpu_name'].apply(normalize_cpu_gpu)
    chunk['ram_combined'] = chunk['ram_type'] + '_' + chunk['ram_size'].astype(str)
    chunk['hdd_combined'] = chunk['storage_type'] + '_' + chunk['hdd_size'].astype(str)

    for col in ['cpu_price', 'gpu_price', 'ram_price', 'hdd_price']:
        chunk[col] = pd.to_numeric(chunk[col], errors='coerce')

    chunk['ram_size'] = chunk['ram_size'].apply(lambda x: float(x.replace('GB', '').replace('gb', '')) if isinstance(x, str) and 'GB' in x.upper() else (float(str(x).replace('MB', '').replace('mb', '')) / 1000 if isinstance(x, str) and 'MB' in x.upper() else float(x)))
    chunk['hdd_size'] = chunk['hdd_size'].apply(lambda x: float(x.replace('TB', '').replace('tb', '')) * 1000 if isinstance(x, str) and 'TB' in x.upper() else (
        float(x.replace('GB', '').replace('gb', '')) if isinstance(x, str) and 'GB' in x.upper() else float(x)
    ))

    chunk['total_price'] = chunk['cpu_price'] + chunk['gpu_price'] + chunk['ram_price'] + chunk['hdd_price']

    return chunk[['cpu_brand', 'gpu_brand', 'ram_combined', 'hdd_combined', 'ram_size', 'hdd_size', 'total_price']]

def parallel_process_chunks(chunks):
    with multiprocessing.Pool(N_CORES) as pool:
        processed_chunks = pool.map(process_chunk, chunks)
    return processed_chunks

def find_best_match(input_name, name_list):
    best_match, score = process.extractOne(input_name, name_list)
    return best_match if score > 70 else None

def predict_price(model, scaler, feature_names, cpu_name, gpu_name, ram_type, ram_size, hdd_type, hdd_size, cpu_names, gpu_names):
    cpu_match = find_best_match(cpu_name, cpu_names)
    gpu_match = find_best_match(gpu_name, gpu_names)

    if cpu_match is None or gpu_match is None:
        return None

    input_data = pd.DataFrame({
        'cpu_brand': [normalize_cpu_gpu(cpu_match)],
        'gpu_brand': [normalize_cpu_gpu(gpu_match)],
        'ram_combined': [f"{ram_type}_{ram_size}"],
        'hdd_combined': [f"{hdd_type}_{hdd_size}"],
        'ram_size': [float(ram_size)],
        'hdd_size': [float(hdd_size) if 'TB' not in hdd_size.upper() else float(hdd_size.replace('TB', '').strip()) * 1000]
    })

    input_data_encoded = pd.get_dummies(input_data, columns=['cpu_brand', 'gpu_brand', 'ram_combined', 'hdd_combined'])

    for col in feature_names:
        if col not in input_data_encoded.columns:
            input_data_encoded[col] = 0

    input_data_encoded = input_data_encoded[feature_names]
    input_data_scaled = scaler.transform(input_data_encoded)

    predicted_price = model.predict(input_data_scaled)
    return predicted_price[0]

def save_chunk_to_csv(chunk, output_file, mode='a'):
    chunk.to_csv(output_file, mode=mode, header=(mode=='w'), index=False)

def check_memory_usage():
    memory_usage = psutil.virtual_memory().percent
    if memory_usage > MEMORY_THRESHOLD:
        print(f"메모리 사용량이 {memory_usage}%로 높습니다. 가비지 컬렉션을 실행합니다.")
        gc.collect()
        time.sleep(5)  # 5초 대기

def check_disk_space(path):
    total, used, free = shutil.disk_usage(path)
    free_gb = free // (2**30)
    return free_gb > DISK_THRESHOLD

def train_model(X, y):
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    model = xgb.XGBRegressor(random_state=42, n_jobs=N_CORES)
    model.fit(X_train_scaled, y_train)

    y_pred = model.predict(X_test_scaled)
    mse = mean_squared_error(y_test, y_pred)
    rmse = np.sqrt(mse)
    r2 = r2_score(y_test, y_pred)

    return model, scaler, rmse, r2

def main():
    start_time = time.time()
    start_datetime = datetime.now()
    print(f"프로그램 시작: {start_datetime.strftime('%Y-%m-%d %H:%M:%S.%f')}")

    db = get_db_connection()

    query = """
    SELECT
        cpu.cpu_name AS cpu_name, cpu.price AS cpu_price,
        gpu.videocard_name AS gpu_name, gpu.price AS gpu_price,
        ram.ramtype AS ram_type, ram.memory_size AS ram_size, ram.price AS ram_price,
        hdd.storage_type AS storage_type, hdd.size AS hdd_size, hdd.price AS hdd_price
    FROM api___ai_suggested_price_of_cpu cpu
    JOIN api___ai_suggested_price_of_gpu gpu
    JOIN api___ai_suggested_price_of_ram ram
    JOIN api___ai_suggested_price_of_hdd hdd
    """

    total_rows = 0
    batch_count = 0
    cpu_names = set()
    gpu_names = set()

    final_model = None
    final_scaler = None
    best_rmse = float('inf')

    while True:
        batch_start = total_rows
        batch_end = batch_start + BATCH_SIZE

        output_file = f'processed_data_batch_{batch_count % MAX_BATCHES}.csv'
        chunks = []
        for chunk in pd.read_sql_query(f"{query} LIMIT {BATCH_SIZE} OFFSET {batch_start}", db, chunksize=CHUNK_SIZE):
            chunks.append(chunk)
            if len(chunks) == N_CORES:
                processed_chunks = parallel_process_chunks(chunks)
                for processed_chunk in processed_chunks:
                    total_rows += len(processed_chunk)
                    cpu_names.update(chunk['cpu_name'])
                    gpu_names.update(chunk['gpu_name'])

                for processed_chunk in processed_chunks:
                    save_chunk_to_csv(processed_chunk, output_file, mode='a' if os.path.exists(output_file) else 'w')

                chunks = []
                check_memory_usage()

                if not check_disk_space('.'):
                    print("디스크 공간이 부족합니다. 이전 배치 파일을 삭제합니다.")
                    for i in range(max(0, batch_count - MAX_BATCHES), batch_count):
                        try:
                            os.remove(f'processed_data_batch_{i % MAX_BATCHES}.csv')
                        except FileNotFoundError:
                            pass  # 파일이 이미 삭제되었거나 존재하지 않는 경우 무시

        if not chunks:  # 더 이상 처리할 데이터가 없음
            break

        # 남은 청크 처리
        if chunks:
            processed_chunks = parallel_process_chunks(chunks)
            for processed_chunk in processed_chunks:
                save_chunk_to_csv(processed_chunk, output_file, mode='a' if os.path.exists(output_file) else 'w')
                total_rows += len(processed_chunk)
                cpu_names.update(chunk['cpu_name'])
                gpu_names.update(chunk['gpu_name'])
            total_rows_percent = round(total_rows / 19344333888 * 100, 5)
            print(f"Processed {total_rows:,} rows ({total_rows_percent}%)")

        print(f"배치 {batch_count} 처리 완료. 총 처리된 행 수: {total_rows:,}")

        # 현재 배치에 대해 모델 학습
        df = pd.read_csv(output_file)
        X = df.drop('total_price', axis=1)
        y = df['total_price']

        X = pd.get_dummies(X, columns=['cpu_brand', 'gpu_brand', 'ram_combined', 'hdd_combined'])

        model, scaler, rmse, r2 = train_model(X, y)

        print(f"배치 {batch_count} 모델 성능: RMSE = {rmse:.2f}, R^2 = {r2:.4f}")

        if rmse < best_rmse:
            best_rmse = rmse
            final_model = model
            final_scaler = scaler
            print(f"새로운 최고 모델 발견 (RMSE: {best_rmse:.2f})")

        batch_count += 1

    print(f"총 처리된 행 수: {total_rows:,}")

    # 최종 모델과 전처리 파이프라인 저장
    joblib.dump(final_model, 'final_model.joblib')
    joblib.dump(final_scaler, 'scaler.joblib')
    joblib.dump(X.columns, 'feature_names.joblib')
    joblib.dump(list(cpu_names), 'cpu_names.joblib')
    joblib.dump(list(gpu_names), 'gpu_names.joblib')

    print("모델과 전처리 파이프라인이 저장되었습니다.")
    print("AI 모델 학습을 완료하였습니다.")

    # 메모리 사용량 출력
    print(f"현재 메모리 사용량: {psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024:.2f} MB")

    # 연결 종료
    db.close()

    # 임시 파일 정리
    for i in range(MAX_BATCHES):
        try:
            os.remove(f'processed_data_batch_{i}.csv')
        except FileNotFoundError:
            pass  # 파일이 이미 삭제되었거나 존재하지 않는 경우 무시

    end_time = time.time()
    end_datetime = datetime.now()
    print(f"프로그램 종료: {end_datetime.strftime('%Y-%m-%d %H:%M:%S.%f')}")
    execution_time = end_time - start_time
    print(f"프로그램 총 실행 시간: {execution_time:.2f}초")

if __name__ == "__main__":
    main()
