"""Self-contained definition of the hybrid RoBERTa + linguistic-features model.

The released checkpoint ``hybrid_model_best.pt`` is a plain PyTorch
``state_dict`` for the ``HybridClassifier`` defined here. Import this module to
rebuild the architecture and load the weights.
"""
import torch
import torch.nn as nn
from transformers import RobertaModel

CONFIG = {
    "roberta_model": "roberta-base",
    "roberta_dim": 768,
    "ling_dim": 25,
    "ling_hidden": 256,
    "hidden_dim": 768,
    "dropout": 0.15,
    "freeze_bottom_layers": 6,
}

LING_FEATURE_NAMES = [
    "msttr", "avg_word_len", "hapax_ratio", "function_ratio", "punct_density",
    "char_entropy", "burstiness", "repetition_ratio", "avg_sent_len",
    "sent_len_std", "noun_ratio", "verb_ratio", "adj_ratio", "adv_ratio",
    "pron_ratio", "pos_diversity", "avg_tree_depth", "max_tree_depth",
    "sub_clause_ratio", "dm_density", "sent_len_cv", "fp_ratio",
    "num_sentences", "words_per_sent", "perplexity",
]


class FeatureAttention(nn.Module):
    def __init__(self, num_features):
        super().__init__()
        self.attention = nn.Sequential(
            nn.Linear(num_features, num_features),
            nn.Tanh(),
            nn.Linear(num_features, num_features),
            nn.Softmax(dim=-1),
        )

    def forward(self, features):
        attn_weights = self.attention(features)
        weighted_features = features * attn_weights
        return weighted_features, attn_weights


class GatedFusion(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.gate = nn.Sequential(
            nn.Linear(dim * 2, dim),
            nn.Sigmoid(),
        )

    def forward(self, cls_emb, ling_emb):
        combined = torch.cat([cls_emb, ling_emb], dim=-1)
        g = self.gate(combined)
        fused = g * cls_emb + (1 - g) * ling_emb
        return fused, g


class HybridClassifier(nn.Module):
    """RoBERTa CLS embedding gated-fused with attended linguistic features."""

    def __init__(self, config=CONFIG):
        super().__init__()
        self.roberta = RobertaModel.from_pretrained(config["roberta_model"])
        if config["freeze_bottom_layers"] > 0:
            for param in self.roberta.embeddings.parameters():
                param.requires_grad = False
            for i in range(config["freeze_bottom_layers"]):
                for param in self.roberta.encoder.layer[i].parameters():
                    param.requires_grad = False
        self.feature_attention = FeatureAttention(config["ling_dim"])
        self.ling_projection = nn.Sequential(
            nn.Linear(config["ling_dim"], config["ling_hidden"]),
            nn.LayerNorm(config["ling_hidden"]),
            nn.GELU(),
            nn.Dropout(config["dropout"]),
            nn.Linear(config["ling_hidden"], config["roberta_dim"]),
            nn.LayerNorm(config["roberta_dim"]),
            nn.GELU(),
            nn.Dropout(config["dropout"]),
        )
        self.gated_fusion = GatedFusion(config["roberta_dim"])
        self.classifier = nn.Sequential(
            nn.Linear(config["roberta_dim"], config["hidden_dim"]),
            nn.GELU(),
            nn.Dropout(config["dropout"]),
            nn.Linear(config["hidden_dim"], 2),
        )
        self.dropout = nn.Dropout(config["dropout"])
        self.config = config

    def forward(self, input_ids, attention_mask, ling_features, return_gate=False):
        outputs = self.roberta(input_ids=input_ids, attention_mask=attention_mask)
        cls_embedding = outputs.last_hidden_state[:, 0, :]
        cls_embedding = self.dropout(cls_embedding)
        attended_features, attn_weights = self.feature_attention(ling_features)
        ling_proj = self.ling_projection(attended_features)
        fused, gate_values = self.gated_fusion(cls_embedding, ling_proj)
        logits = self.classifier(fused)
        if return_gate:
            return logits, gate_values, attn_weights
        return logits


def load_model(checkpoint_path="hybrid_model_best.pt", device="cpu"):
    """Build the model and load the released weights. Returns an eval-mode model."""
    model = HybridClassifier(CONFIG)
    state_dict = torch.load(checkpoint_path, map_location=device)
    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()
    return model