"""ERINE + CRF 模型"""
import torch
import torch.nn as nn
from torchcrf import CRF
from transformers import AutoModel, AutoTokenizer

from config import ERNIE_LOCAL, BIO_LABELS, CHECKPOINT, CHECKPOINT_FROZEN, CHECKPOINT_FC2


class ErnieCRF(nn.Module):
    """ERNIE 3.0 encoder + Linear + CRF"""

    def __init__(self, model_path, num_labels):
        super().__init__()
        self.ernie = AutoModel.from_pretrained(model_path)
        self.fc = nn.Linear(self.ernie.config.hidden_size, num_labels)
        self.crf = CRF(num_labels, batch_first=True)

    def forward(self, input_ids, attention_mask, labels=None):
        mask = attention_mask.bool()
        hidden = self.ernie(input_ids, attention_mask=attention_mask).last_hidden_state
        emissions = self.fc(hidden)
        if labels is not None:
            return -self.crf(emissions, labels, mask=mask, reduction="mean")
        return self.crf.decode(emissions, mask=mask)


class ErnieCRF2(nn.Module):
    """ERNIE 3.0 encoder + 双层FC(hidden→3*hidden→3) + CRF"""

    def __init__(self, model_path, num_labels, hidden_factor=3):
        super().__init__()
        self.ernie = AutoModel.from_pretrained(model_path)
        self.hidden_size = self.ernie.config.hidden_size
        mid_size = self.hidden_size * hidden_factor
        self.fc1 = nn.Linear(self.hidden_size, mid_size)
        self.fc2 = nn.Linear(mid_size, num_labels)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.1)
        self.crf = CRF(num_labels, batch_first=True)

    def forward(self, input_ids, attention_mask, labels=None):
        mask = attention_mask.bool()
        hidden = self.ernie(input_ids, attention_mask=attention_mask).last_hidden_state
        x = self.dropout(self.relu(self.fc1(hidden)))
        emissions = self.fc2(x)
        if labels is not None:
            return -self.crf(emissions, labels, mask=mask, reduction="mean")
        return self.crf.decode(emissions, mask=mask)


def load_model(device="cuda", frozen=False, fc2=False):
    """加载训练好的模型和 tokenizer"""
    tokenizer = AutoTokenizer.from_pretrained(ERNIE_LOCAL)
    if fc2:
        model = ErnieCRF2(ERNIE_LOCAL, len(BIO_LABELS), hidden_factor=2).to(device)
        ckpt = CHECKPOINT_FC2
    else:
        model = ErnieCRF(ERNIE_LOCAL, len(BIO_LABELS)).to(device)
        ckpt = CHECKPOINT_FROZEN if frozen else CHECKPOINT
    model.load_state_dict(torch.load(ckpt, map_location=device, weights_only=True))
    model.eval()
    return model, tokenizer