"""
DeBERTa-v3 Encoder Wrapper with Instruction Token
Reverse-engineers architectural elements from model.config — no invented values.
Author: Ahmad Ali Parr · Trust: Bel Esprit D'Accord Irrevocable Trust
"""
import torch
from transformers import AutoTokenizer, AutoModel, AutoConfig

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load DeBERTa-v3-base
model_name = "microsoft/deberta-v3-base"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
config = AutoConfig.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name, config=config).to(device)
model.eval()

def print_config_summary(cfg):
    print("\n=== DeBERTa-v3 Configuration Summary ===")
    for key, value in sorted(cfg.to_dict().items()):
        if isinstance(value, list) and len(value) > 10:
            print(f"{key}: <list of length {len(value)}>")
        else:
            print(f"{key}: {value}")
    print("=" * 50)

print_config_summary(config)


class DeBERTaEncoder(torch.nn.Module):
    """Thin wrapper: exposes DeBERTa encoder with BERT-like API."""
    def __init__(self, pretrained_model_name_or_path):
        super().__init__()
        full_model = AutoModel.from_pretrained(pretrained_model_name_or_path)
        self.embeddings = full_model.embeddings
        self.encoder    = full_model.encoder
        self.pooler     = full_model.pooler
        self.config     = full_model.config

    def forward(self, input_ids, attention_mask=None, token_type_ids=None,
                position_ids=None, inputs_embeds=None,
                output_attentions=False, output_hidden_states=False, return_dict=True):
        embedding_output = self.embeddings(
            input_ids=input_ids,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
            past_key_values_length=0,
        )
        encoder_outputs = self.encoder(
            embedding_output,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]
        pooled_output   = self.pooler(sequence_output) if self.pooler else None
        if return_dict:
            return {
                "last_hidden_state": sequence_output,
                "pooler_output": pooled_output,
                "hidden_states": encoder_outputs[1] if output_hidden_states else None,
                "attentions":    encoder_outputs[2] if output_attentions else None,
            }
        return (sequence_output, pooled_output)


encoder = DeBERTaEncoder(model_name).to(device)
encoder.eval()

# Add instruction token
INSTRUCTION_TOKEN = "[INST]"
if INSTRUCTION_TOKEN not in tokenizer.get_vocab():
    tokenizer.add_tokens([INSTRUCTION_TOKEN])
    model.resize_token_embeddings(len(tokenizer))
    encoder.embeddings.word_embeddings = model.get_input_embeddings()

inst_token_id = tokenizer.convert_tokens_to_ids(INSTRUCTION_TOKEN)
print(f"Instruction token id: {inst_token_id}")

# Sanity check
sample = f"{INSTRUCTION_TOKEN} Explain quantum entanglement."
inputs = tokenizer(sample, return_tensors="pt", truncation=True,
                   max_length=128, padding="max_length").to(device)
with torch.no_grad():
    out = encoder(**inputs, return_dict=True)
print("last_hidden_state:", out["last_hidden_state"].shape)