import torch
import torch.nn as nn
from transformers import AutoTokenizer, BertModel, BertPreTrainedModel, BertConfig
from pathlib import Path
import logging

# Konfigurasi logging untuk menampilkan pesan error yang jelas
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Definisi pemetaan (bisa juga diletakkan di sini agar terpusat)
ASPECT_MAP = {
    0: 'Fasilitas', 1: 'Harga', 2: 'Hotel', 3: 'Kamar', 4: 'Kolam Renang',
    5: 'Layanan', 6: 'Lokasi', 7: 'Makanan', 8: 'Sarapan', 9: 'Staf'
}
SENTIMENT_MAP = {0: 'Negative', 1: 'Positive'}

# Definisi kelas model, sama seperti sebelumnya
class IndoBERTABSA(BertPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.num_aspect_labels = len(ASPECT_MAP)
        self.num_sentiment_labels = len(SENTIMENT_MAP)
        self.bert = BertModel(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.aspect_classifier = nn.Linear(config.hidden_size, self.num_aspect_labels)
        self.sentiment_classifier = nn.Linear(config.hidden_size, self.num_sentiment_labels)
        self.init_weights()

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None):
        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
        )
        pooled_output = outputs[1]
        pooled_output = self.dropout(pooled_output)
        aspect_logits = self.aspect_classifier(pooled_output)
        sentiment_logits = self.sentiment_classifier(pooled_output)
        return aspect_logits, sentiment_logits

def load_model_and_tokenizer():
    """
    Fungsi terpusat untuk memuat tokenizer dan model.
    Mengembalikan (model, tokenizer) jika berhasil, atau (None, None) jika gagal.
    """
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    logging.info(f"Menggunakan perangkat: {device}")

    try:
        tokenizer = AutoTokenizer.from_pretrained("indobenchmark/indobert-base-p2")
        logging.info("Tokenizer berhasil dimuat.")
    except Exception as e:
        logging.error(f"GAGAL memuat tokenizer: {e}")
        return None, None

    try:
        config = BertConfig.from_pretrained("indobenchmark/indobert-base-p2")
        logging.info("Konfigurasi BERT berhasil dimuat.")
    except Exception as e:
        logging.error(f"GAGAL memuat konfigurasi BERT: {e}")
        return None, None

    model = IndoBERTABSA(config)
    model_save_path = Path('absa_indobert_model_75.pth')

    if not model_save_path.exists():
        logging.error(f"File model tidak ditemukan di '{model_save_path}'")
        return None, None
        
    try:
        model.load_state_dict(torch.load(model_save_path, map_location=device))
        model.to(device)
        model.eval()
        logging.info(f"Model berhasil dimuat dari {model_save_path} dan disetel ke mode evaluasi.")
        return model, tokenizer
    except Exception as e:
        logging.error(f"GAGAL memuat bobot model dari {model_save_path}: {e}")
        return None, None