"""HuggingFace wrapper for the NucEngram genomic language model.

A char-level DNA masked language model (A/C/G/T/N, 9-token vocab) built on a
ModernBERT encoder with an n-gram memory module injected into two early layers.
`AutoModel.from_pretrained(..., trust_remote_code=True)` returns the encoder;
`last_hidden_state` is a per-nucleotide embedding you can pool + probe/fine-tune.
"""
from __future__ import annotations
import numpy as np
import torch
import torch.nn as nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import BaseModelOutput, MaskedLMOutput

from .configuration_nucengram import NucEngramConfig
from .engram import EngramConfig
from .model_modernbert import ModernBertGenomicConfig
from .model_modernbert_engram import ModernBertEngramWrapper
from .tokenizer import encode as _encode, PAD_ID

_VOCAB = {i: c for i, c in zip(range(9), "_XXXACGTN")}  # display only


class NucEngramPreTrainedModel(PreTrainedModel):
    config_class = NucEngramConfig
    base_model_prefix = "net"
    _no_split_modules = ["ModernBertEncoderLayer", "Engram"]

    def _init_weights(self, module):
        pass  # weights come from the pretrained checkpoint

    # --- tokenisation helper (the tokenizer is a fixed char->id map) ---
    @staticmethod
    def encode(sequences, max_length=8192, device="cpu"):
        """str | list[str] -> (input_ids, attention_mask) LongTensors."""
        if isinstance(sequences, str):
            sequences = [sequences]
        B = len(sequences)
        L = min(max(len(s) for s in sequences), max_length)
        ids = np.full((B, L), PAD_ID, dtype=np.int64)
        for i, s in enumerate(sequences):
            e = _encode(s).astype(np.int64)[:L]
            ids[i, :len(e)] = e
        ids = torch.from_numpy(ids)
        return ids.to(device), (ids != PAD_ID).long().to(device)


class NucEngramModel(NucEngramPreTrainedModel):
    """Encoder — returns per-token `last_hidden_state` (use for embeddings / FT)."""

    def __init__(self, config: NucEngramConfig):
        super().__init__(config)
        mb = ModernBertGenomicConfig(**config.backbone)
        eng = EngramConfig(**config.engram)
        self.net = ModernBertEngramWrapper(mb, eng)
        self.post_init()

    def forward(self, input_ids, attention_mask=None, **kwargs):
        if attention_mask is None:
            attention_mask = (input_ids != PAD_ID).long()
        # the engram hooks read the raw ids off the wrapper
        self.net._current_input_ids = input_ids
        out = self.net.model.model(input_ids=input_ids, attention_mask=attention_mask)
        return BaseModelOutput(last_hidden_state=out.last_hidden_state)

    @torch.no_grad()
    def embed(self, sequences, pooling="mean", max_length=8192):
        """Convenience: sequence(s) -> pooled embedding [B, hidden]."""
        self.eval()
        dev = next(self.parameters()).device
        ids, am = self.encode(sequences, max_length=max_length, device=dev)
        h = self.forward(ids, am).last_hidden_state.float()
        if pooling == "cls":
            return h[:, 0]
        m = am.float().unsqueeze(-1)
        return (h * m).sum(1) / m.sum(1).clamp_min(1.0)


class NucEngramForMaskedLM(NucEngramPreTrainedModel):
    """Same weights, exposing the MLM head (returns `logits` over the 9-token vocab)."""

    def __init__(self, config: NucEngramConfig):
        super().__init__(config)
        mb = ModernBertGenomicConfig(**config.backbone)
        eng = EngramConfig(**config.engram)
        self.net = ModernBertEngramWrapper(mb, eng)
        self.post_init()

    def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
        if attention_mask is None:
            attention_mask = (input_ids != PAD_ID).long()
        self.net._current_input_ids = input_ids
        out = self.net.model(input_ids=input_ids, attention_mask=attention_mask)
        logits = out.logits if hasattr(out, "logits") else out
        loss = None
        if labels is not None:
            loss = nn.functional.cross_entropy(
                logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index=-100)
        return MaskedLMOutput(loss=loss, logits=logits)