"""
Stage 3 — Classifier model.

Architecture (JQL §4.3):
  - Encoder: Snowflake/snowflake-arctic-embed-m-v2.0 (305M params, frozen)
  - Head: Linear(hidden_size, hidden_size) → ReLU → Linear(hidden_size, 1)
  - Input: [CLS] embedding (first token of last_hidden_state)
  - Output: scalar regression score (interpreted as 0–5 range)

The encoder is fully frozen; only the head is trained.
"""

from pathlib import Path

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

# ── Backbone (pinned; do not change without spec update) ──────────────────────
BACKBONE = "Snowflake/snowflake-arctic-embed-m-v2.0"
HIDDEN_SIZE = 768  # arctic-embed-m-v2.0 hidden size; also the head's input/hidden width


def build_head(hidden_size: int = HIDDEN_SIZE) -> nn.Sequential:
    """MLP regression head on the [CLS] embedding (JQL §4.3): Linear → ReLU → Linear(1)."""
    return nn.Sequential(
        nn.Linear(hidden_size, hidden_size),
        nn.ReLU(),
        nn.Linear(hidden_size, 1),
    )


def load_head_only(checkpoint_path: str | Path, hidden_size: int = HIDDEN_SIZE) -> nn.Sequential:
    """
    Load just the regression head from a full MarketingClassifier checkpoint.

    Skips downloading/instantiating the frozen 305M-param encoder — use this
    when only scoring pre-computed [CLS] embeddings (eval_compare.py, eval_ensemble.py).
    """
    head = build_head(hidden_size)
    state = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
    head_state = {k[len("head."):]: v for k, v in state.items() if k.startswith("head.")}
    head.load_state_dict(head_state)
    return head


class MarketingClassifier(nn.Module):
    def __init__(self, backbone: str = BACKBONE):
        super().__init__()
        # arctic-embed-m-v2.0's custom modeling defaults to use_memory_efficient_attention=True
        # (requires xformers).  Disable that and route to the model's built-in GteSdpaAttention
        # (uses torch.nn.functional.scaled_dot_product_attention — flash-attention on A100,
        # O(N) memory vs O(N²) for eager).  Requires transformers 4.46.x; tested on 4.46.3.
        config = AutoConfig.from_pretrained(backbone, trust_remote_code=True)
        config.use_memory_efficient_attention = False
        self.encoder = AutoModel.from_pretrained(
            backbone,
            config=config,
            add_pooling_layer=False,
            trust_remote_code=True,
            attn_implementation="sdpa",
        )
        hidden_size = self.encoder.config.hidden_size
        for param in self.encoder.parameters():
            param.requires_grad = False
        # MLP regression head on [CLS] embedding (JQL §4.3)
        self.head = build_head(hidden_size)

    def encode(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
    ) -> torch.Tensor:
        """Returns (batch_size, hidden_size) [CLS] embeddings (no head)."""
        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        return outputs.last_hidden_state[:, 0]  # (B, H)

    def score_from_embedding(self, embedding: torch.Tensor) -> torch.Tensor:
        """Returns (batch_size,) scores from pre-computed [CLS] embeddings."""
        return self.head(embedding).squeeze(-1)  # (B,)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
    ) -> torch.Tensor:
        """Returns (batch_size,) scalar predictions."""
        outputs = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
        )
        cls_emb = outputs.last_hidden_state[:, 0]  # (B, H)
        return self.head(cls_emb).squeeze(-1)  # (B,)