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

# 전역 변수로 CPU 코어 수 설정
N_CORES = multiprocessing.cpu_count()

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_type'] = chunk['ram_type'].fillna('Unknown')
    chunk['storage_type'] = chunk['storage_type'].str.lower()

    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 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.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
    """

    chunk_size = 10000
    total_rows = 0

    # 데이터 누적을 위한 변수들
    X_list = []
    y_list = []

    # 청크를 모아서 한 번에 처리
    chunks = []
    for chunk in pd.read_sql_query(query, db, chunksize=chunk_size):
        chunks.append(chunk)
        if len(chunks) == N_CORES:
            processed_chunks = parallel_process_chunks(chunks)
            for processed_chunk in processed_chunks:
                X_chunk = processed_chunk.drop(['cpu_name', 'gpu_name', 'cpu_price', 'gpu_price', 'ram_price', 'hdd_price', 'total_price'], axis=1)
                y_chunk = processed_chunk['total_price']
                X_chunk = pd.get_dummies(X_chunk, columns=['cpu_brand', 'gpu_brand', 'ram_type', 'storage_type'])
                X_list.append(X_chunk)
                y_list.append(y_chunk)
                total_rows += len(processed_chunk)
            print(f"Processed {total_rows} rows")
            chunks = []
            gc.collect()

    # 남은 청크 처리
    if chunks:
        processed_chunks = parallel_process_chunks(chunks)
        for processed_chunk in processed_chunks:
            X_chunk = processed_chunk.drop(['cpu_name', 'gpu_name', 'cpu_price', 'gpu_price', 'ram_price', 'hdd_price', 'total_price'], axis=1)
            y_chunk = processed_chunk['total_price']
            X_chunk = pd.get_dummies(X_chunk, columns=['cpu_brand', 'gpu_brand', 'ram_type', 'storage_type'])
            X_list.append(X_chunk)
            y_list.append(y_chunk)
            total_rows += len(processed_chunk)
        print(f"Processed {total_rows} rows")

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

    # 모든 청크의 컬럼을 일치시키기
    all_columns = pd.concat([X.columns for X in X_list]).unique()
    for i, X in enumerate(X_list):
        missing_cols = set(all_columns) - set(X.columns)
        for col in missing_cols:
            X_list[i][col] = 0

    # 데이터 결합
    X = pd.concat(X_list, axis=0, ignore_index=True)
    y = pd.concat(y_list, ignore_index=True)

    del X_list, y_list
    gc.collect()

    # 학습 및 테스트 데이터 분할
    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)

    # 모델 정의 및 학습
    models = {
        'Linear Regression': LinearRegression(n_jobs=N_CORES),
        'Random Forest': RandomForestRegressor(n_estimators=100, random_state=42, n_jobs=N_CORES),
        'XGBoost': xgb.XGBRegressor(random_state=42, n_jobs=N_CORES)
    }

    # 하이퍼파라미터 튜닝을 위한 그리드
    param_grids = {
        'Random Forest': {
            'n_estimators': [100, 200],
            'max_depth': [10, 20],
            'min_samples_split': [2, 5],
            'min_samples_leaf': [1, 2]
        },
        'XGBoost': {
            'n_estimators': [100, 200],
            'max_depth': [3, 4],
            'learning_rate': [0.01, 0.1]
        }
    }

    best_models = {}

    for name, model in models.items():
        if name in param_grids:
            grid_search = GridSearchCV(model, param_grids[name], cv=5, scoring='neg_mean_squared_error', n_jobs=N_CORES)
            grid_search.fit(X_train_scaled, y_train)
            best_models[name] = grid_search.best_estimator_
        else:
            model.fit(X_train_scaled, y_train)
            best_models[name] = model

        y_pred = best_models[name].predict(X_test_scaled)
        mse = mean_squared_error(y_test, y_pred)
        rmse = np.sqrt(mse)
        mae = mean_absolute_error(y_test, y_pred)
        r2 = r2_score(y_test, y_pred)
        mape = np.mean(np.abs((y_test - y_pred) / y_test)) * 100

        print(f"\n{name} 모델 평가 결과:")
        print(f"테스트 세트 RMSE: {rmse:.2f}")
        print(f"테스트 세트 MAE: {mae:.2f}")
        print(f"테스트 세트 R^2: {r2:.4f}")
        print(f"테스트 세트 MAPE: {mape:.2f}%")

    # 최종 모델 선택 (예: XGBoost)
    final_model = best_models['XGBoost']

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

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

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

    # 연결 종료
    db.close()

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