# 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.memory_size AS ram_size, ram.ramtype AS ram_type, 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
# query 에서 각각의 input 폼필드에서 cpu.cpu_name AS cpu_name 과 gpu.videocard_name AS gpu_name 과 ram.ramtype AS ram_type + ram.memory_size AS ram_size 과 hdd.storage_type AS storage_type + hdd.size AS hdd_size 에 해당하는 각각의 폼필드에 입력된 값을 이 파일에 의해 생성되는 학습모델 생성파일을 활용해서 cpu_price + gpu_price + ram_price + hdd_price 를 합산한 결과 데이터를 활용하는 데 학습모델 또는 예측모델을 사용할려고 해. 단 이떄 입력 폼필드에 의해 입력되는 데이터값중에  cpu_name 과  gpu_name 은 데이터베이스에 저장되어 있는 컬럼값과 100% 일치하지 않을 수도 있으므로 입력받은 부분적인 데이터를 통해 일치할 확률이 높은 cpu_name 으로 예측해서 신뢰도를 향샹시키도록 해야 해. 이러한 의도와 조건들을 충족하도록 한 소스코드를 현재 소스코드를 근거로 하여 수정한 후에 전체 소스코드를 제시해줘'

# 1. CPU와 GPU 이름에 대한 부분 일치 검색 기능 추가
# 2. RAM 타입과 크기를 결합한 특성 생성
# 3. HDD 타입과 크기를 결합한 특성 생성
# 4. 예측을 위한 별도의 함수 추가
# 5. 각 청크를 처리한 후 즉시 CSV 파일로 저장하여 메모리 사용을 최소화하고 또 메모리 사용량을 주기적으로 확인해서 RAM용량이 부족해서 프로그램이 중단되는 일이 생기지 않게 함
# 6. 중간 결과를 저장하는 CSV 파일을 여러 개의 작은 파일로 나누어 저장합니다.
# 7. 일정 크기의 데이터를 처리한 후 모델 학습을 수행하고, 그 결과를 누적합니다.
# 8. 디스크 공간을 주기적으로 확인하고, 부족할 경우 이전 중간 결과를 삭제합니다.

import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor
import xgboost as xgb
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from sklearn.model_selection import GridSearchCV
import joblib
from db_config import get_db_connection
import gc
import time
from datetime import datetime
import psutil
import os
import multiprocessing
from functools import partial
from fuzzywuzzy import process
import shutil

# 전역 변수 설정
N_CORES = multiprocessing.cpu_count()
MEMORY_THRESHOLD = 80  # 메모리 사용량 임계값 (%)
DISK_THRESHOLD = 10  # 디스크 여유 공간 임계값 (GB)
CHUNK_SIZE = 10000
BATCH_SIZE = 1000000  # 한 번에 처리할 행 수

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

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_name': [cpu_match],
        'gpu_name': [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['cpu_brand'] = input_data['cpu_name'].apply(normalize_cpu_gpu)
    input_data['gpu_brand'] = input_data['gpu_name'].apply(normalize_cpu_gpu)

    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'):
    columns = ['cpu_name', 'cpu_price', 'gpu_name', 'gpu_price', 'ram_type', 'ram_size', 'ram_price',
               'storage_type', 'hdd_size', 'hdd_price', 'cpu_brand', 'gpu_brand', 'ram_combined',
               'hdd_combined', 'total_price']
    chunk.to_csv(output_file, mode=mode, header=(mode=='w'), index=False, columns=columns)

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

        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(processed_chunk['cpu_name'])
                    gpu_names.update(processed_chunk['gpu_name'])

                output_file = f'processed_data_batch_{batch_count}.csv'
                for processed_chunk in processed_chunks:
                    save_chunk_to_csv(processed_chunk, output_file, mode='a')

                chunks = []
                check_memory_usage()

                if not check_disk_space('.'):
                    print("디스크 공간이 부족합니다. 이전 배치 파일을 삭제합니다.")
                    for i in range(max(0, batch_count - 1)):
                        os.remove(f'processed_data_batch_{i}.csv')

        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')
                total_rows += len(processed_chunk)
                cpu_names.update(processed_chunk['cpu_name'])
                gpu_names.update(processed_chunk['gpu_name'])

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

        # 현재 배치에 대해 모델 학습
        columns = ['cpu_name', 'cpu_price', 'gpu_name', 'gpu_price', 'ram_type', 'ram_size', 'ram_price',
                   'storage_type', 'hdd_size', 'hdd_price', 'cpu_brand', 'gpu_brand', 'ram_combined',
                   'hdd_combined', 'total_price']
        df = pd.read_csv(output_file, names=columns, header=None)
        print("데이터프레임 열:", df.columns.tolist())  # 열 이름 출력

        # 삭제할 열 목록
        columns_to_drop = ['cpu_name', 'gpu_name', 'cpu_price', 'gpu_price', 'ram_price', 'hdd_price', 'ram_type', 'storage_type']

        # 실제로 존재하는 열만 삭제
        columns_to_drop = [col for col in columns_to_drop if col in df.columns]
        X = df.drop(columns_to_drop, axis=1)

        if 'total_price' in df.columns:
            y = df['total_price']
        else:
            print("'total_price' 열이 없습니다. 데이터를 확인해 주세요.")
            break

        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(batch_count):
        os.remove(f'processed_data_batch_{i}.csv')

    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()
