"""The 6.53M-parameter head; BGE-M3 is not embedded in these weights."""

from dataclasses import dataclass
from typing import Optional

import torch
from torch import nn
from transformers import PreTrainedModel
from transformers.utils import ModelOutput

from .configuration_scripture_genome import ScriptureGenomeConfig


@dataclass
class ScriptureGenomeOutput(ModelOutput):
    genome: Optional[torch.Tensor] = None


class ScriptureGenomeHeadModel(PreTrainedModel):
    config_class = ScriptureGenomeConfig
    base_model_prefix = "head"
    main_input_name = "embeddings"

    def __init__(self, config):
        super().__init__(config)
        self.head = nn.Sequential(
            nn.Linear(config.input_dim, config.hidden_dim),
            nn.GELU(),
            nn.Linear(config.hidden_dim, config.output_dim),
            nn.Sigmoid(),
        )
        self.post_init()

    def forward(self, embeddings, return_dict=None):
        """Accept normalized BGE-M3 CLS embeddings [batch,1024], not token IDs."""
        if embeddings.ndim != 2 or embeddings.shape[-1] != self.config.input_dim:
            raise ValueError("embeddings must have shape [batch, 1024]")
        if not embeddings.is_floating_point() or not torch.isfinite(embeddings).all():
            raise ValueError("embeddings must be finite floating-point values")
        genome = self.head(embeddings)
        use_dict = self.config.use_return_dict if return_dict is None else return_dict
        return ScriptureGenomeOutput(genome=genome) if use_dict else (genome,)


ScriptureGenomeHeadModel.register_for_auto_class("AutoModel")