"""Task heads for StrataBERT."""

from __future__ import annotations

import torch
from torch import nn


class StrataBertClassificationHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dropout = nn.Dropout(config.classifier_dropout)
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.out_proj = nn.Linear(config.hidden_size, config.num_labels)

    def forward(self, pooled: torch.Tensor) -> torch.Tensor:
        x = self.dropout(pooled)
        x = torch.nn.functional.gelu(self.dense(x))
        x = self.dropout(x)
        return self.out_proj(x)


class StrataBertTokenClassificationHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dropout = nn.Dropout(config.classifier_dropout)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.classifier(self.dropout(hidden_states))


class StrataBertLMHead(nn.Module):
    def __init__(self, config, embedding_weight: torch.nn.Parameter | None = None):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps)
        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
        self.decoder.bias = self.bias
        if embedding_weight is not None:
            self.decoder.weight = embedding_weight

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        x = torch.nn.functional.gelu(self.dense(hidden_states))
        x = self.norm(x)
        return self.decoder(x)


class StrataBertRTDHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dropout = nn.Dropout(config.classifier_dropout)
        self.classifier = nn.Linear(config.hidden_size, 1)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.classifier(self.dropout(hidden_states)).squeeze(-1)