import numpy as np
import tensorflow as tf
from tensorflow import keras
from keras import layers

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

# 1. 데이터 로드 및 전처리
print("\n1. 데이터 로드 및 전처리를 시작합니다...")
# num_words=10000: 가장 빈도가 높은 1만 개의 단어만 사용
(x_train, y_train), (x_test, y_test) = keras.datasets.imdb.load_data(num_words=10000)

print(f"학습 데이터 개수: {len(x_train)}")
print(f"테스트 데이터 개수: {len(x_test)}")

# 문장의 길이를 동일하게 맞추기 위해 패딩(padding) 처리 (maxlen=256)
x_train = keras.preprocessing.sequence.pad_sequences(x_train, maxlen=256)
x_test = keras.preprocessing.sequence.pad_sequences(x_test, maxlen=256)
print("데이터 전처리가 완료되었습니다.")

# 2. LSTM 모델 생성, 학습 및 저장
print("\n2. LSTM 모델 학습을 시작합니다...")

# LSTM 모델 아키텍처 정의
lstm_model = keras.Sequential([
    layers.Embedding(input_dim=10000, output_dim=128),
    layers.LSTM(64),
    layers.Dense(1, activation="sigmoid")
])

# 모델 컴파일
lstm_model.compile(
    loss="binary_crossentropy",
    optimizer="adam",
    metrics=["accuracy"]
)

print("\n--- LSTM 모델 구조 ---")
lstm_model.summary()

# 모델 학습
batch_size = 128
epochs = 1 # 예제이므로 epoch를 줄여서 실행 시간을 단축합니다.
history_lstm = lstm_model.fit(
    x_train, y_train,
    batch_size=batch_size,
    epochs=epochs,
    validation_data=(x_test, y_test)
)

# 모델 평가
score_lstm = lstm_model.evaluate(x_test, y_test, verbose=0)
print(f"\nLSTM 모델 테스트 결과 -> Loss: {score_lstm[0]:.4f}, Accuracy: {score_lstm[1]:.4f}\n")

# 학습된 LSTM 모델 저장
lstm_model.save("lstm_model.keras")
print("LSTM 모델이 'lstm_model.keras' 파일로 저장되었습니다.")


# 3. GRU 모델 생성, 학습 및 저장
print("\n3. GRU 모델 학습을 시작합니다...")

# GRU 모델 아키텍처 정의
gru_model = keras.Sequential([
    layers.Embedding(input_dim=10000, output_dim=128),
    layers.GRU(64),
    layers.Dense(1, activation="sigmoid")
])

# 모델 컴파일
gru_model.compile(
    loss="binary_crossentropy",
    optimizer="adam",
    metrics=["accuracy"]
)

print("\n--- GRU 모델 구조 ---")
gru_model.summary()


# 모델 학습
history_gru = gru_model.fit(
    x_train, y_train,
    batch_size=batch_size,
    epochs=epochs,
    validation_data=(x_test, y_test)
)

# 모델 평가
score_gru = gru_model.evaluate(x_test, y_test, verbose=0)
print(f"\nGRU 모델 테스트 결과 -> Loss: {score_gru[0]:.4f}, Accuracy: {score_gru[1]:.4f}")


# 학습된 GRU 모델 저장
gru_model.save("gru_model.keras")
print("GRU 모델이 'gru_model.keras' 파일로 저장되었습니다.")