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

# 가장 빈도가 높은 1만개 단어만 사용하여 데이터셋 로드
(x_train, y_train), (x_test, y_test) = keras.datasets.imdb.load_data(num_words=10000)

print(f"패딩 전 첫 번째 리뷰 길이: {len(x_train[0])}")

# 모든 시퀀스의 길이를 256으로 통일
# maxlen보다 길면 잘라내고, 짧으면 앞부분을 0으로 채움 (pre-padding)
x_train = keras.preprocessing.sequence.pad_sequences(x_train, maxlen=256)
x_test = keras.preprocessing.sequence.pad_sequences(x_test, maxlen=256)

print(f"패딩 후 첫 번째 리뷰 길이: {len(x_train[0])}")

# 모델 아키텍처 정의
model = keras.Sequential([
    # 1. 단어 임베딩 층
    # input_dim: 전체 단어 집합의 크기 (가장 빈번한 1만개 단어)
    # output_dim: 각 단어를 표현할 벡터의 차원 (32차원)
    keras.layers.Embedding(input_dim=10000, output_dim=32),
    
    # 2. RNN 층
    # units: 은닉 상태 벡터의 차원 (32차원)
    keras.layers.SimpleRNN(32),
    
    # 3. 최종 분류기(Classifier)
    # units: 출력 뉴런의 수 (긍정/부정 1개)
    # activation: 출력 값을 0~1 사이 확률로 변환 (이진 분류)
    keras.layers.Dense(1, activation="sigmoid"),
])


model.compile(
    # 손실 함수: 예측이 정답과 얼마나 다른지 측정.
    # 이진 분류(0 또는 1) 문제이므로 binary_crossentropy가 가장 적합.
    loss="binary_crossentropy",
    
    # 옵티마이저: 손실을 최소화하기 위해 모델의 가중치를 업데이트하는 알고리즘.
    # Adam은 현재 가장 널리 쓰이고 성능이 좋은 옵티마이저 중 하나.
    optimizer="adam",
    
    # 평가지표: 훈련 과정을 모니터링할 지표. 정확도를 사용.
    metrics=["accuracy"]
)

batch_size = 128
epochs = 10

# 모델 학습 실행
# validation_data를 지정하여 매 에포크마다 테스트 데이터로 성능을 검증
history = model.fit(
    x_train, y_train,
    batch_size=batch_size,
    epochs=epochs,
    validation_data=(x_test, y_test)
)

# 학습 완료 후 최종 성능 평가
score = model.evaluate(x_test, y_test, verbose=0)
print(f"\nTest loss: {score[0]:.4f}")
print(f"Test accuracy: {score[1]:.4f}")

# 모델의 구조, 가중치, 학습 설정을 모두 '.keras' 파일 하나에 저장
model.save("my_rnn_model_imdb.keras")

# 모델 구조 요약 출력
model.summary()

# 저장된 모델 불러오기
loaded_model = keras.models.load_model("my_rnn_model_imdb.keras")