# Self-contained HuggingFace Hub loading shim for NeuralOps BERT (phantom)
import os
from pathlib import Path
import torch
import torch.nn as nn
from transformers import BertConfig, BertModel, BertForSequenceClassification
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download

# Compatibility shim: PyTorch < 2.0 with Accelerate >= 1.0 (named_parameters remove_duplicate)
_orig_named_params = nn.Module.named_parameters
def _compat_named_parameters(self, prefix: str = '', recurse: bool = True, **kwargs):
    return _orig_named_params(self, prefix=prefix, recurse=recurse)
nn.Module.named_parameters = _compat_named_parameters

class NeuralOpsBertConfig(BertConfig):
    model_type = "bert"

def _load_attention_weights_into(model, repo_or_dir):
    p = Path(repo_or_dir)
    if p.is_dir() and (p / "attention_weights.safetensors").exists():
        w_path = str(p / "attention_weights.safetensors")
    else:
        w_path = hf_hub_download(str(repo_or_dir), "attention_weights.safetensors")
    weights = load_file(w_path)
    model.load_state_dict(weights, strict=False)

class NeuralOpsBertModel(BertModel):
    config_class = NeuralOpsBertConfig
    def __init__(self, config, add_pooling_layer=True):
        super().__init__(config, add_pooling_layer=add_pooling_layer)
        from neuralops.api import replace_attention
        replace_attention(self, attention="phantom", copy_weights=False)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        config = kwargs.pop("config", None) or cls.config_class.from_pretrained(pretrained_model_name_or_path)
        base = BertModel.from_pretrained("bert-base-uncased", config=config, **kwargs)
        from neuralops.api import replace_attention
        replace_attention(base, attention="phantom", copy_weights=False)
        _load_attention_weights_into(base, pretrained_model_name_or_path)
        return base

class NeuralOpsBertForSequenceClassification(BertForSequenceClassification):
    config_class = NeuralOpsBertConfig
    def __init__(self, config):
        super().__init__(config)
        from neuralops.api import replace_attention
        replace_attention(self, attention="phantom", copy_weights=False)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        config = kwargs.pop("config", None) or cls.config_class.from_pretrained(pretrained_model_name_or_path)
        base = BertForSequenceClassification.from_pretrained("bert-base-uncased", config=config, **kwargs)
        from neuralops.api import replace_attention
        replace_attention(base, attention="phantom", copy_weights=False)
        _load_attention_weights_into(base, pretrained_model_name_or_path)
        return base