import torch
import torch.nn as nn
from transformers import AutoModel

class HierarchicalClassifier(nn.Module):
    def __init__(self, base_model_name, num_labels=2, dropout_prob=0.1):
        super().__init__()

        self.encoder = AutoModel.from_pretrained(base_model_name)

        hidden_size = self.encoder.config.hidden_size

        self.hidden_size = hidden_size
        self.num_labels = num_labels

        self.chunk_attention = nn.Linear(hidden_size, 1)

        self.dropout = nn.Dropout(dropout_prob)

        self.classifier = nn.Linear(hidden_size, num_labels)

    def forward(self, input_ids, attention_mask, chunk_mask):

        B, K, L = input_ids.shape

        flat_input_ids = input_ids.view(B * K, L)
        flat_attention_mask = attention_mask.view(B * K, L)

        outputs = self.encoder(
            flat_input_ids,
            attention_mask=flat_attention_mask
        )

        chunk_cls = outputs.last_hidden_state[:, 0, :]

        chunk_embeds = chunk_cls.view(
            B,
            K,
            self.hidden_size
        )

        attn_scores = self.chunk_attention(
            chunk_embeds
        ).squeeze(-1)

        attn_scores = attn_scores.masked_fill(
            chunk_mask == 0,
            float("-inf")
        )

        attn_weights = torch.softmax(
            attn_scores,
            dim=-1
        )

        doc_embed = (
            chunk_embeds *
            attn_weights.unsqueeze(-1)
        ).sum(dim=1)

        doc_embed = self.dropout(doc_embed)

        logits = self.classifier(doc_embed)

        return logits