"""
Frox AI Morph 1.1 — Core Language Model

Improvements over Morph 1.0:
  - 64K vocab (was 32K) for better multilingual + code coverage
  - 16K context window (was 8K) via YaRN RoPE
  - QK-norm enabled throughout
  - Sliding window (even layers) + full attn (odd layers) interleaved
  - Depth-scaled residual init
  - Pre-computed causal mask (cached for reuse)
  - HF GenerationMixin compatible (used by PEFT, vLLM, TRL)
  - save() / from_saved() / from_pretrained() / push_to_hub() helpers
"""
from __future__ import annotations
import json
import math
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint

from config.model_config import MorphConfig, MorphTextConfig
from model.attention.gqa import MorphDecoderLayer, MorphRMSNorm


# ── Output types ──────────────────────────────────────────────────

@dataclass
class MorphModelOutput:
    last_hidden_state: torch.Tensor
    past_key_values:   Optional[Tuple] = None
    hidden_states:     Optional[Tuple] = None
    attentions:        Optional[Tuple] = None


@dataclass
class MorphCausalLMOutput:
    """
    PEFT / HuggingFace Trainer compatible output.
    Dict-style access required by PEFT internals.
    """
    loss:             Optional[torch.Tensor] = None
    logits:           Optional[torch.Tensor] = None
    past_key_values:  Optional[Tuple] = None
    hidden_states:    Optional[Tuple] = None
    attentions:       Optional[Tuple] = None

    def __getitem__(self, key: str):       return getattr(self, key)
    def __setitem__(self, key: str, v):    setattr(self, key, v)
    def get(self, key: str, default=None): return getattr(self, key, default)
    def __contains__(self, key: str):
        return hasattr(self, key) and getattr(self, key) is not None
    def keys(self):
        return [k for k in ("loss","logits","past_key_values","hidden_states","attentions")
                if getattr(self, k, None) is not None]


# ── Backbone ──────────────────────────────────────────────────────

class MorphModel(nn.Module):
    """
    Frox Morph 1.1 transformer backbone.
    Pure decoder — multimodal tokens injected via projectors.
    """

    def __init__(self, config: MorphTextConfig):
        super().__init__()
        self.config = config
        self.padding_idx = config.pad_token_id

        # Embedding table covers vocab + reserved special tokens
        self.embed_tokens = nn.Embedding(
            config.total_vocab_size,
            config.hidden_size,
            padding_idx=self.padding_idx,
        )

        # Transformer layers
        self.layers = nn.ModuleList([
            MorphDecoderLayer(
                hidden_size=config.hidden_size,
                num_heads=config.num_attention_heads,
                num_kv_heads=config.num_key_value_heads,
                head_dim=config.head_dim,
                intermediate_size=config.intermediate_size,
                max_position_embeddings=config.max_position_embeddings,
                rope_theta=config.rope_theta,
                rope_scaling_factor=config.rope_scaling_factor,
                rms_norm_eps=config.rms_norm_eps,
                layer_idx=i,
                qk_norm=config.qk_norm,
                use_sliding_window=config.use_sliding_window,
                sliding_window_size=config.sliding_window_size,
                init_std=config.init_std,
                num_hidden_layers=config.num_hidden_layers,
            )
            for i in range(config.num_hidden_layers)
        ])

        self.norm = MorphRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.gradient_checkpointing = False
        self._init_embeddings()

    def _init_embeddings(self):
        """Standard embedding init. Larger init std for larger vocab."""
        std = self.config.init_std
        nn.init.normal_(self.embed_tokens.weight, mean=0.0, std=std)
        if self.embed_tokens.padding_idx is not None:
            self.embed_tokens.weight.data[self.embed_tokens.padding_idx].zero_()

    def gradient_checkpointing_enable(self, **kwargs):
        self.gradient_checkpointing = True

    def gradient_checkpointing_disable(self):
        self.gradient_checkpointing = False

    def get_input_embeddings(self) -> nn.Embedding:
        return self.embed_tokens

    def set_input_embeddings(self, value: nn.Embedding):
        self.embed_tokens = value

    def forward(
        self,
        input_ids:          Optional[torch.LongTensor]  = None,
        attention_mask:     Optional[torch.Tensor]       = None,
        position_ids:       Optional[torch.LongTensor]  = None,
        past_key_values:    Optional[List[Tuple]]        = None,
        inputs_embeds:      Optional[torch.FloatTensor] = None,
        use_cache:          bool = True,
        output_attentions:  bool = False,
        output_hidden_states: bool = False,
    ) -> MorphModelOutput:

        if self.gradient_checkpointing and self.training:
            use_cache = False  # incompatible with grad checkpointing

        # Embeddings
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        hidden_states = inputs_embeds
        B, S, _ = hidden_states.shape

        # Position IDs
        if position_ids is None:
            past_len = past_key_values[0][0].shape[2] if past_key_values else 0
            position_ids = torch.arange(
                past_len, past_len + S,
                device=hidden_states.device,
            ).unsqueeze(0).expand(B, -1)

        # Build causal mask
        past_len = past_key_values[0][0].shape[2] if past_key_values else 0
        causal_mask = self._build_causal_mask(
            attention_mask, hidden_states.dtype, hidden_states.device, S, past_len
        )

        all_hidden_states = () if output_hidden_states else None
        all_attentions    = () if output_attentions    else None
        next_cache        = () if use_cache            else None

        for i, layer in enumerate(self.layers):
            if output_hidden_states:
                all_hidden_states += (hidden_states,)

            past_kv = past_key_values[i] if past_key_values is not None else None

            if self.gradient_checkpointing and self.training:
                def _ckpt_forward(l):
                    def fn(hs, mask, pos):
                        return l(
                            hidden_states=hs,
                            attention_mask=mask,
                            position_ids=pos,
                            past_key_value=None,
                            use_cache=False,
                            output_attentions=output_attentions,
                        )
                    return fn
                layer_outputs = torch.utils.checkpoint.checkpoint(
                    _ckpt_forward(layer),
                    hidden_states, causal_mask, position_ids,
                    use_reentrant=False,
                )
            else:
                layer_outputs = layer(
                    hidden_states=hidden_states,
                    attention_mask=causal_mask,
                    position_ids=position_ids,
                    past_key_value=past_kv,
                    use_cache=use_cache,
                    output_attentions=output_attentions,
                )

            hidden_states = layer_outputs[0]
            if output_attentions:
                all_attentions += (layer_outputs[1],)
            if use_cache:
                next_cache += (layer_outputs[-1],)

        hidden_states = self.norm(hidden_states)
        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        return MorphModelOutput(
            last_hidden_state=hidden_states,
            past_key_values=next_cache,
            hidden_states=all_hidden_states,
            attentions=all_attentions,
        )

    def _build_causal_mask(
        self,
        attention_mask: Optional[torch.Tensor],
        dtype: torch.dtype,
        device: torch.device,
        seq_len: int,
        past_len: int,
    ) -> Optional[torch.Tensor]:
        """
        4D causal mask [B_or_1, 1, S_q, S_k].
        Never becomes 5D (was the bug in Morph 1.0).
        """
        total_len = seq_len + past_len
        min_val = torch.finfo(dtype).min

        causal = torch.full(
            (seq_len, total_len), fill_value=min_val,
            dtype=dtype, device=device,
        )
        causal = torch.triu(causal, diagonal=past_len + 1)
        causal = causal[None, None, :, :]   # [1, 1, S_q, S_k]

        if attention_mask is not None:
            pad_mask = (1.0 - attention_mask[:, None, None, :].to(dtype)) * min_val
            causal = causal + pad_mask      # [B, 1, S_q, S_k]

        return causal


# ── Causal LM ────────────────────────────────────────────────────

class MorphForCausalLM(nn.Module):
    """
    Frox Morph 1.1 — full causal language model.

    HuggingFace / PEFT / TRL compatible:
      ✓ prepare_inputs_for_generation()
      ✓ can_generate()
      ✓ get/set input/output embeddings
      ✓ gradient_checkpointing_enable/disable
      ✓ forward() accepts return_dict + **kwargs
    """

    def __init__(self, config: MorphTextConfig):
        super().__init__()
        self.config = config
        self.model  = MorphModel(config)

        # LM head: same dim as embedding table
        self.lm_head = nn.Linear(
            config.hidden_size, config.total_vocab_size, bias=False
        )

        # Tie embeddings (saves ~400M params at 64K vocab)
        if config.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight

    # ── Embedding accessors ───────────────────────────────────────

    def get_input_embeddings(self)  -> nn.Embedding: return self.model.embed_tokens
    def set_input_embeddings(self, v): self.model.embed_tokens = v
    def get_output_embeddings(self) -> nn.Linear:    return self.lm_head
    def set_output_embeddings(self, v): self.lm_head = v

    # ── Gradient checkpointing ────────────────────────────────────

    def gradient_checkpointing_enable(self, **kwargs):
        self.model.gradient_checkpointing_enable()

    def gradient_checkpointing_disable(self):
        self.model.gradient_checkpointing_disable()

    # ── HF generation compatibility ───────────────────────────────

    def can_generate(self) -> bool:
        return True

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.LongTensor,
        past_key_values=None,
        attention_mask=None,
        inputs_embeds=None,
        **kwargs,
    ) -> dict:
        if past_key_values is not None:
            input_ids = input_ids[:, -1:]   # only the new token

        model_inputs: dict = {
            "input_ids":       input_ids,
            "past_key_values": past_key_values,
            "use_cache":       kwargs.get("use_cache", True),
            "attention_mask":  attention_mask,
        }
        if inputs_embeds is not None and past_key_values is None:
            model_inputs.pop("input_ids")
            model_inputs["inputs_embeds"] = inputs_embeds

        return model_inputs

    # ── Forward ───────────────────────────────────────────────────

    def forward(
        self,
        input_ids:             Optional[torch.LongTensor]  = None,
        attention_mask:        Optional[torch.Tensor]       = None,
        position_ids:          Optional[torch.LongTensor]  = None,
        past_key_values:       Optional[List[Tuple]]        = None,
        inputs_embeds:         Optional[torch.FloatTensor] = None,
        labels:                Optional[torch.LongTensor]  = None,
        use_cache:             bool = True,
        output_attentions:     bool = False,
        output_hidden_states:  bool = False,
        return_dict:           bool = True,   # PEFT compatibility
        **kwargs,                             # absorb PEFT extra kwargs
    ) -> MorphCausalLMOutput:

        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
        )

        hidden_states = outputs.last_hidden_state
        logits = self.lm_head(hidden_states).float()  # always float32

        loss = None
        if labels is not None:
            # Shift for next-token prediction
            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),
                ignore_index=-100,
            )

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

    # ── Generation (used in edge / demo / inference engine) ──────

    @torch.no_grad()
    def generate(
        self,
        input_ids:        torch.LongTensor,
        attention_mask:   Optional[torch.Tensor] = None,
        max_new_tokens:   int   = 512,
        temperature:      float = 0.7,
        top_p:            float = 0.9,
        top_k:            int   = 50,
        repetition_penalty: float = 1.1,
        eos_token_id:     Optional[int] = None,
        pad_token_id:     Optional[int] = None,
        use_cache:        bool  = True,
        do_sample:        bool  = True,
        stream_callback=None,   # NEW 1.1: optional per-token callback
    ) -> torch.LongTensor:

        eos = eos_token_id if eos_token_id is not None else self.config.eos_token_id
        pad = pad_token_id if pad_token_id is not None else self.config.pad_token_id

        B = input_ids.shape[0]
        generated = input_ids.clone()
        past_key_values = None
        finished = torch.zeros(B, dtype=torch.bool, device=input_ids.device)

        for step in range(max_new_tokens):
            curr_input = generated[:, -1:] if past_key_values is not None else generated

            out = self.forward(
                input_ids=curr_input,
                attention_mask=attention_mask,
                past_key_values=past_key_values,
                use_cache=use_cache,
            )
            logits = out.logits[:, -1, :]
            past_key_values = out.past_key_values

            # Repetition penalty
            if repetition_penalty != 1.0:
                for b in range(B):
                    for tid in set(generated[b].tolist()):
                        if logits[b, tid] < 0:
                            logits[b, tid] *= repetition_penalty
                        else:
                            logits[b, tid] /= repetition_penalty

            if temperature != 1.0:
                logits = logits / temperature

            if top_k > 0:
                top_k_vals, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < top_k_vals[:, -1:]] = float("-inf")

            if do_sample and top_p < 1.0:
                sorted_logits, sorted_idx = torch.sort(logits, descending=True)
                cum_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
                remove = cum_probs - F.softmax(sorted_logits, dim=-1) > top_p
                sorted_logits[remove] = float("-inf")
                logits = torch.zeros_like(logits).scatter_(1, sorted_idx, sorted_logits)

            if do_sample:
                probs = F.softmax(logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)
            else:
                next_token = logits.argmax(dim=-1, keepdim=True)

            next_token = torch.where(
                finished.unsqueeze(-1),
                torch.full_like(next_token, pad),
                next_token,
            )
            generated = torch.cat([generated, next_token], dim=-1)

            if attention_mask is not None:
                attention_mask = torch.cat([
                    attention_mask,
                    torch.ones(B, 1, device=attention_mask.device),
                ], dim=-1)

            finished = finished | (next_token.squeeze(-1) == eos)

            # NEW 1.1: stream callback for real-time output
            if stream_callback is not None:
                for b in range(B):
                    if not finished[b]:
                        stream_callback(b, next_token[b, 0].item(), step)

            if finished.all():
                break

        return generated

    # ── Utilities ─────────────────────────────────────────────────

    def param_count(self) -> dict:
        total     = sum(p.numel() for p in self.parameters())
        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
        return {
            "total":              total,
            "trainable":          trainable,
            "total_billions":     round(total     / 1e9, 3),
            "trainable_billions": round(trainable / 1e9, 3),
        }

    def save(self, path: str):
        from dataclasses import asdict
        p = Path(path)
        p.mkdir(parents=True, exist_ok=True)
        torch.save(self.state_dict(), p / "model.pt")
        with open(p / "config.json", "w") as f:
            json.dump(asdict(self.config), f, indent=2)
        print(f"✓ Morph 1.1 LM saved to {path}")

    @classmethod
    def from_saved(cls, path: str, device: str = "cpu") -> "MorphForCausalLM":
        from utils.common import require_checkpoint_dir
        p = require_checkpoint_dir(path)
        with open(p / "config.json") as f:
            cfg_dict = json.load(f)
        config = MorphTextConfig(**cfg_dict)
        model  = cls(config)
        state  = torch.load(p / "model.pt", map_location=device, weights_only=True)
        model.load_state_dict(state, strict=False)
        return model

    @classmethod
    def from_config(cls, config: MorphTextConfig) -> "MorphForCausalLM":
        """Create model with random weights from config."""
        return cls(config)

    @classmethod
    def from_morph_1_checkpoint(cls, path: str) -> "MorphForCausalLM":
        """
        Load a Morph 1.0 checkpoint into a Morph 1.1 model.
        Handles vocab size mismatch (32K → 64K) by zero-padding embeddings.
        """
        p = Path(path)
        with open(p / "config.json") as f:
            old_cfg = json.load(f)

        # Build 1.1 config with same architecture but upgraded vocab
        new_cfg = MorphTextConfig(**old_cfg)
        new_cfg.vocab_size = 64000
        new_cfg.qk_norm    = True
        new_cfg.version    = "1.1.0"

        model = cls(new_cfg)

        # Load old weights, skip mismatched embed table
        old_state = torch.load(p / "model.pt", map_location="cpu", weights_only=True)
        new_state = model.state_dict()

        for name, param in old_state.items():
            if name not in new_state:
                continue
            if param.shape == new_state[name].shape:
                new_state[name] = param
            elif "embed_tokens" in name or "lm_head" in name:
                # Pad vocabulary dimension: copy old rows, leave new rows at init
                old_rows = param.shape[0]
                new_state[name][:old_rows] = param
                print(f"  Padded {name}: {param.shape} → {new_state[name].shape}")

        model.load_state_dict(new_state)
        print(f"✓ Morph 1.1 loaded from Morph 1.0 checkpoint: {path}")
        return model