#!/usr/bin/env python3
"""
SynapCore Phobos-Scratch v1 — DeepInfer (DeepSeek-V4.1) Architectural Modules
(`phobos-1/models/phobos_scratch_v1/deepinfer_modules.py`)
=============================================================================
Pure-PyTorch, A100 (`SM80`) & training-compatible implementations of the 5 core
architectural innovations from `https://github.com/deepseek-ai/DeepInfer`:

  1. `DeepInferNoAuxGate` (`v4_modules/moe.py` -> `TopKGate`):
       Uses `sqrtsoftplus` scoring + `noaux_tc` (`e_score_correction_bias`) dynamic
       auxiliary-loss-free load balancing and normalized top-k probabilities.
  2. `DeepInferEngramMemory` (`v4_modules/engram.py` & `engram_metadata.py`):
       Prime-modulo causal 2-gram & 3-gram hash lookup tables with gated KV projection
       fused directly into transformer hidden states.
  3. `DeepInferHyperConnection` (`v4_modules/hc_blocks.py` & `mhc_modeling.py`):
       Multi-Head Hyper-Connections (`hc_mult = 4` parallel residual streams) with
       learned dynamic sigmoid pre-mix, post-mix, and residual stream coupling.
  4. `DeepInferMarkovMTPHead` (`v4_modules/heads.py` & `mtp/dflash.py`):
       `MarkovGPTHead` low-rank first-order bigram logit correction (`markov_w1/w2`)
       + `DFlash` next-token draft projection (`MTP`) + `ConfidenceHead`.
  5. `DeepInferStructuredMasker` (`structure_output/mask.py` & `matcher.py`):
       Step-wise token logit bitmasking guaranteeing balanced `<think>...</think>`
       and `<tool_call>...</tool_call>` blocks and preventing premature `<|eos|>`.

Strictly <= 320 lines.
"""

from __future__ import annotations

from typing import Dict, List, Optional, Tuple

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


def sqrtsoftplus(x: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
    """DeepInfer `TopKGate` scoring function: `sqrt(softplus(x) + eps)`."""
    return torch.sqrt(F.softplus(x.float()) + eps).to(x.dtype)


class DeepInferNoAuxGate(nn.Module):
    """
    DeepInfer `TopKGate` (`v4_modules/moe.py`) with `sqrtsoftplus` activation and
    `noaux_tc` (`e_score_correction_bias`) auxiliary-loss-free load balancing.
    """

    def __init__(self, num_branches: int = 4, num_leaves: int = 8, bias_lr: float = 0.0015) -> None:
        super().__init__()
        self.num_branches = num_branches
        self.num_leaves = num_leaves
        self.bias_lr = bias_lr
        self.register_buffer("branch_correction_bias", torch.zeros(num_branches, dtype=torch.float32))
        self.register_buffer("leaf_correction_bias", torch.zeros(num_leaves, dtype=torch.float32))

    def score_and_balance(
        self, branch_logits: torch.Tensor, leaf_logits: torch.Tensor, top_k_leaves: int = 2
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Computes `sqrtsoftplus` scores, applies `e_score_correction_bias` for top-k choice
        (`noaux_tc`), updates correction biases during training, and returns normalized weights.
        """
        b_scores = sqrtsoftplus(branch_logits)
        l_scores = sqrtsoftplus(leaf_logits)

        b_choice = b_scores + self.branch_correction_bias.to(b_scores.dtype)
        l_choice = l_scores + self.leaf_correction_bias.to(l_scores.dtype)

        _, top_l_idx = torch.topk(l_choice, k=top_k_leaves, dim=-1)
        top_l_raw = torch.gather(l_scores, dim=-1, index=top_l_idx)
        top_l_weights = top_l_raw / (top_l_raw.sum(dim=-1, keepdim=True) + 1e-6)

        if self.training:
            with torch.no_grad():
                b_top1 = b_choice.argmax(dim=-1).reshape(-1)
                b_freq = torch.bincount(b_top1, minlength=self.num_branches).float() / max(1, b_top1.numel())
                b_target = 1.0 / float(self.num_branches)
                self.branch_correction_bias.add_(self.bias_lr * torch.sign(b_target - b_freq)).clamp_(-0.4, 0.4)

                l_flat = top_l_idx.reshape(-1)
                l_freq = torch.bincount(l_flat, minlength=self.num_leaves).float() / max(1, l_flat.numel())
                l_target = 1.0 / float(self.num_leaves)
                self.leaf_correction_bias.add_(self.bias_lr * torch.sign(l_target - l_freq)).clamp_(-0.4, 0.4)

        b_probs = b_scores / (b_scores.sum(dim=-1, keepdim=True) + 1e-6)
        l_probs = l_scores / (l_scores.sum(dim=-1, keepdim=True) + 1e-6)
        return b_probs, l_probs, top_l_idx, top_l_weights


class DeepInferEngramMemory(nn.Module):
    """
    DeepInfer `Engram` (`v4_modules/engram.py` & `engram_metadata.py`):
    Hashes causal 2-grams `(x_{t-1}, x_t)` and 3-grams `(x_{t-2}, x_{t-1}, x_t)` into
    distinct prime-sized tables, projects through a fused KV linear layer, and gates into `h`.
    """

    def __init__(
        self,
        d_model: int,
        layer_id: int,
        ngram_dim: int = 48,
        prime_2gram: int = 8191,
        prime_3gram: int = 8209,
    ) -> None:
        super().__init__()
        self.prime_2gram = prime_2gram
        self.prime_3gram = prime_3gram
        # Deterministic odd multipliers per layer (`NgramHashMapping` convention)
        m0 = 2 * (10007 * (layer_id + 1) + 31) + 1
        m1 = 2 * (20011 * (layer_id + 1) + 97) + 1
        m2 = 2 * (30013 * (layer_id + 1) + 193) + 1
        self.register_buffer("multipliers", torch.tensor([m0, m1, m2], dtype=torch.long), persistent=False)

        self.embed_2gram = nn.Embedding(prime_2gram, ngram_dim)
        self.embed_3gram = nn.Embedding(prime_3gram, ngram_dim)
        nn.init.normal_(self.embed_2gram.weight, mean=0.0, std=0.02)
        nn.init.normal_(self.embed_3gram.weight, mean=0.0, std=0.02)

        self.fused_kv_projection = nn.Linear(ngram_dim * 2, d_model * 2, bias=False)
        self.weight_hidden = nn.Parameter(torch.ones(d_model) * 0.5)
        self.weight_embed = nn.Parameter(torch.ones(d_model) * 0.5)
        self.res_scale = nn.Parameter(torch.tensor(0.15))

    def forward(self, h: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor:
        """Applies causal 2-gram & 3-gram hash lookup and gated fusion to `h` (`[B, T, D]`)."""
        ids_64 = input_ids.long()
        prev1 = F.pad(ids_64[:, :-1], (1, 0), value=0)
        prev2 = F.pad(ids_64[:, :-2], (2, 0), value=0)

        m0, m1, m2 = self.multipliers[0], self.multipliers[1], self.multipliers[2]
        h2 = torch.remainder(ids_64 * m0 + prev1 * m1, self.prime_2gram)
        h3 = torch.remainder(ids_64 * m0 + prev1 * m1 + prev2 * m2, self.prime_3gram)

        ngram_feat = torch.cat([self.embed_2gram(h2), self.embed_3gram(h3)], dim=-1).to(h.dtype)
        k_ng, v_ng = self.fused_kv_projection(ngram_feat).chunk(2, dim=-1)

        h_norm = h * torch.rsqrt(h.pow(2).mean(dim=-1, keepdim=True) + 1e-6)
        gate = torch.sigmoid(h_norm * self.weight_hidden.to(h.dtype) + k_ng * self.weight_embed.to(h.dtype))
        return h + self.res_scale.to(h.dtype) * (gate * v_ng)


class DeepInferHyperConnection(nn.Module):
    """
    DeepInfer `SigmoidHCBlockMixin` (`v4_modules/hc_blocks.py` & `mhc_modeling.py`):
    Maintains `hc_mult = 4` parallel residual streams `[B, T, M, D]` and learns dynamic
    pre-layer mixing, post-layer writeback, and cross-stream residual coupling.
    """

    def __init__(self, d_model: int, layer_id: int, hc_mult: int = 4) -> None:
        super().__init__()
        self.hc_mult = hc_mult
        self.d_model = d_model
        # Dynamic summary projection from each stream `D -> 2+M` scores with symmetry-breaking init
        self.stream_score = nn.Linear(d_model, 2 + hc_mult, bias=True)
        nn.init.normal_(self.stream_score.weight, mean=0.0, std=0.02 / (d_model ** 0.5))
        nn.init.zeros_(self.stream_score.bias)

        base_pre = torch.linspace(-1.0, 1.0, steps=hc_mult)
        base_pre[layer_id % hc_mult] = 1.5
        self.base_pre = nn.Parameter(base_pre)

        base_post = torch.full((hc_mult,), -1.0)
        base_post[layer_id % hc_mult] = 1.5
        self.base_post = nn.Parameter(base_post)
        self.base_res = nn.Parameter(torch.eye(hc_mult) * 2.5 - 1.25)
        self.dynamic_scale = nn.Parameter(torch.tensor(0.25))

    def pre_mix(self, trunk: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Mixes `[B, T, M, D]` trunk into sub-layer input `[B, T, D]` and returns `(x_in, post_w, res_mat)`."""
        raw = self.stream_score(trunk) * self.dynamic_scale.to(trunk.dtype)  # [B, T, M, 2+M]
        pre_logits = raw[..., 0] + self.base_pre.to(trunk.dtype)  # [B, T, M]
        post_logits = raw[..., 1] + self.base_post.to(trunk.dtype)  # [B, T, M]
        res_logits = raw[..., 2:] + self.base_res.to(trunk.dtype)  # [B, T, M, M]

        pre_w = F.softmax(pre_logits.float(), dim=-1).to(trunk.dtype)
        post_w = torch.sigmoid(post_logits)
        res_mat = F.softmax(res_logits.float(), dim=-1).to(trunk.dtype)

        x_in = torch.einsum("btm,btmd->btd", pre_w, trunk)
        return x_in, post_w, res_mat

    def post_mix(
        self, sublayer_out: torch.Tensor, trunk: torch.Tensor, post_w: torch.Tensor, res_mat: torch.Tensor
    ) -> torch.Tensor:
        """Writes `sublayer_out` (`[B, T, D]`) back into the `M`-stream residual trunk (`[B, T, M, D]`)."""
        mixed_trunk = torch.einsum("btij,btjd->btid", res_mat, trunk)
        return mixed_trunk + post_w.unsqueeze(-1) * sublayer_out.unsqueeze(-2)


class DeepInferMarkovMTPHead(nn.Module):
    """
    DeepInfer `MarkovGPTHead` + `DFlash` Next-Token MTP & `ConfidenceHead`
    (`v4_modules/heads.py` & `models/mtp/dflash.py`).
    """

    def __init__(self, vocab_size: int, d_model: int, markov_rank: int = 64) -> None:
        super().__init__()
        self.markov_rank = markov_rank
        self.markov_w1 = nn.Embedding(vocab_size, markov_rank)
        self.markov_w2 = nn.Linear(markov_rank, vocab_size, bias=False)
        nn.init.normal_(self.markov_w1.weight, mean=0.0, std=0.02)
        nn.init.zeros_(self.markov_w2.weight)
        self.markov_scale = nn.Parameter(torch.tensor(0.25))

        # DFlash lightweight next-2 token (MTP) projection + ConfidenceHead
        self.mtp_proj = nn.Linear(d_model + markov_rank, d_model, bias=False)
        self.confidence_head = nn.Linear(d_model + markov_rank, 1, bias=False)

    def forward_logits_and_mtp(
        self,
        h_norm: torch.Tensor,
        input_ids: torch.Tensor,
        lm_head_weight: torch.Tensor,
        pad_id: int = 0,
        compute_mtp_loss: bool = False,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Computes main logits with `MarkovGPTHead` first-order bigram correction, token
        confidence probability (`ConfidenceHead`), and optional `DFlash` 2-step MTP loss.
        """
        m_emb = self.markov_w1(input_ids).to(h_norm.dtype)  # [B, T, markov_rank]
        conf_in = torch.cat([h_norm, m_emb], dim=-1)
        fused_w = torch.cat(
            [lm_head_weight, self.markov_scale.to(lm_head_weight.dtype) * self.markov_w2.weight],
            dim=-1,
        )
        logits = F.linear(conf_in, fused_w)
        confidence = torch.sigmoid(self.confidence_head(conf_in)).squeeze(-1)

        mtp_loss = h_norm.new_zeros(())
        if compute_mtp_loss and input_ids.size(1) > 2:
            # Predict token t+2 from (h_t, markov_emb(x_{t+1})) only on non-pad targets
            next1_emb = m_emb[:, 1:-1, :]
            mtp_in = torch.cat([h_norm[:, :-2, :], next1_emb], dim=-1)
            mtp_h = F.silu(self.mtp_proj(mtp_in)).reshape(-1, h_norm.size(-1))
            targets_t2 = input_ids[:, 2:].reshape(-1)
            valid_mask = targets_t2 != pad_id
            if valid_mask.any():
                valid_h = mtp_h[valid_mask]
                valid_t = targets_t2[valid_mask]
                chunk_sz = 16384
                total_loss = h_norm.new_zeros((), dtype=torch.float32)
                for s in range(0, valid_h.size(0), chunk_sz):
                    e = min(valid_h.size(0), s + chunk_sz)
                    c_logits = F.linear(valid_h[s:e], lm_head_weight).float()
                    total_loss = total_loss + F.cross_entropy(c_logits, valid_t[s:e], reduction="sum")
                mtp_loss = total_loss / float(valid_h.size(0))

        return logits, confidence, mtp_loss

    def forward_fused_training_losses(
        self,
        h_norm: torch.Tensor,
        input_ids: torch.Tensor,
        lm_head_weight: torch.Tensor,
        pad_id: int = 0,
        assistant_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, float, torch.Tensor, float]:
        """
        Memory-efficient chunked training loss for large A100 batches (`65K` tokens/step):
        Computes assistant-salience next-token CE, raw CE, and 2-step DFlash MTP loss only on
        non-pad tokens in 16,384-token chunks without allocating full `[B, T, V]` tensors.
        """
        m_emb = self.markov_w1(input_ids).to(h_norm.dtype)
        conf_in = torch.cat([h_norm, m_emb], dim=-1)
        mean_conf = float(torch.sigmoid(self.confidence_head(conf_in)).detach().mean().item())
        fused_w = torch.cat(
            [lm_head_weight, self.markov_scale.to(lm_head_weight.dtype) * self.markov_w2.weight],
            dim=-1,
        )

        shift_in = conf_in[:, :-1, :].reshape(-1, conf_in.size(-1))
        shift_t = input_ids[:, 1:].reshape(-1)
        shift_asst = assistant_mask[:, 1:].reshape(-1).float() if assistant_mask is not None else None
        valid_1 = shift_t != pad_id
        salience_ce = h_norm.new_zeros((), dtype=torch.float32)
        raw_ce_val = 0.0
        chunk_sz = 16384
        if valid_1.any():
            vin = shift_in[valid_1]
            vt = shift_t[valid_1]
            vasst = shift_asst[valid_1] if shift_asst is not None else None
            w_num = h_norm.new_zeros((), dtype=torch.float32)
            w_den = h_norm.new_zeros((), dtype=torch.float32)
            raw_sum = h_norm.new_zeros((), dtype=torch.float32)
            for s in range(0, vin.size(0), chunk_sz):
                e = min(vin.size(0), s + chunk_sz)
                c_logits = F.linear(vin[s:e], fused_w).float()
                pt_ce = F.cross_entropy(c_logits, vt[s:e], reduction="none")
                with torch.no_grad():
                    sw = torch.ones_like(pt_ce)
                    sw = torch.where(pt_ce >= 2.2, torch.full_like(sw, 2.0), sw)
                    sw = torch.where(pt_ce <= 0.25, torch.full_like(sw, 0.25), sw)
                    if vasst is not None:
                        sw = sw * (0.20 + 1.80 * vasst[s:e])
                w_num = w_num + (pt_ce * sw).sum()
                w_den = w_den + sw.sum()
                raw_sum = raw_sum + pt_ce.detach().sum()
            salience_ce = w_num / w_den.clamp_min(1.0)
            raw_ce_val = float((raw_sum / float(vin.size(0))).item())

        mtp_loss = h_norm.new_zeros((), dtype=torch.float32)
        if input_ids.size(1) > 2:
            mtp_in = torch.cat([h_norm[:, :-2, :], m_emb[:, 1:-1, :]], dim=-1)
            mtp_h = F.silu(self.mtp_proj(mtp_in)).reshape(-1, h_norm.size(-1))
            targets_t2 = input_ids[:, 2:].reshape(-1)
            valid_2 = targets_t2 != pad_id
            if valid_2.any():
                vh2 = mtp_h[valid_2]
                vt2 = targets_t2[valid_2]
                t2_sum = h_norm.new_zeros((), dtype=torch.float32)
                for s in range(0, vh2.size(0), chunk_sz):
                    e = min(vh2.size(0), s + chunk_sz)
                    c_logits2 = F.linear(vh2[s:e], lm_head_weight).float()
                    t2_sum = t2_sum + F.cross_entropy(c_logits2, vt2[s:e], reduction="sum")
                mtp_loss = t2_sum / float(vh2.size(0))

        return salience_ce, raw_ce_val, mtp_loss, mean_conf


class DeepInferStructuredMasker:
    """
    DeepInfer `XGrammarUtils` / `XGrammarMatcher` (`structure_output/mask.py`):
    Applies step-wise logit bitmasking during generation to enforce well-formed
    `<think>...</think>` and `<tool_call>...</tool_call>` blocks and block premature `<|eos|>`.
    """

    @staticmethod
    def apply_step_constraints_inplace(
        logits: torch.Tensor,
        generated_seq: List[int],
        eos_id: Optional[int],
        think_start_id: int = 5,
        think_end_id: int = 6,
        tool_start_id: int = 7,
        tool_end_id: int = 8,
        min_tokens_before_eos: int = 8,
        steps_generated: int = 0,
    ) -> None:
        """Modifies `logits` (`[B, V]`) in-place to enforce structural tag state machine rules."""
        n_think_open = sum(1 for t in generated_seq if t == think_start_id)
        n_think_close = sum(1 for t in generated_seq if t == think_end_id)
        in_think = n_think_open > n_think_close

        n_tool_open = sum(1 for t in generated_seq if t == tool_start_id)
        n_tool_close = sum(1 for t in generated_seq if t == tool_end_id)
        in_tool = n_tool_open > n_tool_close

        if in_think:
            # Cannot open nested <think> or emit <|eos|> before closing </think>
            logits[:, think_start_id] = -1e9
            if eos_id is not None:
                logits[:, eos_id] = -1e9
        else:
            # Cannot emit orphan </think> when not inside <think>
            logits[:, think_end_id] = -1e9

        if in_tool:
            logits[:, tool_start_id] = -1e9
            if eos_id is not None:
                logits[:, eos_id] = -1e9
        else:
            logits[:, tool_end_id] = -1e9

        if eos_id is not None and steps_generated < min_tokens_before_eos:
            logits[:, eos_id] = -1e9