import os
import sys
import pickle
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# 기본 경로 설정
BASE_DIR = Path(__file__).resolve().parent
MODEL_DIR = BASE_DIR / "saved_models"

KC_DIR = MODEL_DIR / "kcbert_web"
DEBERTA_DIR = MODEL_DIR / "deberta_web"
CNN_PATH = MODEL_DIR / "char_cnn_web.pt"
VOCAB_PATH = MODEL_DIR / "vocab.pkl"

# Char-CNN 네트워크 구조 정의
class MultiScaleCharCNNEncoder(nn.Module):
    def __init__(self, vocab_size, embed_dim, num_filters, filter_sizes):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.convs = nn.ModuleList([
            nn.Conv1d(embed_dim, num_filters, kernel_size=fs)
            for fs in filter_sizes
        ])
        self.fc = nn.Linear(len(filter_sizes) * num_filters, 256)

    def forward(self, x):
        x = self.embedding(x).transpose(1, 2)
        conved = [torch.relu(conv(x)) for conv in self.convs]
        pooled = [torch.max_pool1d(c, c.shape[2]).squeeze(2) for c in conved]
        return self.fc(torch.cat(pooled, dim=1))

class SiameseNetwork(nn.Module):
    def __init__(self, encoder):
        super().__init__()
        self.encoder = encoder
        self.classifier = nn.Sequential(
            nn.Linear(256 * 2, 128),
            nn.ReLU(),
            nn.Linear(128, 2)
        )

    def forward(self, a, b):
        a_vec = self.encoder(a)
        b_vec = self.encoder(b)
        return self.classifier(torch.cat((a_vec, b_vec), dim=1))

# 파일 존재 검증
def check_model_files():
    required_paths = {
        "KcBERT 모델 폴더": KC_DIR,
        "DeBERTa 모델 폴더": DEBERTA_DIR,
        "Char-CNN 가중치": CNN_PATH,
        "문자 vocab 파일": VOCAB_PATH,
    }
    missing = [f"{name}: {path}" for name, path in required_paths.items() if not path.exists()]
    if missing:
        raise FileNotFoundError("필수 모델 파일 또는 폴더가 없습니다.\n\n" + "\n".join(missing))

# 모델 로드 핵심 함수
def load_all_models():
    check_model_files()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    with open(VOCAB_PATH, "rb") as f:
        vocab = pickle.load(f)

    # 1. KcBERT
    kc_tokenizer = AutoTokenizer.from_pretrained(str(KC_DIR), local_files_only=True)
    kc_model = AutoModelForSequenceClassification.from_pretrained(str(KC_DIR), local_files_only=True).to(device)
    kc_model.eval()

    # 2. DeBERTa
    deberta_tokenizer = AutoTokenizer.from_pretrained(str(DEBERTA_DIR), local_files_only=True)
    deberta_model = AutoModelForSequenceClassification.from_pretrained(str(DEBERTA_DIR), local_files_only=True).to(device)
    deberta_model.eval()

    # 3. Char-CNN
    cnn_encoder = MultiScaleCharCNNEncoder(vocab_size=len(vocab), embed_dim=128, num_filters=128, filter_sizes=[3, 5, 7])
    cnn_model = SiameseNetwork(cnn_encoder).to(device)
    cnn_model.load_state_dict(torch.load(str(CNN_PATH), map_location=device))
    cnn_model.eval()

    return {
        "device": device,
        "vocab": vocab,
        "kc_tokenizer": kc_tokenizer,
        "kc_model": kc_model,
        "deberta_tokenizer": deberta_tokenizer,
        "deberta_model": deberta_model,
        "cnn_model": cnn_model,
    }

# Char-CNN 전용 토크나이저
def tokenize_char(text, vocab, max_len=256):
    encoded = [vocab.get(char, 1) for char in str(text)]
    encoded = encoded[:max_len] + [0] * max(0, max_len - len(encoded))
    return torch.tensor(encoded[:max_len], dtype=torch.long)

# 최종 앙상블 예측 실행 함수
def predict_ensemble(text_a, text_b, bundle):
    device = bundle["device"]
    vocab = bundle["vocab"]
    kc_tokenizer = bundle["kc_tokenizer"]
    kc_model = bundle["kc_model"]
    deberta_tokenizer = bundle["deberta_tokenizer"]
    deberta_model = bundle["deberta_model"]
    cnn_model = bundle["cnn_model"]

    # 모델 가중치 비율 세팅 (KcBERT 40%, CNN 20%, DeBERTa 40%)
    weights = np.array([0.4, 0.2, 0.4])

    with torch.no_grad():
        # KcBERT 추론
        kc_inputs = kc_tokenizer(text_a, text_b, return_tensors="pt", truncation=True, max_length=128, padding="max_length")
        kc_inputs = {k: v.to(device) for k, v in kc_inputs.items()}
        kc_prob = torch.softmax(kc_model(**kc_inputs).logits, dim=-1).cpu().numpy()[0]

        # Char-CNN 추론
        a_idx = tokenize_char(text_a, vocab).unsqueeze(0).to(device)
        b_idx = tokenize_char(text_b, vocab).unsqueeze(0).to(device)
        cnn_prob = torch.softmax(cnn_model(a_idx, b_idx), dim=-1).cpu().numpy()[0]

        # DeBERTa 추론
        deberta_inputs = deberta_tokenizer(text_a, text_b, return_tensors="pt", truncation=True, max_length=128, padding="max_length")
        deberta_inputs = {k: v.to(device) for k, v in deberta_inputs.items()}
        deberta_prob = torch.softmax(deberta_model(**deberta_inputs).logits, dim=-1).cpu().numpy()[0]

    # 가중 평균 결합
    final_prob = (kc_prob * weights[0]) + (cnn_prob * weights[1]) + (deberta_prob * weights[2])
    pred_label = int(np.argmax(final_prob))

    return {
        "pred_label": pred_label,
        "same_prob": float(final_prob[0]),
        "diff_prob": float(final_prob[1]),
        "kc_prob": kc_prob,
        "cnn_prob": cnn_prob,
        "deberta_prob": deberta_prob,
        "device": str(device),
        "model_dir": str(MODEL_DIR)
    }