import gradio as gr
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, models
import os

# 1. 모델 아키텍처 재정의 (저장된 가중치를 불러오기 위함)
class DynamicConv2D(layers.Layer):
    def __init__(self, k=3, **kwargs):
        super().__init__(**kwargs)
        assert k % 2 == 1
        self.k = k
        self.generator = layers.Dense(k * k)

    def call(self, x):
        B, H, W, C = tf.shape(x)[0], tf.shape(x)[1], tf.shape(x)[2], tf.shape(x)[3]
        kernels = self.generator(x)
        kernels = tf.nn.softmax(kernels, axis=-1)
        pad = (self.k - 1) // 2
        x_pad = tf.pad(x, [[0, 0], [pad, pad], [pad, pad], [0, 0]])
        patches = tf.image.extract_patches(
            images=x_pad,
            sizes=[1, self.k, self.k, 1],
            strides=[1, 1, 1, 1],
            rates=[1, 1, 1, 1],
            padding='VALID'
        )
        patches = tf.reshape(patches, [B, H, W, self.k * self.k, C])
        kernels_exp = tf.expand_dims(kernels, axis=-1)
        return tf.reduce_sum(patches * kernels_exp, axis=3)

def build_dynamic_model(input_shape=(28, 28, 1), num_classes=26):
    inputs = layers.Input(shape=input_shape)
    
    # 일반 Conv로 기초 특징 추출
    x = layers.Conv2D(32, (3, 3), activation='relu', padding='same')(inputs)
    x = layers.BatchNormalization()(x)
    
    # Dynamic Convolution 적용
    x = DynamicConv2D(k=3)(x)
    x = layers.Activation('relu')(x)
    x = layers.MaxPooling2D((2, 2))(x)
    
    x = layers.Conv2D(64, (3, 3), activation='relu', padding='same')(x)
    x = DynamicConv2D(k=3)(x)
    x = layers.Activation('relu')(x)
    x = layers.GlobalAveragePooling2D()(x)
    
    x = layers.Dense(128, activation='relu')(x)
    outputs = layers.Dense(num_classes, activation='softmax')(x)
    
    return models.Model(inputs, outputs)

# 2. 모델 로드 (파일 경로 확인 필요)
model = build_dynamic_model()
# 파일 경로를 절대 경로로 설정
weights_path = 'dynamic_conv_alphabet.weights.h5'

if os.path.exists(weights_path):
    model.load_weights(weights_path)
    print(f"가중치 로드 성공: {weights_path}")
else:
    print(f"오류: {weights_path} 파일을 찾을 수 없습니다. 먼저 학습을 수행하세요.")
# 3. 예측 함수 정의

# ... (상단 모델 정의 및 로드 부분은 동일)

def classify_alphabet(sketchpad):
    # 수정: sketchpad가 None이거나 데이터가 없는 경우 처리
    if sketchpad is None or (isinstance(sketchpad, dict) and sketchpad.get("composite") is None):
        return {"상태": "글씨를 기다리는 중..."}
    
    try:
        # 3-1. 이미지 데이터 추출
        # Gradio 버전에 따라 구조가 다를 수 있으므로 안전하게 접근
        img_data = sketchpad["composite"]
        
        # 만약 이미지가 투명도가 없는 3채널이라면 흑백 전환, 4채널이면 Alpha 사용
        if img_data.shape[-1] == 4:
            img = img_data[:, :, 3] # Alpha 채널 (글씨 부분)
        else:
            img = tf.image.rgb_to_grayscale(img_data)[:, :, 0]

        # 3-2. 전처리
        img = tf.cast(img, tf.float32)
        img = tf.image.resize(tf.expand_dims(img, axis=-1), (28, 28))
        
        # EMNIST 데이터 방향에 맞게 전치(Transpose)
        img = tf.image.transpose(img) 
        
        img = img / 255.0
        img = tf.expand_dims(img, axis=0)

        # 3-3. 추론
        preds = model.predict(img, verbose=0)[0]
        
        results = {}
        for i in range(26):
            char = chr(ord('A') + i)
            results[char] = float(preds[i])
            
        return results
    except Exception as e:
        return {"에러": str(e)}

# 4. Gradio 인터페이스 설정 수정
interface = gr.Interface(
    fn=classify_alphabet,
    inputs=gr.Sketchpad(label="알파벳을 그려보세요 (A-Z)", type="numpy"),
    outputs=gr.Label(num_top_classes=3, label="예측 결과"),
    title="Dynamic Conv Alphabet Recognizer",
    live=True
)

if __name__ == "__main__":
    # 수정: SSR 모드와 Hot Reload 관련 에러를 방지하기 위해 설정을 추가합니다.
    interface.launch(
        server_name="0.0.0.0", 
        ssr_mode=False  # 로그에 나온 SSR 관련 이슈 방지
    )