import resource
import re
import pandas as pd
import mysql.connector
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import joblib
from used_pc_price_from_api_cpu_gpu_ram_hdd_table_entire_update_for_item_save_joblib import ensure_connection, process_na_price, clean_price, print_memory_usage
from db_config import get_db_connection
import numpy as np
import os
import psutil
import xgboost as xgb
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import GridSearchCV

# CPU와 GPU 모델 정규화 함수
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'

# db 객체 생성
db = get_db_connection()
cursor = db.cursor()

# 여기서부터 모델 학습 부분
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
LIMIT 1000000

"""

# 쿼리 실행 전 연결 확인
ensure_connection(db)

# 주요 연산 전후로 메모리 사용량 출력
print_memory_usage()

# 데이터를 저장할 빈 리스트 생성
chunks = []

chunk_size = 10000  # 적절한 크기로 조정하세요
max_rows = 1000000  # 최대 처리할 행 수
total_rows = 0

for chunk in pd.read_sql_query(query, db, chunksize=chunk_size):
    # 각 청크에 대해 데이터 처리 수행
    chunk['cpu_brand'] = chunk['cpu_name'].apply(normalize_cpu_gpu)
    chunk['gpu_brand'] = chunk['gpu_name'].apply(normalize_cpu_gpu)
    chunks.append(chunk)

    total_rows += len(chunk)
    print(f"처리된 행 수: {total_rows}")
    if total_rows >= max_rows:
        print(f"최대 {max_rows}행 처리 완료")
        break

# 모든 청크를 하나의 데이터프레임으로 결합
df = pd.concat(chunks, ignore_index=True)

# RAM 타입과 저장장치 타입 처리
df['ram_type'] = df['ram_type'].fillna('Unknown')
df['storage_type'] = df['storage_type'].str.lower()

# 특성 엔코딩
df = pd.get_dummies(df, columns=['cpu_brand', 'gpu_brand', 'ram_type', 'storage_type'])

# 가격 열을 숫자형으로 변환
price_columns = ['cpu_price', 'gpu_price', 'ram_price', 'hdd_price']
for col in price_columns:
    df[col] = pd.to_numeric(df[col], errors='coerce')

# RAM 크기와 HDD 크기를 숫자로 변환
df['ram_size'] = df['ram_size'].apply(lambda x: float(x.replace('GB', '').replace('gb', '')) if 'GB' in str(x).upper() else (float(str(x).replace('MB', '').replace('mb', '')) / 1000 if 'MB' in str(x).upper() else float(x)))
df['hdd_size'] = df['hdd_size'].apply(lambda x: float(x.replace('TB', '').replace('tb', '')) * 1000 if 'TB' in str(x).upper() else float(x.replace('GB', '').replace('gb', '')))

# 가격 합산으로 총 가격 계산
df['total_price'] = df['cpu_price'] + df['gpu_price'] + df['ram_price'] + df['hdd_price']

# 특성과 타겟 분리
X = df.drop(['cpu_name', 'gpu_name', 'cpu_price', 'gpu_price', 'ram_price', 'hdd_price', 'total_price'], axis=1)
y = df['total_price']

# 학습 및 테스트 데이터 분할
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(),
    'Random Forest': RandomForestRegressor(n_estimators=100, random_state=42),
    'XGBoost': xgb.XGBRegressor(random_state=42)
}

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

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=-1)
        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} (기준: <100, 낮을수록 좋음)")
    print(f"테스트 세트 MAE: {mae:.2f} (기준: <50, 낮을수록 좋음)")
    print(f"테스트 세트 R^2: {r2:.4f} (기준: >0.8, 높을수록 좋음)")
    print(f"테스트 세트 MAPE: {mape:.2f}% (기준: <10%, 낮을수록 좋음)")

    # 교차 검증
    cv_scores = cross_val_score(best_models[name], X_train_scaled, y_train, cv=5, scoring='neg_mean_squared_error')
    cv_rmse = np.sqrt(-cv_scores)
    print(f"교차 검증 RMSE: {cv_rmse.mean():.2f} (+/- {cv_rmse.std() * 2:.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()
