import os
import sys
import itertools
import numpy as np
import tensorflow as tf
from sklearn.model_selection import train_test_split

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.model import MalConv
from src.utils import preprocess_dataset

def hyperparameter_search(csv_path, 
                         param_grid=None, 
                         max_length=2**20,
                         epochs=5,
                         validation_split=0.2):
    """
    그리드 서치를 통한 하이퍼파라미터 최적화
    
    Args:
        csv_path: 훈련 데이터 CSV 경로
        param_grid: 하이퍼파라미터 그리드
        max_length: 최대 입력 길이
        epochs: 훈련 에포크 수
        validation_split: 검증 데이터 비율
    """
    
    if param_grid is None:
        param_grid = {
            'embedding_size': [8, 16],
            'num_filters': [64, 128],
            'fc_size': [64, 128],
            'learning_rate': [0.001, 0.0001]
        }
    
    print("데이터 로딩 중...")
    X, y = preprocess_dataset(csv_path, max_length)
    X_train, X_val, y_train, y_val = train_test_split(
        X, y, test_size=validation_split, random_state=42, stratify=y
    )
    
    # 모든 하이퍼파라미터 조합 생성
    param_names = list(param_grid.keys())
    param_values = list(param_grid.values())
    param_combinations = list(itertools.product(*param_values))
    
    best_score = 0
    best_params = None
    results = []
    
    print(f"총 {len(param_combinations)}개의 조합을 테스트합니다.")
    
    for i, params in enumerate(param_combinations):
        param_dict = dict(zip(param_names, params))
        print(f"\n[{i+1}/{len(param_combinations)}] 테스트 중: {param_dict}")
        
        try:
            # 모델 생성
            model = MalConv(
                max_input_length=max_length,
                embedding_size=param_dict['embedding_size'],
                num_filters=param_dict['num_filters'],
                fc_size=param_dict['fc_size']
            )
            
            # 컴파일
            model.compile(
                optimizer=tf.keras.optimizers.Adam(
                    learning_rate=param_dict['learning_rate']
                ),
                loss='binary_crossentropy',
                metrics=['accuracy']
            )
            
            # 더미 입력으로 모델 빌드
            dummy_input = np.zeros((1, max_length), dtype=np.uint8)
            _ = model(dummy_input)
            
            # 훈련
            history = model.fit(
                X_train, y_train,
                batch_size=16,
                epochs=epochs,
                validation_data=(X_val, y_val),
                verbose=0
            )
            
            # 평가
            val_loss, val_acc = model.evaluate(X_val, y_val, verbose=0)
            
            result = {
                'params': param_dict,
                'val_accuracy': val_acc,
                'val_loss': val_loss
            }
            results.append(result)
            
            print(f"검증 정확도: {val_acc:.4f}")
            
            # 최고 성능 업데이트
            if val_acc > best_score:
                best_score = val_acc
                best_params = param_dict
                print(f"새로운 최고 성능! 정확도: {best_score:.4f}")
            
        except Exception as e:
            print(f"에러 발생: {e}")
            continue
    
    print("\n" + "="*50)
    print("하이퍼파라미터 튜닝 완료")
    print("="*50)
    print(f"최고 성능: {best_score:.4f}")
    print(f"최적 하이퍼파라미터: {best_params}")
    
    # 결과 정렬
    results.sort(key=lambda x: x['val_accuracy'], reverse=True)
    
    print("\n상위 5개 결과:")
    for i, result in enumerate(results[:5]):
        print(f"{i+1}. 정확도: {result['val_accuracy']:.4f}, "
              f"파라미터: {result['params']}")
    
    return best_params, results

def main():
    csv_path = "Input/sample_data.csv"  # 실제 데이터 경로로 변경
    
    # 커스텀 하이퍼파라미터 그리드
    param_grid = {
        'embedding_size': [8, 16],
        'num_filters': [64, 128],
        'fc_size': [64, 128],
        'learning_rate': [0.001, 0.0001]
    }
    
    best_params, results = hyperparameter_search(
        csv_path=csv_path,
        param_grid=param_grid,
        epochs=3  # 빠른 테스트를 위해 에포크 수 감소
    )
    
    print(f"\n최적 하이퍼파라미터로 모델을 다시 훈련하세요:")
    print(f"python src/train.py {csv_path} --epochs 10")

if __name__ == "__main__":
    main()