# coding=utf-8
"""
StellarAI: Lightweight Multimodal Large Language Model
Hugging Face compatible implementation
"""

import math
from typing import Optional, Tuple, Dict, Any, List

import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers import PreTrainedModel, GenerationConfig, GenerationMixin
from transformers.modeling_outputs import (
    CausalLMOutputWithPast,
    BaseModelOutputWithPast,
)
from transformers.utils import add_start_docstrings, logging

from configuration_stellarai import StellarAIConfig

logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "StellarAIConfig"
_CHECKPOINT_FOR_DOC = "StellarAI/stellarai-tiny"

STELLARAI_START_DOCSTRING = r"""
    StellarAI Model: A lightweight multimodal large language model.

    Parameters:
        config ([`StellarAIConfig`]): Model configuration class with all the parameters.
"""


# ============================================================
#  RoPE Rotary Position Embedding
# ============================================================
class RoPE(nn.Module):
    """Rotary Position Embedding"""

    def __init__(self, head_dim: int, max_seq_len: int = 2048, base: float = 10000.0):
        super().__init__()
        self.head_dim = head_dim
        self.base = base
        half = head_dim // 2
        inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
        positions = torch.arange(max_seq_len).float()
        angles = torch.outer(positions, inv_freq)
        cos = torch.zeros(max_seq_len, head_dim)
        sin = torch.zeros(max_seq_len, head_dim)
        cos[:, 0::2] = torch.cos(angles)
        cos[:, 1::2] = torch.cos(angles)
        sin[:, 0::2] = torch.sin(angles)
        sin[:, 1::2] = torch.sin(angles)
        self.register_buffer("cos", cos, persistent=False)
        self.register_buffer("sin", sin, persistent=False)

    def forward(self, x: torch.Tensor, offset: int = 0) -> torch.Tensor:
        *_, S, Dh = x.shape
        cos = self.cos[offset:offset + S]
        sin = self.sin[offset:offset + S]
        x1 = x[..., 0::2]
        x2 = x[..., 1::2]
        cos_half = cos[:, 0::2]
        sin_half = sin[:, 0::2]
        while cos_half.dim() < x1.dim():
            cos_half = cos_half.unsqueeze(0)
            sin_half = sin_half.unsqueeze(0)
        out1 = x1 * cos_half - x2 * sin_half
        out2 = x1 * sin_half + x2 * cos_half
        result = torch.empty_like(x)
        result[..., 0::2] = out1
        result[..., 1::2] = out2
        return result


# ============================================================
#  Multi-Head Self Attention
# ============================================================
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1, rope: Optional[RoPE] = None):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        self.scale = 1.0 / math.sqrt(self.head_dim)
        self.Wq = nn.Linear(d_model, d_model, bias=True)
        self.Wk = nn.Linear(d_model, d_model, bias=True)
        self.Wv = nn.Linear(d_model, d_model, bias=True)
        self.Wo = nn.Linear(d_model, d_model, bias=True)
        self.dropout = nn.Dropout(dropout)
        self.rope = rope

    def forward(
        self, x: torch.Tensor, mask: Optional[torch.Tensor] = None,
        kv_cache: Optional[Dict] = None,
    ) -> Tuple[torch.Tensor, Optional[Dict]]:
        B, S, D = x.shape
        Q = self.Wq(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
        K = self.Wk(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
        V = self.Wv(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)

        offset = 0
        if kv_cache is not None and "K" in kv_cache:
            offset = kv_cache["K"].shape[2]
        if self.rope is not None:
            Q = self.rope(Q.transpose(1, 2), offset=offset).transpose(1, 2)
            K = self.rope(K.transpose(1, 2), offset=offset).transpose(1, 2)

        if kv_cache is not None:
            if "K" in kv_cache:
                K = torch.cat([kv_cache["K"], K], dim=2)
                V = torch.cat([kv_cache["V"], V], dim=2)
            kv_cache["K"] = K
            kv_cache["V"] = V

        scores = torch.matmul(Q, K.transpose(-2, -1)) * self.scale
        Sq = Q.shape[2]
        Sk = K.shape[2]
        if mask is not None:
            scores = scores.masked_fill(mask[:Sq, :Sk], float('-inf'))
        else:
            causal = torch.triu(torch.ones(Sq, Sk, device=x.device, dtype=torch.bool), diagonal=1)
            scores = scores.masked_fill(causal, float('-inf'))

        attn = F.softmax(scores, dim=-1)
        attn = self.dropout(attn)
        out = torch.matmul(attn, V)
        out = out.transpose(1, 2).contiguous().view(B, S, D)
        return self.Wo(out), kv_cache


# ============================================================
#  Cross Attention
# ============================================================
class CrossAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        self.scale = 1.0 / math.sqrt(self.head_dim)
        self.Wq = nn.Linear(d_model, d_model, bias=True)
        self.Wk = nn.Linear(d_model, d_model, bias=True)
        self.Wv = nn.Linear(d_model, d_model, bias=True)
        self.Wo = nn.Linear(d_model, d_model, bias=True)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, ctx: torch.Tensor) -> torch.Tensor:
        B, Sx, D = x.shape
        Sc = ctx.shape[1]
        Q = self.Wq(x).view(B, Sx, self.num_heads, self.head_dim).transpose(1, 2)
        K = self.Wk(ctx).view(B, Sc, self.num_heads, self.head_dim).transpose(1, 2)
        V = self.Wv(ctx).view(B, Sc, self.num_heads, self.head_dim).transpose(1, 2)
        scores = torch.matmul(Q, K.transpose(-2, -1)) * self.scale
        attn = F.softmax(scores, dim=-1)
        attn = self.dropout(attn)
        out = torch.matmul(attn, V)
        out = out.transpose(1, 2).contiguous().view(B, Sx, D)
        return self.Wo(out)


# ============================================================
#  Transformer Block (Pre-LN)
# ============================================================
class TransformerBlock(nn.Module):
    def __init__(self, d_model: int, num_heads: int, ff_dim: int,
                 dropout: float = 0.1, rope: Optional[RoPE] = None, eps: float = 1e-6):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, num_heads, dropout, rope)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, ff_dim),
            nn.GELU(),
            nn.Linear(ff_dim, d_model),
        )
        self.ln1 = nn.LayerNorm(d_model, eps=eps)
        self.ln2 = nn.LayerNorm(d_model, eps=eps)
        self.dropout = nn.Dropout(dropout)

    def forward(
        self, x: torch.Tensor, mask: Optional[torch.Tensor] = None,
        kv_cache: Optional[Dict] = None,
    ) -> Tuple[torch.Tensor, Optional[Dict]]:
        attn_out, kv_cache = self.attn(self.ln1(x), mask, kv_cache)
        x = x + self.dropout(attn_out)
        x = x + self.dropout(self.ffn(self.ln2(x)))
        return x, kv_cache


# ============================================================
#  Fusion Block (Self + Cross Attention)
# ============================================================
class FusionBlock(nn.Module):
    def __init__(self, d_model: int, num_heads: int, ff_dim: int,
                 dropout: float = 0.1, rope: Optional[RoPE] = None, eps: float = 1e-6):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout, rope)
        self.cross_attn = CrossAttention(d_model, num_heads, dropout)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, ff_dim),
            nn.GELU(),
            nn.Linear(ff_dim, d_model),
        )
        self.ln1 = nn.LayerNorm(d_model, eps=eps)
        self.ln2 = nn.LayerNorm(d_model, eps=eps)
        self.ln3 = nn.LayerNorm(d_model, eps=eps)
        self.dropout = nn.Dropout(dropout)

    def forward(
        self, x: torch.Tensor, ctx: Optional[torch.Tensor],
        mask: Optional[torch.Tensor] = None, kv_cache: Optional[Dict] = None,
    ) -> Tuple[torch.Tensor, Optional[Dict]]:
        attn_out, kv_cache = self.self_attn(self.ln1(x), mask, kv_cache)
        x = x + self.dropout(attn_out)
        if ctx is not None:
            x = x + self.dropout(self.cross_attn(self.ln2(x), ctx))
        x = x + self.dropout(self.ffn(self.ln3(x)))
        return x, kv_cache


# ============================================================
#  Vision Encoder (CNN + ViT Hybrid)
# ============================================================
class VisionEncoder(nn.Module):
    def __init__(self, cfg: StellarAIConfig):
        super().__init__()
        self.cfg = cfg
        vc = cfg.vision_cfg
        D = cfg.d_model

        layers = []
        c_in = vc["num_channels"]
        for c_out in vc["cnn_channels"]:
            layers.append(nn.Conv2d(c_in, c_out, 3, stride=2, padding=1))
            layers.append(nn.BatchNorm2d(c_out))
            layers.append(nn.SiLU())
            c_in = c_out
        self.cnn = nn.Sequential(*layers)
        self.cnn_proj = nn.Linear(c_in, D)

        self.pos_emb = nn.Parameter(torch.randn(vc["vision_num_patches"], D) * 0.02)

        self.blocks = nn.ModuleList([
            TransformerBlock(D, vc["vision_num_heads"], vc["vision_ff_dim"],
                             cfg.dropout, rope=None, eps=cfg.layer_norm_eps)
            for _ in range(vc["vision_num_layers"])
        ])
        self.final_ln = nn.LayerNorm(D, eps=cfg.layer_norm_eps)

    def forward(self, images: torch.Tensor) -> torch.Tensor:
        x = self.cnn(images)
        B, C, Hp, Wp = x.shape
        target = int(math.sqrt(self.cfg.vision_cfg["vision_num_patches"]))
        if Hp != target:
            x = F.adaptive_avg_pool2d(x, (target, target))
        x = x.flatten(2).transpose(1, 2)
        x = self.cnn_proj(x)
        x = x + self.pos_emb.unsqueeze(0)
        for blk in self.blocks:
            x, _ = blk(x, mask=None)
        return self.final_ln(x)


# ============================================================
#  StellarAI Core Model
# ============================================================
@add_start_docstrings(
    "The bare StellarAI Model outputting raw hidden-states.",
    STELLARAI_START_DOCSTRING,
)
class StellarAIModel(PreTrainedModel):
    config_class = StellarAIConfig
    base_model_prefix = "model"
    _no_split_modules = ["TransformerBlock", "FusionBlock"]

    def __init__(self, config: StellarAIConfig):
        super().__init__(config)
        self.config = config
        D = config.d_model
        fc = config.mm_fusion_cfg

        self.wte = nn.Embedding(config.vocab_size, D)
        self.text_rope = RoPE(D // config.num_attention_heads,
                              config.max_position_embeddings, config.rope_theta)

        self.text_blocks = nn.ModuleList([
            TransformerBlock(D, config.num_attention_heads, config.intermediate_size,
                             config.dropout, rope=self.text_rope, eps=config.layer_norm_eps)
            for _ in range(config.num_hidden_layers)
        ])
        self.text_final_ln = nn.LayerNorm(D, eps=config.layer_norm_eps)

        self.vision_encoder = VisionEncoder(config)

        self.fusion_blocks = nn.ModuleList([
            FusionBlock(D, fc["fusion_num_heads"], fc["fusion_ff_dim"],
                        config.dropout, rope=self.text_rope, eps=config.layer_norm_eps)
            for _ in range(fc["fusion_num_layers"])
        ])
        self.fusion_final_ln = nn.LayerNorm(D, eps=config.layer_norm_eps)

        self.gradient_checkpointing = False
        self.post_init()

    def get_input_embeddings(self):
        return self.wte

    def set_input_embeddings(self, value):
        self.wte = value

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        images: Optional[torch.FloatTensor] = None,
        image_positions: Optional[List] = None,
        use_cache: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> BaseModelOutputWithPast:
        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("Cannot specify both input_ids and inputs_embeds")
        if input_ids is None and inputs_embeds is None:
            raise ValueError("Must specify either input_ids or inputs_embeds")

        if inputs_embeds is None:
            inputs_embeds = self.wte(input_ids)

        B, S, D = inputs_embeds.shape
        device = inputs_embeds.device

        # Attention mask -> causal mask
        if attention_mask is not None and attention_mask.dim() == 2:
            mask = ~(attention_mask[:, None, None, :].bool())
        else:
            mask = None

        # Vision encoding
        vision_features = None
        if images is not None:
            vision_features = self.vision_encoder(images)
            if image_positions is not None:
                for b in range(B):
                    if b < len(image_positions):
                        s, e = image_positions[b]
                        n = min(e - s, vision_features.shape[1])
                        inputs_embeds[b, s:s + n] = vision_features[b, :n]

        # Text Transformer
        # HF may pass DynamicCache (not subscriptable). Our layers use dict-per-layer format.
        # Safest: if use_cache is requested but past_key_values is not a plain sequence of
        # layer-level cache dicts, fall back to no-cache by treating past as all-None and
        # still producing present_kv (the dict-style) so the next iteration can consume it.
        if past_key_values is not None:
            # Accept both tuple[dict, ...] and list[dict, ...]; reject DynamicCache & friends.
            if not isinstance(past_key_values, (tuple, list)) or not all(
                (isinstance(x, dict) or x is None) for x in past_key_values
            ):
                past_key_values = None

        x = inputs_embeds
        all_hidden_states = () if output_hidden_states else None
        present_kv = [] if use_cache else None

        past = past_key_values if past_key_values is not None else [None] * (
            self.config.num_hidden_layers + self.config.mm_fusion_cfg["fusion_num_layers"]
        )

        for i, blk in enumerate(self.text_blocks):
            if output_hidden_states:
                all_hidden_states += (x,)
            kv = past[i] if use_cache else None
            if use_cache and kv is None:
                kv = {}
            x, kv_out = blk(x, mask, kv_cache=kv)
            if use_cache:
                present_kv.append(kv_out)

        x = self.text_final_ln(x)

        # Fusion layers
        offset = self.config.num_hidden_layers
        ctx = vision_features
        for i, fblk in enumerate(self.fusion_blocks):
            if output_hidden_states:
                all_hidden_states += (x,)
            kv = past[offset + i] if use_cache else None
            if use_cache and kv is None:
                kv = {}
            x, kv_out = fblk(x, ctx, mask, kv_cache=kv)
            if use_cache:
                present_kv.append(kv_out)

        x = self.fusion_final_ln(x)

        if output_hidden_states:
            all_hidden_states += (x,)

        if not return_dict:
            return tuple(v for v in [x, tuple(present_kv) if use_cache else None, all_hidden_states, None] if v is not None)

        return BaseModelOutputWithPast(
            last_hidden_state=x,
            past_key_values=tuple(present_kv) if use_cache else None,
            hidden_states=all_hidden_states,
            attentions=None,
        )


# ============================================================
#  StellarAI for Causal Language Modeling
# ============================================================
class StellarAIForCausalLM(PreTrainedModel, GenerationMixin):
    r"""
    StellarAI Model with a language modeling head on top (for causal LM).
    """
    config_class = StellarAIConfig
    base_model_prefix = "model"
    _no_split_modules = ["TransformerBlock", "FusionBlock"]
    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config: StellarAIConfig):
        super().__init__(config)
        self.model = StellarAIModel(config)
        D = config.d_model
        self.lm_bias = nn.Parameter(torch.zeros(config.vocab_size)) if config.lm_head_bias else None
        self.post_init()

    def get_output_embeddings(self):
        # Weight tying: output shares embedding weights
        return self.model.wte

    def set_output_embeddings(self, new_embeddings):
        self.model.wte = new_embeddings

    def get_input_embeddings(self):
        return self.model.wte

    def set_input_embeddings(self, value):
        self.model.wte = value

    def tie_weights(self, recompute_mapping: bool = False, missing_keys=None, **kwargs):
        """Weight tying: LM head shares embedding weights"""
        if self.config.use_weight_tying:
            # weight is implicitly shared since we use F.linear(model.wte.weight)
            pass

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        images: Optional[torch.FloatTensor] = None,
        image_positions: Optional[List] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            images=images,
            image_positions=image_positions,
            use_cache=use_cache,
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
            return_dict=return_dict,
        )

        hidden_states = outputs[0]
        logits = F.linear(hidden_states, self.model.wte.weight)
        if self.lm_bias is not None:
            logits = logits + self.lm_bias

        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, self.config.vocab_size),
                shift_labels.view(-1),
                ignore_index=-100,
            )

        if not return_dict:
            output = (logits,) + outputs[1:]
            return ((loss,) + output) if loss is not None else output

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    @torch.no_grad()
    def generate_text(
        self,
        prompt: str,
        tokenizer=None,
        max_new_tokens: int = 100,
        temperature: float = 0.7,
        top_k: int = 40,
        device: torch.device = None,
    ) -> str:
        """Convenience method for text generation."""
        if tokenizer is None:
            raise ValueError("tokenizer is required for generate_text()")
        if device is None:
            device = next(self.parameters()).device
        self.eval()

        ids = tokenizer.encode(prompt)
        # Ensure [BOS] at start if using SimpleTokenizer
        if hasattr(tokenizer, 'bos_id') and ids[0] != tokenizer.bos_id:
            ids = [tokenizer.bos_id] + ids
        input_ids = torch.tensor([ids], dtype=torch.long, device=device)
        eos_id = tokenizer.eos_id if hasattr(tokenizer, 'eos_id') else 2

        for _ in range(max_new_tokens):
            if input_ids.shape[1] >= self.config.max_position_embeddings:
                break
            out = self.forward(input_ids)
            logits = out.logits[:, -1, :self.config.vocab_size]
            if temperature <= 0:
                next_id = logits.argmax(dim=-1, keepdim=True)
            else:
                logits = logits / temperature
                v, _ = torch.topk(logits, min(top_k, logits.shape[-1]), dim=-1)
                logits[logits < v[:, [-1]]] = float('-inf')
                probs = F.softmax(logits, dim=-1)
                next_id = torch.multinomial(probs, num_samples=1)
            input_ids = torch.cat([input_ids, next_id], dim=1)
            if next_id.item() == eos_id:
                break

        result_ids = input_ids[0].tolist()
        return tokenizer.decode(result_ids, skip_special=True) if hasattr(tokenizer, 'decode') else str(result_ids)

    # ===== HuggingFace GenerationMixin required hooks =====
    def prepare_inputs_for_generation(
        self,
        input_ids,
        past_key_values=None,
        attention_mask=None,
        inputs_embeds=None,
        images=None,
        image_positions=None,
        **kwargs,
    ):
        """Transform generation-time inputs into forward() arguments. Handles KV cache truncation."""
        if past_key_values is not None:
            # Only pass the last token as input; KV holds earlier context
            past_length = 0
            try:
                # past_key_values is tuple of kv_cache dicts, each has K tensor shape (B, H, S, Dh)
                first_cache = past_key_values[0]
                if isinstance(first_cache, dict) and "K" in first_cache:
                    past_length = first_cache["K"].shape[2]
            except Exception:
                past_length = 0
            input_ids = input_ids[:, past_length:]

        return {
            "input_ids": input_ids if inputs_embeds is None else None,
            "inputs_embeds": inputs_embeds,
            "past_key_values": past_key_values,
            "attention_mask": attention_mask,
            "images": images,
            "image_positions": image_positions,
            "use_cache": False,
        }

    @staticmethod
    def _reorder_cache(past_key_values, beam_idx):
        """Reorder KV cache entries for beam search (no-op for greedy / sampling single-beam)."""
        if past_key_values is None:
            return None
        reordered = []
        for layer_cache in past_key_values:
            if isinstance(layer_cache, dict):
                new_cache = {}
                for k, v in layer_cache.items():
                    if isinstance(v, torch.Tensor):
                        new_cache[k] = v.index_select(0, beam_idx.to(v.device))
                    else:
                        new_cache[k] = v
                reordered.append(new_cache)
            else:
                reordered.append(layer_cache)
        return tuple(reordered)