"""
스탠스 분류 추론 예시 코드

Usage:
    from inference import StancePredictor

    predictor = StancePredictor("your-username/korean-news-stance-classifier")
    result = predictor.predict("정부의 새 정책이 경제 성장에 기여할 것으로 기대된다")
    print(result)
"""

import torch
import torch.nn as nn
from transformers import BertModel, AutoTokenizer
from huggingface_hub import hf_hub_download
import json


class StanceClassifier(nn.Module):
    """KoBERT 기반 스탠스 분류 모델"""

    def __init__(self, n_classes=3, dropout=0.3, model_name="skt/kobert-base-v1"):
        super(StanceClassifier, self).__init__()
        self.bert = BertModel.from_pretrained(model_name)
        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(self.bert.config.hidden_size, n_classes)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.pooler_output
        pooled_output = self.dropout(pooled_output)
        return self.classifier(pooled_output)


class StancePredictor:
    """HuggingFace Hub에서 모델을 로드하여 스탠스 예측"""

    def __init__(self, repo_id: str, device: str = None):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.label_names = ["옹호", "중립", "비판"]
        self.label_names_en = ["support", "neutral", "oppose"]

        # 설정 로드
        config_path = hf_hub_download(repo_id, "config.json")
        with open(config_path, "r", encoding="utf-8") as f:
            self.config = json.load(f)

        # 토크나이저 로드 (반드시 monologg/kobert!)
        self.tokenizer = AutoTokenizer.from_pretrained(
            self.config.get("tokenizer", "monologg/kobert"),
            trust_remote_code=True
        )

        # 모델 로드
        model_path = hf_hub_download(repo_id, "model.pth")
        self.model = StanceClassifier(
            n_classes=self.config.get("num_labels", 3),
            dropout=self.config.get("dropout", 0.3),
            model_name=self.config.get("base_model", "skt/kobert-base-v1")
        )
        self.model.load_state_dict(torch.load(model_path, map_location=self.device))
        self.model.to(self.device)
        self.model.eval()

    def predict(self, text: str) -> dict:
        """단일 텍스트 스탠스 예측"""
        inputs = self.tokenizer(
            text,
            return_tensors="pt",
            max_length=self.config.get("max_length", 512),
            truncation=True,
            padding="max_length"
        )

        input_ids = inputs["input_ids"].to(self.device)
        attention_mask = inputs["attention_mask"].to(self.device)

        with torch.no_grad():
            outputs = self.model(input_ids, attention_mask)
            probs = torch.softmax(outputs, dim=1)[0]
            pred = torch.argmax(probs).item()

        return {
            "stance": self.label_names_en[pred],
            "stance_kr": self.label_names[pred],
            "confidence": round(probs[pred].item(), 4),
            "probabilities": {
                "support": round(probs[0].item(), 4),
                "neutral": round(probs[1].item(), 4),
                "oppose": round(probs[2].item(), 4)
            }
        }

    def predict_batch(self, texts: list, batch_size: int = 16) -> list:
        """배치 텍스트 스탠스 예측"""
        results = []

        for i in range(0, len(texts), batch_size):
            batch = texts[i:i + batch_size]
            inputs = self.tokenizer(
                batch,
                return_tensors="pt",
                max_length=self.config.get("max_length", 512),
                truncation=True,
                padding="max_length"
            )

            input_ids = inputs["input_ids"].to(self.device)
            attention_mask = inputs["attention_mask"].to(self.device)

            with torch.no_grad():
                outputs = self.model(input_ids, attention_mask)
                probs = torch.softmax(outputs, dim=1)

            for j in range(len(batch)):
                pred = torch.argmax(probs[j]).item()
                results.append({
                    "stance": self.label_names_en[pred],
                    "stance_kr": self.label_names[pred],
                    "confidence": round(probs[j][pred].item(), 4),
                    "probabilities": {
                        "support": round(probs[j][0].item(), 4),
                        "neutral": round(probs[j][1].item(), 4),
                        "oppose": round(probs[j][2].item(), 4)
                    }
                })

        return results


if __name__ == "__main__":
    # 사용 예시
    predictor = StancePredictor("your-username/korean-news-stance-classifier")

    test_texts = [
        "정부의 새 정책이 경제 성장에 크게 기여할 것으로 기대된다",
        "야당은 졸속 행정이라며 강하게 반발했다",
        "국회에서 법안 심의가 진행되고 있다"
    ]

    for text in test_texts:
        result = predictor.predict(text)
        print(f"텍스트: {text}")
        print(f"결과: {result}")
        print("-" * 50)