import tensorflow as tf
# hf_hub_download 함수를 직접 사용하기 위해 import 합니다.
from huggingface_hub import hf_hub_download 
from PIL import Image, ImageOps
import numpy as np
import os

def load_model_from_hf(model_id, model_filename="model.keras"):
    """
    Hugging Face Hub에서 Keras 모델 파일을 다운로드한 후,
    로컬에 저장된 파일을 이용해 모델을 불러옵니다.
    """
    try:
        print(f"Downloading model '{model_id}' from Hugging Face Hub...")
        # 1. Hugging Face Hub에서 모델 파일 다운로드
        #    이 함수는 다운로드된 파일이 저장된 로컬 경로를 반환합니다.
        model_path = hf_hub_download(repo_id=model_id, filename=model_filename)
        print(f"Model downloaded to: {model_path}")
        
        # 2. 로컬에 저장된 모델 파일 로드
        print("Loading model from local file...")
        model = tf.keras.models.load_model(model_path)
        print("Model loaded successfully!")
        return model
        
    except Exception as e:
        print(f"Error loading model: {e}")
        print("Please check if the model ID and filename are correct on Hugging Face Hub.")
        return None

def preprocess_image(image_path):
    """
    사용자 이미지를 MNIST 데이터셋 형식에 맞게 전처리합니다.
    """
    try:
        # 1. 이미지 열기
        img = Image.open(image_path)
        
        # 2. 흑백(Grayscale)으로 변환
        img = img.convert('L')
        
        # 3. 색상 반전 (MNIST는 흰색 글씨/검은 배경, 사용자는 보통 검은 글씨/흰색 배경)
        if np.mean(np.array(img)) > 128:
            img = ImageOps.invert(img)

        # 4. 28x28 크기로 리사이즈
        img = img.resize((28, 28), Image.Resampling.LANCZOS)
        
        # 5. Numpy 배열로 변환하고 0~1 사이 값으로 정규화
        img_array = np.array(img).astype('float32') / 255.0
        
        # 6. 모델의 입력 형태에 맞게 차원 확장 (1, 28, 28, 1)
        processed_img = np.expand_dims(img_array, axis=0)
        processed_img = np.expand_dims(processed_img, axis=-1)
        
        return processed_img
    except FileNotFoundError:
        print(f"Error: The file '{image_path}' was not found.")
        return None
    except Exception as e:
        print(f"Error processing image: {e}")
        return None

def main():
    # Hugging Face에 업로드된 모델 ID
    model_id = "OneclickAI/CNN_test_Model"
    
    # 모델 로드 (수정된 함수 호출)
    # 이전 train.py에서 model.save("my_keras_model.keras")로 저장했으므로, 
    # Hub에 올라간 파일 이름은 'my_keras_model.keras'일 것입니다. 
    # 만약 다른 이름으로 올렸다면 해당 파일명으로 수정해주세요.
    # (Hugging Face Hub에서는 보통 'model.keras' 라는 이름을 권장합니다)
    model = load_model_from_hf(model_id, model_filename="my_keras_model.keras") 
    
    if model is None:
        return # 모델 로드 실패 시 종료

    # 사용자에게 이미지 경로를 계속해서 입력받음
    while True:
        user_input = input("\nPlease enter the path to your image (or type 'exit' to quit): ")
        
        if user_input.lower() == 'exit':
            break
            
        if not os.path.exists(user_input):
            print(f"File not found at '{user_input}'. Please check the path and try again.")
            continue

        # 이미지 전처리
        processed_image = preprocess_image(user_input)
        
        if processed_image is not None:
            # 모델 예측 수행
            predictions = model.predict(processed_image)
            
            # 가장 높은 확률을 가진 클래스(숫자)를 찾음
            predicted_digit = np.argmax(predictions[0])
            confidence = np.max(predictions[0]) * 100
            
            print("\n--- Prediction Result ---")
            print(f"Predicted Digit: {predicted_digit}")
            print(f"Confidence: {confidence:.2f}%")
            print("-------------------------")

if __name__ == "__main__":
    main()