import math
from types import SimpleNamespace
import json
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint as grad_checkpoint


def rotate_every_two(x):
    x1 = x[..., ::2]
    x2 = x[..., 1::2]
    return torch.stack((-x2, x1), dim=-1).reshape_as(x)


def apply_rotary_pos_emb(q, k, sin, cos):
    # q,k: (B, nh, T, hs)
    q_ = (q * cos) + (rotate_every_two(q) * sin)
    k_ = (k * cos) + (rotate_every_two(k) * sin)
    return q_, k_


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

    def forward(self, seq_len, device):
        t = torch.arange(seq_len, device=device).type_as(self.inv_freq)
        freqs = torch.einsum('i , j -> i j', t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)  # (T, dim)
        sin = emb.sin()[None, None, :, :]
        cos = emb.cos()[None, None, :, :]
        return sin, cos


class RMSNorm(nn.Module):
    """Simple RMSNorm implementation compatible with HF's RMSNorm behavior."""
    def __init__(self, dim, eps=1e-8):
        super().__init__()
        self.eps = eps
        self.scale = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        # x: (B, T, C)
        norm = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return x * norm * self.scale


class MultiHeadAttention(nn.Module):
    def __init__(self, n_embd, n_head, attn_pdrop=0.1, resid_pdrop=0.1, use_rotary=True):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.head_dim = n_embd // n_head
        self.scale = 1.0 / math.sqrt(self.head_dim)

        self.qkv = nn.Linear(n_embd, n_embd * 3, bias=False)
        self.proj = nn.Linear(n_embd, n_embd)
        self.attn_dropout = nn.Dropout(attn_pdrop)
        self.resid_dropout = nn.Dropout(resid_pdrop)
        self.use_rotary = use_rotary
        if use_rotary:
            self.rotary = RotaryEmbedding(self.head_dim)

        # optional flash attention detection
        self.use_flash = False
        try:
            # try common flash attention package
            import flash_attn  # type: ignore
            self.use_flash = True
        except Exception:
            self.use_flash = False

    def forward(self, x, attn_mask=None):
        B, T, C = x.size()
        qkv = self.qkv(x).view(B, T, 3, self.n_head, self.head_dim).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]  # each (B, nh, T, hs)

        if self.use_rotary:
            sin, cos = self.rotary(T, device=x.device)
            q, k = apply_rotary_pos_emb(q, k, sin, cos)

        if self.use_flash:
            # best-effort: if flash attention is available, try to use it (APIs vary by package)
            try:
                # flatten for flash attention calls
                qkv = torch.stack((q, k, v), dim=2)
                # fallback to manual matmul if API unknown
                raise RuntimeError('flash-attn integration placeholder; falling back')
            except Exception:
                att = torch.matmul(q, k.transpose(-2, -1)) * self.scale
        else:
            att = torch.matmul(q, k.transpose(-2, -1)) * self.scale

        causal_mask = torch.tril(torch.ones(T, T, device=x.device)).view(1, 1, T, T)
        att = att.masked_fill(causal_mask == 0, float('-inf'))

        if attn_mask is not None:
            if attn_mask.dim() == 2:
                attn_mask = attn_mask.view(B, 1, 1, T)
            att = att.masked_fill(attn_mask == 0, float('-inf'))

        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)
        y = torch.matmul(att, v)  # (B, nh, T, hs)
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.proj(y)
        y = self.resid_dropout(y)
        return y


class SwiGLU(nn.Module):
    def __init__(self, dim_in, dim_out):
        super().__init__()
        # dim_out is the inner dim; we keep ability to set it equal to dim_in for smaller models
        self.fc1 = nn.Linear(dim_in, dim_out)
        self.fc_gate = nn.Linear(dim_in, dim_out)
        self.fc2 = nn.Linear(dim_out, dim_in)
        self.dropout = nn.Dropout(0.0)

    def forward(self, x):
        return self.fc2(F.silu(self.fc1(x)) * self.fc_gate(x))


class FeedForward(nn.Module):
    def __init__(self, n_embd, mlp_ratio=1.0, pdrop=0.1, inner_dim=None):
        super().__init__()
        # Allow inner_dim override; default reduce to match embedding for compact model
        if inner_dim is None:
            inner = int(n_embd * mlp_ratio)
        else:
            inner = inner_dim
        self.fn = SwiGLU(n_embd, inner)
        self.dropout = nn.Dropout(pdrop)

    def forward(self, x, tag_emb=None):
        # tag_emb is accepted for API compatibility with MoE variants that may use router bias
        return self.dropout(self.fn(x))


class MoEFeedForward(nn.Module):
    """Mixture-of-Experts feedforward: small top-k router routing per token.

    Notes: simplified router for resource-constrained mini models. Uses token-level routing.
    """
    def __init__(self, n_embd, num_experts=4, top_k=1, expert_ctor=None, router_temperature=1.0, aux_coef=0.0, tag_proj_dim=None):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.router_temperature = router_temperature
        self.aux_coef = aux_coef
        assert 1 <= top_k <= num_experts
        if expert_ctor is None:
            expert_ctor = lambda: FeedForward(n_embd)
        self.experts = nn.ModuleList([expert_ctor() for _ in range(num_experts)])
        # lightweight router: linear to num_experts
        self.router = nn.Linear(n_embd, num_experts)
        # optional projection from a tag embedding (B, C) -> (B, num_experts) to bias router logits
        self.tag_proj = nn.Linear(tag_proj_dim, num_experts) if tag_proj_dim is not None else None

    def forward(self, x, tag_emb=None):
        # x: (B, T, C)
        B, T, C = x.size()
        logits = self.router(x)  # (B, T, num_experts)
        # if a tag embedding is provided (B, C) and we have a projection, add it as a bias
        if tag_emb is not None and self.tag_proj is not None:
            # project per-batch tag embedding to expert logits and broadcast to tokens
            # tag_emb: (B, C) -> (B, num_experts) -> (B, 1, num_experts)
            tag_bias = self.tag_proj(tag_emb).unsqueeze(1)
            logits = logits + tag_bias
        # apply temperature to router logits
        if self.router_temperature and self.router_temperature != 1.0:
            probs = F.softmax(logits / float(self.router_temperature), dim=-1)
        else:
            probs = F.softmax(logits, dim=-1)
        # topk indices
        topk = probs.topk(self.top_k, dim=-1)
        indices = topk.indices  # (B, T, top_k)
        weights = topk.values  # (B, T, top_k)

        out = x.new_zeros(B, T, C)
        # naive per-expert dispatch (may be slower but simple)
        for e in range(self.num_experts):
            # mask tokens that route to expert e
            mask = (indices == e)  # (B, T, top_k)
            if not mask.any():
                continue
            # combine along top_k: compute contribution weight per (B,T)
            # for tokens where expert e selected, create input slice
            sel = mask.any(-1)  # (B, T)
            if not sel.any():
                continue
            inp = x[sel]
            expert_out = self.experts[e](inp)
            # add weighted contribution
            # weights for those selected tokens: take max across top_k positions where index==e
            w = torch.zeros(B, T, device=x.device)
            for k in range(self.top_k):
                w = w + (indices[..., k] == e).float() * weights[..., k]
            w_sel = w[sel].unsqueeze(-1)
            out[sel] = out[sel] + expert_out * w_sel

        # compute lightweight auxiliary load-balancing loss (optional)
        self.last_aux_loss = None
        if getattr(self, 'aux_coef', 0.0):
            # average probability mass per expert across tokens
            load = probs.sum(dim=(0, 1)) / (B * T)
            aux = (load * load).sum()
            self.last_aux_loss = aux * float(self.aux_coef)

        return out


class TransformerBlock(nn.Module):
    def __init__(self, n_embd, n_head, mlp_ratio=4, attn_pdrop=0.1, resid_pdrop=0.1, use_rotary=True):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = MultiHeadAttention(n_embd, n_head, attn_pdrop, resid_pdrop, use_rotary=use_rotary)
        self.ln2 = nn.LayerNorm(n_embd)
        self.mlp = FeedForward(n_embd, mlp_ratio, resid_pdrop)

    def forward(self, x, attn_mask=None, tag_emb=None):
        x = x + self.attn(self.ln1(x), attn_mask=attn_mask)
        # allow mlp variants (MoE) to accept tag_emb
        x = x + (self.mlp(self.ln2(x), tag_emb=tag_emb) if hasattr(self.mlp, '__call__') else self.mlp(self.ln2(x)))
        return x


class Hanuman(nn.Module):
    """Hanuman: advanced GPT-like mini model with rotary embeddings and SwiGLU MLP.

    Compatible forward signature with HF GPT2LMHeadModel: forward(input_ids, attention_mask, labels)
    Returns SimpleNamespace(loss=..., logits=...)
    """

    def __init__(self, *, vocab_size, n_positions=4096, n_embd=512, n_layer=8, n_head=8, mlp_ratio=1.0,
                 attn_pdrop=0.1, resid_pdrop=0.1, use_rotary=True, use_rmsnorm=True, use_moe=False,
                 moe_experts=4, moe_top_k=1, gradient_checkpointing=False, use_think_head=False, think_aux_coef=1.0):
        super().__init__()
        self.vocab_size = vocab_size
        self.n_positions = n_positions
        self.n_embd = n_embd

        self.use_rmsnorm = use_rmsnorm
        self.gradient_checkpointing = gradient_checkpointing

        self.wte = nn.Embedding(vocab_size, n_embd)
        self.wpe = nn.Embedding(n_positions, n_embd)
        self.drop = nn.Dropout(0.1)

        self.blocks = nn.ModuleList()
        for _ in range(n_layer):
            blk = TransformerBlock(n_embd, n_head, mlp_ratio, attn_pdrop, resid_pdrop, use_rotary=use_rotary)
            self.blocks.append(blk)

        # final norm: RMSNorm or LayerNorm
        if use_rmsnorm:
            self.ln_f = RMSNorm(n_embd)
        else:
            self.ln_f = nn.LayerNorm(n_embd)

        # optional MoE on top of feedforwards inside blocks: swap block.mlp with MoE variant
        if use_moe:
            for blk in self.blocks:
                blk.mlp = MoEFeedForward(n_embd, num_experts=moe_experts, top_k=moe_top_k,
                                          expert_ctor=lambda: FeedForward(n_embd, mlp_ratio=mlp_ratio, inner_dim=n_embd))

        self.head = nn.Linear(n_embd, vocab_size, bias=False)

        # optional think head for intermediate reasoning outputs (same vocab by default)
        self.use_think_head = use_think_head
        self.think_aux_coef = float(think_aux_coef)
        if use_think_head:
            self.think_head = nn.Linear(n_embd, vocab_size, bias=False)

    def forward(self, input_ids=None, attention_mask=None, labels=None, thought_labels=None):
        B, T = input_ids.size()
        assert T <= self.n_positions, f"Sequence length {T} > model max {self.n_positions}"

        pos = torch.arange(0, T, dtype=torch.long, device=input_ids.device).unsqueeze(0)
        x = self.wte(input_ids) + self.wpe(pos)
        x = self.drop(x)

        # If user provided a special effort tag token (e.g., first token in input), compute tag_emb
        tag_emb = None
        try:
            # detect if first token corresponds to a special think token id set on the model
            if hasattr(self, 'think_token_ids') and isinstance(self.think_token_ids, dict):
                # look for a single-tag indicator in input_ids (assumed at position 0)
                first = input_ids[:, 0]
                # if a known tag id is present, make tag_emb from its token embedding
                for tag, tid in self.think_token_ids.items():
                    if (first == tid).any():
                        tag_emb = self.wte(tid).unsqueeze(0).expand(input_ids.size(0), -1)
                        break
        except Exception:
            tag_emb = None

        for blk in self.blocks:
            if self.gradient_checkpointing and self.training:
                x = grad_checkpoint(blk, x, attention_mask, tag_emb)
            else:
                x = blk(x, attn_mask=attention_mask, tag_emb=tag_emb)

        x = self.ln_f(x)
        logits = self.head(x)

        loss = None
        thought_loss = None
        if labels is not None:
            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
            lm_loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
            loss = lm_loss

        # optional thinking head loss
        thought_logits = None
        if self.use_think_head and thought_labels is not None:
            thought_logits = self.think_head(x)
            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
            thought_loss = loss_fct(thought_logits.view(-1, thought_logits.size(-1)), thought_labels.view(-1))
            if loss is None:
                loss = thought_loss * self.think_aux_coef
            else:
                loss = loss + thought_loss * self.think_aux_coef

        return SimpleNamespace(loss=loss, logits=logits, thought_logits=thought_logits, thought_loss=thought_loss)

    # runtime helpers
    def to_device(self, device):
        self.to(device)

    def enable_fp16(self):
        # cast model params to float16 where safe
        self.half()

    def set_gradient_checkpointing(self, enabled: bool):
        self.gradient_checkpointing = enabled

    # Simple autoregressive generator (CPU/GPU). Not optimized for speed.
    @torch.no_grad()
    def generate(self, input_ids, max_new_tokens=50, temperature=1.0, top_k=0, top_p=0.0, eos_token_id=None):
        device = input_ids.device
        self.eval()
        out = input_ids
        for _ in range(max_new_tokens):
            logits = self.forward(input_ids=out).logits
            next_logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
            if top_k > 0:
                vals, idx = torch.topk(next_logits, top_k)
                probs = torch.zeros_like(next_logits).scatter(1, idx, F.softmax(vals, dim=-1))
            elif top_p > 0.0:
                sorted_logits, sorted_indices = torch.sort(next_logits, descending=True)
                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
                cutoff = cumulative_probs > top_p
                cutoff_index = torch.argmax(cutoff.int(), dim=-1)
                mask = torch.zeros_like(sorted_logits).bool()
                for b in range(sorted_logits.size(0)):
                    mask[b, :cutoff_index[b]+1] = True
                probs = torch.zeros_like(next_logits)
                probs.scatter_(1, sorted_indices, F.softmax(sorted_logits, dim=-1) * mask.float())
            else:
                probs = F.softmax(next_logits, dim=-1)

            next_token = torch.multinomial(probs, num_samples=1)
            out = torch.cat([out, next_token], dim=1)
            if eos_token_id is not None and next_token.item() == eos_token_id:
                break
        return out

    @torch.no_grad()
    def generate_effort(self, input_ids, effort='short', reason_budget=None, temperature=1.0, top_k=0, top_p=0.0, eos_token_id=None):
        """
        Two-phase decoding: generate reasoning tokens inside a <scratch> block up to reason_budget, then generate final answer after <final>.
        effort in {'none','short','medium','long'} maps to default budgets if reason_budget is None.
        This is a simple, synchronous implementation; production should use batched, streaming decodes.
        """
        budget_map = {'none': 0, 'short': 64, 'medium': 256, 'long': 1024}
        if reason_budget is None:
            reason_budget = budget_map.get(effort, 64)

        device = input_ids.device
        model = self
        # phase 1: generate reasoning tokens if budget > 0
        out = input_ids
        if reason_budget > 0:
            for _ in range(reason_budget):
                logits = model.forward(input_ids=out).logits
                next_logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
                probs = F.softmax(next_logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)
                out = torch.cat([out, next_token], dim=1)
        # phase 2: generate final answer until eos or short fixed length
        final_out = out
        for _ in range(128):
            logits = model.forward(input_ids=final_out).logits
            next_logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
            probs = F.softmax(next_logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            final_out = torch.cat([final_out, next_token], dim=1)
            if eos_token_id is not None and next_token.item() == eos_token_id:
                break
        return final_out

    # Utilities to play nice with train.py expectations
    def save_pretrained(self, out_dir: str, use_safetensors: bool = False):
        os.makedirs(out_dir, exist_ok=True)
        # save state and a small config
        model_path = os.path.join(out_dir, 'pytorch_model.bin')
        cfg = {
            'vocab_size': self.vocab_size,
            'n_positions': self.n_positions,
            'n_embd': self.n_embd,
            'n_layer': len(self.blocks),
            'n_head': self.blocks[0].attn.n_head if len(self.blocks) else 0,
        }
        with open(os.path.join(out_dir, 'config.json'), 'w', encoding='utf-8') as f:
            json.dump(cfg, f)

        if use_safetensors:
            try:
                from safetensors.torch import save_file as safe_save
                state = {k: v.cpu() for k, v in self.state_dict().items()}
                safe_save(state, os.path.join(out_dir, 'pytorch_model.safetensors'))
                return
            except Exception:
                # fallback to torch.save if safetensors isn't available
                pass

        torch.save(self.state_dict(), model_path)

    @classmethod
    def from_pretrained(cls, in_dir: str, map_location=None):
        with open(os.path.join(in_dir, 'config.json'), 'r', encoding='utf-8') as f:
            cfg = json.load(f)
        model = cls(
            vocab_size=cfg.get('vocab_size', 32000),
            n_positions=cfg.get('n_positions', 1024),
            n_embd=cfg.get('n_embd', 768),
            n_layer=cfg.get('n_layer', 12),
            n_head=cfg.get('n_head', 12),
        )
        # Prefer safetensors if present
        safetensors_path = os.path.join(in_dir, 'pytorch_model.safetensors')
        bin_path = os.path.join(in_dir, 'pytorch_model.bin')
        if os.path.exists(safetensors_path):
            try:
                from safetensors.torch import load_file as safe_load
                state = safe_load(safetensors_path, device=map_location or 'cpu')
            except Exception:
                state = torch.load(safetensors_path, map_location=map_location)
        elif os.path.exists(bin_path):
            state = torch.load(bin_path, map_location=map_location)
        else:
            raise FileNotFoundError(f'No model file found in {in_dir}')

        # state is a mapping of tensors
        model.load_state_dict(state)
        return model

    def resize_token_embeddings(self, new_vocab_size: int):
        old_wte = self.wte
        old_vocab, emb_dim = old_wte.weight.shape
        if new_vocab_size == old_vocab:
            return
        new_wte = nn.Embedding(new_vocab_size, emb_dim)
        # copy existing weights
        with torch.no_grad():
            new_wte.weight[:old_vocab] = old_wte.weight
        self.wte = new_wte

        new_head = nn.Linear(emb_dim, new_vocab_size, bias=False)
        with torch.no_grad():
            new_head.weight[:,:old_vocab] = self.head.weight
        self.head = new_head


def build_from_config(config):
    # Build Hanuman from a GPT2Config-like object with mini-model defaults
    return Hanuman(
        vocab_size=getattr(config, 'vocab_size', 32000),
        n_positions=getattr(config, 'n_positions', getattr(config, 'n_ctx', 4096)),
        n_embd=getattr(config, 'n_embd', 512),
        n_layer=getattr(config, 'n_layer', 8),
        n_head=getattr(config, 'n_head', 8),
        mlp_ratio=getattr(config, 'mlp_ratio', 1.0),
        use_rmsnorm=getattr(config, 'use_rmsnorm', True),
        use_moe=getattr(config, 'use_moe', False),
        moe_experts=getattr(config, 'moe_experts', 4),
        moe_top_k=getattr(config, 'moe_top_k', 1),
    gradient_checkpointing=getattr(config, 'gradient_checkpointing', False),
    use_think_head=getattr(config, 'use_think_head', False),
    think_aux_coef=getattr(config, 'think_aux_coef', 1.0),
    )