# Simple Hugging Face compatible wrapper without complex inheritance

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PretrainedConfig
from transformers.modeling_outputs import CausalLMOutputWithPast

# ============================================================
# CONFIG
# ============================================================

class PersadianNanoV4Config:
    def __init__(self):
        self.hidden_size = 512
        self.intermediate_size = 1024
        self.num_hidden_layers = 12
        self.num_attention_heads = 8
        self.num_key_value_heads = 4
        self.num_experts = 4
        self.num_experts_per_tok = 2
        self.progressive_experts = True
        self.min_experts = 1
        self.max_experts = 2
        self.use_hyper_connection = True
        self.use_compressed_attention = True
        self.compress_ratio = 4
        self.vocab_size = 50257
        self.max_position_embeddings = 2048
        self.bos_token_id = 50256
        self.eos_token_id = 50256
        self.rope_theta = 10000.0
        self.dropout = 0.1
        self.layer_norm_eps = 1e-5
        self.attention_dropout = 0.0
        self.use_flash_attention = True
        self.torch_dtype = "float16"

# ============================================================
# ROTARY EMBEDDING
# ============================================================

class RotaryEmbedding(nn.Module):
    def __init__(self, dim, base=10000.0):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)

    def forward(self, x, position_ids):
        batch, seq_len, num_heads, head_dim = x.shape
        freqs = torch.einsum("bi,j->bij", position_ids.float(), self.inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        cos = emb.cos()[:, :, None, :]
        sin = emb.sin()[:, :, None, :]
        x1 = x[..., ::2]
        x2 = x[..., 1::2]
        rotated = torch.stack((-x2, x1), dim=-1).flatten(-2)
        return (x * cos) + (rotated * sin)

# ============================================================
# ADAPTIVE HYPER CONNECTION
# ============================================================

class AdaptiveHyperConnection(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.router = nn.Sequential(
            nn.Linear(hidden_size, hidden_size // 4),
            nn.SiLU(),
            nn.Linear(hidden_size // 4, hidden_size),
            nn.Sigmoid()
        )

    def forward(self, x, residual):
        weight = self.router(x.mean(dim=1, keepdim=True))
        return residual + (x * weight)

# ============================================================
# ATTENTION
# ============================================================

class CompressedSparseAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.head_dim = config.hidden_size // config.num_attention_heads
        
        self.q_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        
        compressed_dim = config.hidden_size // config.compress_ratio
        self.compressor = nn.Linear(config.hidden_size, compressed_dim)
        self.k_proj_compressed = nn.Linear(compressed_dim, self.num_key_value_heads * self.head_dim, bias=False)
        self.v_proj_compressed = nn.Linear(compressed_dim, self.num_key_value_heads * self.head_dim, bias=False)
        
        self.rotary_emb = RotaryEmbedding(self.head_dim, config.rope_theta)

    def forward(self, hidden_states, attention_mask=None, position_ids=None):
        batch_size, seq_len, _ = hidden_states.shape
        q = self.q_proj(hidden_states).view(batch_size, seq_len, self.num_heads, self.head_dim)
        
        if self.config.use_compressed_attention and seq_len > 512:
            compressed = self.compressor(hidden_states)
            k = self.k_proj_compressed(compressed)
            v = self.v_proj_compressed(compressed)
        else:
            k = self.k_proj(hidden_states)
            v = self.v_proj(hidden_states)
        
        k = k.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim)
        v = v.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim)
        
        repeat_factor = self.num_heads // self.num_key_value_heads
        k = k.repeat_interleave(repeat_factor, dim=2)
        v = v.repeat_interleave(repeat_factor, dim=2)
        
        q = self.rotary_emb(q, position_ids)
        k = self.rotary_emb(k, position_ids)
        
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        
        if self.config.use_flash_attention and hasattr(F, "scaled_dot_product_attention"):
            attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask, is_causal=True)
        else:
            scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
            if attention_mask is not None:
                scores = scores + attention_mask
            attn_output = torch.matmul(F.softmax(scores, dim=-1), v)
        
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.hidden_size)
        return self.o_proj(attn_output)

# ============================================================
# MIXTURE OF EXPERTS
# ============================================================

class MixtureOfExperts(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.num_experts = config.num_experts
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(config.hidden_size, config.intermediate_size),
                nn.SiLU(),
                nn.Linear(config.intermediate_size, config.hidden_size),
                nn.Dropout(config.dropout)
            ) for _ in range(config.num_experts)
        ])
        self.router = nn.Linear(config.hidden_size, config.num_experts, bias=False)
        self.aux_loss_coef = 0.01

    def forward(self, hidden_states, progressive_factor=1.0):
        batch_size, seq_len, hidden_size = hidden_states.shape
        hidden_states_flat = hidden_states.view(-1, hidden_size)
        router_logits = self.router(hidden_states_flat)
        routing_weights = F.softmax(router_logits, dim=-1)
        
        if self.config.progressive_experts:
            k = int(round(self.config.min_experts + (self.config.max_experts - self.config.min_experts) * progressive_factor))
            k = max(self.config.min_experts, min(k, self.config.max_experts))
        else:
            k = self.config.num_experts_per_tok
        
        top_k_weights, top_k_indices = torch.topk(routing_weights, k, dim=-1)
        top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True)
        
        final_hidden = torch.zeros_like(hidden_states_flat)
        for expert_idx in range(self.num_experts):
            expert_mask = (top_k_indices == expert_idx).any(dim=-1)
            if expert_mask.any():
                expert_input = hidden_states_flat[expert_mask]
                expert_output = self.experts[expert_idx](expert_input)
                expert_weights = top_k_weights[expert_mask].mean(dim=-1, keepdim=True)
                final_hidden[expert_mask] += expert_output * expert_weights
        
        router_probs = routing_weights.mean(dim=0)
        aux_loss = torch.var(router_probs)
        return final_hidden.view(batch_size, seq_len, hidden_size), aux_loss * self.aux_loss_coef

# ============================================================
# DECODER LAYER
# ============================================================

class PersadianNanoV4DecoderLayer(nn.Module):
    def __init__(self, config, layer_idx):
        super().__init__()
        self.config = config
        self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.self_attn = CompressedSparseAttention(config)
        self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.moe = MixtureOfExperts(config)
        if config.use_hyper_connection:
            self.hc_attention = AdaptiveHyperConnection(config.hidden_size)
            self.hc_moe = AdaptiveHyperConnection(config.hidden_size)

    def forward(self, hidden_states, attention_mask=None, position_ids=None, progressive_factor=1.0):
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        attn_output = self.self_attn(hidden_states, attention_mask, position_ids)
        
        if self.config.use_hyper_connection:
            hidden_states = self.hc_attention(attn_output, residual)
        else:
            hidden_states = residual + attn_output
        
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        moe_output, aux_loss = self.moe(hidden_states, progressive_factor)
        
        if self.config.use_hyper_connection:
            hidden_states = self.hc_moe(moe_output, residual)
        else:
            hidden_states = residual + moe_output
        
        return hidden_states, aux_loss

# ============================================================
# MAIN MODEL
# ============================================================

class PersadianNanoV4Model(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList([PersadianNanoV4DecoderLayer(config, i) for i in range(config.num_hidden_layers)])
        self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.lm_head.weight = self.embed_tokens.weight
        self.apply(self._init_weights)
    
    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
    
    def forward(self, input_ids, attention_mask=None, progressive_factor=1.0):
        batch_size, seq_len = input_ids.shape
        hidden_states = self.embed_tokens(input_ids)
        position_ids = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)
        
        if attention_mask is None:
            mask = torch.triu(torch.full((seq_len, seq_len), float("-inf"), device=input_ids.device), diagonal=1)
            attention_mask = mask.unsqueeze(0).unsqueeze(0)
        
        total_aux_loss = torch.tensor(0.0, device=input_ids.device)
        for layer in self.layers:
            hidden_states, aux_loss = layer(hidden_states, attention_mask, position_ids, progressive_factor)
            total_aux_loss += aux_loss
        
        hidden_states = self.norm(hidden_states)
        logits = self.lm_head(hidden_states)
        return logits, total_aux_loss
    
    def generate(self, input_ids, max_new_tokens=50, temperature=0.7, top_k=50):
        self.eval()
        for _ in range(max_new_tokens):
            logits, _ = self.forward(input_ids)
            next_token_logits = logits[:, -1, :] / temperature
            if top_k > 0:
                values, _ = torch.topk(next_token_logits, top_k)
                min_values = values[:, -1].unsqueeze(-1)
                next_token_logits = torch.where(next_token_logits < min_values, 
                                               torch.full_like(next_token_logits, float("-inf")), 
                                               next_token_logits)
            probs = F.softmax(next_token_logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            input_ids = torch.cat([input_ids, next_token], dim=1)
            if next_token.item() == self.config.eos_token_id:
                break
        return input_ids

# ============================================================
# HF COMPATIBLE WRAPPER (SIMPLE VERSION)
# ============================================================

class PersadianNanoV4ConfigHF(PretrainedConfig):
    model_type = "persadian_nano_v4"
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        for key, value in kwargs.items():
            setattr(self, key, value)

class PersadianNanoV4ForCausalLM(nn.Module):
    """Simple HF-compatible wrapper - no complex inheritance"""
    config_class = PersadianNanoV4ConfigHF
    
    def __init__(self, config):
        super().__init__()
        # Create original config
        original_config = PersadianNanoV4Config()
        for key, value in config.__dict__.items():
            if hasattr(original_config, key):
                setattr(original_config, key, value)
        self.config = original_config
        self.model = PersadianNanoV4Model(original_config)
    
    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        """Load model from Hugging Face hub"""
        from transformers import AutoConfig
        import torch
        
        # Load config
        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True)
        
        # Create model
        model = cls(config)
        
        # Load weights
        import os
        from safetensors.torch import load_file
        
        # Try to load weights
        weight_files = [
            f"{pretrained_model_name_or_path}/pytorch_model.bin",
            f"{pretrained_model_name_or_path}/model.safetensors"
        ]
        
        for weight_file in weight_files:
            if os.path.exists(weight_file):
                if weight_file.endswith('.safetensors'):
                    state_dict = load_file(weight_file)
                else:
                    state_dict = torch.load(weight_file, map_location='cpu')
                model.load_state_dict(state_dict, strict=False)
                break
        
        return model
    
    def forward(self, input_ids, attention_mask=None, labels=None):
        logits, aux_loss = self.model(input_ids, attention_mask)
        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
        return CausalLMOutputWithPast(loss=loss, logits=logits)
    
    def generate(self, input_ids, **kwargs):
        return self.model.generate(input_ids, **kwargs)
    
    def eval(self):
        self.model.eval()
        return self
    
    def to(self, device):
        self.model = self.model.to(device)
        return self