# test.py (오류 수정 최종 코드)

import numpy as np
import tensorflow as tf
from tensorflow import keras
# from_pretrained_keras 대신 hf_hub_download를 사용합니다.
from huggingface_hub import hf_hub_download

print("TensorFlow 버전:", tf.__version__)

# 1. Hugging Face Hub에서 모델 파일 다운로드 후 로드
REPO_ID = "OneclickAI/LSTM_GUE_test_Model"
print(f"\n'{REPO_ID}' 저장소에서 모델 파일의 위치를 확인합니다...")

try:
    # 1단계: hf_hub_download로 파일의 로컬 캐시 경로를 가져옵니다.
    # 파일이 이미 다운로드 되었다면, 다운로드를 생략하고 경로만 즉시 반환합니다.
    print("LSTM 모델 경로 확인 중...")
    lstm_model_path = hf_hub_download(repo_id=REPO_ID, filename="lstm_model.keras")
    print(f"LSTM 모델 파일 위치: {lstm_model_path}")
    
    print("GRU 모델 경로 확인 중...")
    gru_model_path = hf_hub_download(repo_id=REPO_ID, filename="gru_model.keras")
    print(f"GRU 모델 파일 위치: {gru_model_path}")

    # 2단계: 다운로드된 파일 경로를 Keras의 표준 load_model 함수로 직접 로드합니다.
    print("\nKeras로 모델을 로드합니다...")
    lstm_model = keras.models.load_model(lstm_model_path)
    gru_model = keras.models.load_model(gru_model_path)
    
    print("모델을 성공적으로 로드했습니다.")

except Exception as e:
    print(f"모델 로딩 중 오류 발생: {e}")
    print("인터넷 연결 및 저장소 ID, 파일명을 확인해주세요.")
    exit()

# IMDB 데이터셋의 단어 인덱스 로드 ('단어': 정수)
word_index = keras.datasets.imdb.get_word_index()

# 2. 예측할 리뷰 문장 정의
review1 = "This movie was fantastic and wonderful. I really enjoyed it."
review2 = "It was a complete waste of time. The plot was terrible and the acting was bad."

# 3. 문장 전처리 함수
def preprocess_text(text, word_index, maxlen=256):
    """
    텍스트를 모델이 이해할 수 있는 정수 시퀀스로 변환하고 패딩합니다.
    """
    # 문장을 소문자로 변환하고 단어 단위로 분할
    tokens = text.lower().split()
    
    # 각 단어를 정수 인덱스로 변환 (word_index에 없으면 2번 인덱스'<unk>' 사용)
    token_indices = [word_index.get(word, 2) for word in tokens]
    
    # 시퀀스 패딩
    padded_sequence = keras.preprocessing.sequence.pad_sequences([token_indices], maxlen=maxlen)
    
    return padded_sequence

# 4. 모델 예측 및 결과 출력 함수
def predict_review(review_text, model, model_name):
    """
    전처리된 텍스트를 사용하여 감성 분석을 수행하고 결과를 출력합니다.
    """
    # 문장 전처리
    processed_review = preprocess_text(review_text, word_index)
    
    # 예측 수행
    prediction = model.predict(processed_review, verbose=0) # 예측 시 로그 출력을 끔
    positive_probability = prediction[0][0] * 100
    
    print(f"--- {model_name} 모델 예측 결과 ---")
    print(f"리뷰: '{review_text}'")
    print(f"긍정 확률: {positive_probability:.2f}%")
    if positive_probability > 50:
        print("결과: 긍정적인 리뷰입니다.")
    else:
        print("결과: 부정적인 리뷰입니다.")
    print("-" * 30)

# 5. 각 리뷰에 대해 두 모델로 예측 수행
print("\n" + "="*40)
print("첫 번째 리뷰 예측 시작")
print("="*40)
predict_review(review1, lstm_model, "LSTM")
predict_review(review1, gru_model, "GRU")


print("\n" + "="*40)
print("두 번째 리뷰 예측 시작")
print("="*40)
predict_review(review2, lstm_model, "LSTM")
predict_review(review2, gru_model, "GRU")