import torch
import torch.nn as nn
import torch.nn.functional as F
import pickle
import joblib
import numpy as np

# SimpleClassifier 클래스 정의
class SimpleClassifier(nn.Module):
    def __init__(self, input_size, num_classes):
        super(SimpleClassifier, self).__init__()
        self.fc1 = nn.Linear(input_size, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, num_classes)
        self.dropout = nn.Dropout(0.3)
        
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = F.relu(self.fc2(x))
        x = self.dropout(x)
        x = self.fc3(x)
        return x

def test_current_model():
    """현재 모델 테스트"""
    print("=== 현재 모델 테스트 ===")
    
    try:
        # 설정 로드
        with open('config.json', 'r', encoding='utf-8') as f:
            import json
            config = json.load(f)
        
        id2label = config.get('id2label', {})
        print(f"라벨 수: {len(id2label)}")
        
        # 모델 로드
        input_size = 3000  # TF-IDF 특성 수
        num_classes = len(id2label)
        model = SimpleClassifier(input_size, num_classes)
        model.load_state_dict(torch.load('pytorch_model.bin', map_location='cpu'))
        
        # 벡터라이저 로드
        vectorizer = joblib.load('vectorizer.pkl')
        
        model.eval()
        
        # 테스트 단어들 (환원철 포함)
        test_words = ["철ㄹ", "CaO", "해면철", "등류", "환원철"]
        
        for word in test_words:
            print(f"\n{'='*50}")
            print(f"입력: '{word}'")
            print(f"{'='*50}")
            
            # TF-IDF 벡터화
            word_vector = vectorizer.transform([word]).toarray()
            word_tensor = torch.FloatTensor(word_vector)
            
            with torch.no_grad():
                outputs = model(word_tensor)
                probabilities = F.softmax(outputs, dim=1)
                
                # 상위 5개 예측
                top_probs, top_indices = torch.topk(probabilities, 5, dim=1)
                
                print(f"최대 확률: {probabilities.max().item():.4f} ({probabilities.max().item()*100:.1f}%)")
                print(f"상위 5개 예측:")
                
                for i in range(5):
                    label_id = top_indices[0][i].item()
                    probability = top_probs[0][i].item()
                    label = id2label.get(str(label_id), f"Unknown_{label_id}")
                    print(f"  {i+1}. {label}: {probability:.4f} ({probability*100:.1f}%)")
                    
    except Exception as e:
        print(f"에러 발생: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    test_current_model()