import gradio as gr
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
import torch
from PIL import Image
import numpy as np

# OCR 모델 및 프로세서 초기화
processor = TrOCRProcessor.from_pretrained('microsoft/trocr-base-handwritten')
model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-base-handwritten')

# 정답 및 해설 데이터베이스
answer_key = {
    "1": {
        "answer": "민주주의",
        "explanation": "민주주의는 국민이 주인이 되어 나라의 중요한 일을 결정하는 제도입니다. 우리나라는 민주주의 국가로, 국민들이 투표를 통해 대표자를 선출하고 중요한 결정에 참여합니다."
    },
    "2": {
        "answer": "삼권분립",
        "explanation": "삼권분립은 입법부, 행정부, 사법부로 권력을 나누어 서로 견제와 균형을 이루게 하는 제도입니다. 이를 통해 한 쪽에 권력이 집중되는 것을 막을 수 있습니다."
    },
    "3": {
        "answer": "지방자치제도",
        "explanation": "지방자치제도는 지역의 일을 그 지역 주민들이 직접 결정하고 처리하는 제도입니다. 주민들이 직접 지방자치단체장과 지방의회 의원을 선출합니다."
    }
}

def preprocess_image(image):
    """이미지 전처리 함수"""
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    return image

def recognize_text(image):
    """손글씨 인식 함수"""
    image = preprocess_image(image)
    pixel_values = processor(image, return_tensors="pt").pixel_values
    
    with torch.no_grad():
        generated_ids = model.generate(pixel_values)
    
    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return generated_text

def grade_answer(question_number, student_answer):
    """답안 채점 함수"""
    correct_answer = answer_key[question_number]["answer"]
    explanation = answer_key[question_number]["explanation"]
    
    # 답안 비교 (띄어쓰기, 대소문자 무시)
    is_correct = student_answer.replace(" ", "").lower() == correct_answer.replace(" ", "").lower()
    
    return {
        "정답 여부": "정답" if is_correct else "오답",
        "정답": correct_answer,
        "해설": explanation
    }

def process_submission(image, question_number):
    """전체 처리 함수"""
    if not image or not question_number:
        return "이미지와 문제 번호를 모두 입력해주세요."
    
    # 손글씨 인식
    recognized_text = recognize_text(image)
    
    # 채점 및 해설
    result = grade_answer(question_number, recognized_text)
    
    # 결과 포맷팅
    output = f"""
인식된 답안: {recognized_text}
채점 결과: {result['정답 여부']}
정답: {result['정답']}

[해설]
{result['해설']}
    """
    
    return output

# Gradio 인터페이스 생성
iface = gr.Interface(
    fn=process_submission,
    inputs=[
        gr.Image(label="답안 이미지를 업로드하세요", type="numpy"),
        gr.Dropdown(choices=["1", "2", "3"], label="문제 번호를 선택하세요")
    ],
    outputs=gr.Textbox(label="채점 결과"),
    title="초등학교 사회 시험지 채점 프로그램",
    description="손글씨로 작성된 사회 시험 답안을 채점하고 해설을 제공하는 프로그램입니다.",
)

if __name__ == "__main__":
    iface.launch()