"""M4 — Mixture of Recursions (token-level dynamic depth) for ARCHON.

NeurIPS 2025 (https://github.com/raymin0223/mixture_of_recursions)
+ ARCHON innovation: combine with MTP-5 task_conditional.

Idea: ARCHON 18L stack is REUSED N times per token, where N depends on a
lightweight router. Easy tokens use 1× (~9L effective), hard tokens use 4×
(~72L effective). Same weights, dynamic compute.

For ARCHON 282M: estimate equivalent of 1B-3B at hard tokens with no extra params.

Architecture change required: router head. After warm SFT, train router with
distillation from Claude on (token, optimal_depth) pairs.

PHASE D — requires architecture mod + retraining router (~$50 GPU, 12-16h V100).
"""
from __future__ import annotations

from dataclasses import dataclass

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


@dataclass
class MoRConfig:
    """MoR hyper-params for ARCHON."""

    n_recursion_buckets: int = 3  # [1×, 2×, 4×]
    bucket_recursions: tuple = (1, 2, 4)
    router_hidden: int = 128
    aux_loss_weight: float = 0.01  # Encourage balanced bucket usage


class MoRRouter(nn.Module):
    """Token-level depth router. Output: softmax over n buckets."""

    def __init__(self, hidden_dim: int, cfg: MoRConfig = MoRConfig()):
        super().__init__()
        self.cfg = cfg
        self.proj1 = nn.Linear(hidden_dim, cfg.router_hidden)
        self.proj2 = nn.Linear(cfg.router_hidden, cfg.n_recursion_buckets)

    def forward(self, h: torch.Tensor) -> torch.Tensor:
        """h [B, T, D] -> bucket_probs [B, T, n_buckets]."""
        x = F.gelu(self.proj1(h))
        return F.softmax(self.proj2(x), dim=-1)


class ArchonMoRWrapper(nn.Module):
    """Wrap ArchonBrain with MoR recursion logic.

    Forward:
        1. Embed -> h
        2. Router predicts bucket per token from initial h
        3. For each bucket, run h through layers N times (N = bucket_recursion)
        4. Mix outputs weighted by router probs
        5. Apply final norm + lm_head + MTP heads as usual
    """

    def __init__(self, archon_model, cfg: MoRConfig = MoRConfig()):
        super().__init__()
        self.archon = archon_model
        self.cfg = cfg
        self.router = MoRRouter(archon_model.config.hidden_dim, cfg)

    def forward(self, input_ids: torch.Tensor, targets: torch.Tensor = None):
        h0 = self.archon.embed(input_ids)
        bucket_probs = self.router(h0)  # [B, T, n_buckets]

        bucket_outputs = []
        for bucket_idx, n_recursions in enumerate(self.cfg.bucket_recursions):
            h = h0
            for _ in range(n_recursions):
                for layer in self.archon.layers:
                    h = layer(h)
            bucket_outputs.append(h)  # [B, T, D]

        # Mix: weighted sum
        stacked = torch.stack(bucket_outputs, dim=-2)  # [B, T, n_buckets, D]
        weights = bucket_probs.unsqueeze(-1)  # [B, T, n_buckets, 1]
        h_mixed = (stacked * weights).sum(dim=-2)  # [B, T, D]

        h_mixed = self.archon.norm(h_mixed)
        logits = self.archon.lm_head(h_mixed)

        # Aux balance loss: encourage uniform bucket usage
        aux = -(bucket_probs * (bucket_probs + 1e-9).log()).sum(dim=-1).mean()
        aux_loss = -self.cfg.aux_loss_weight * aux  # maximize entropy

        loss = None
        if targets is not None:
            ntp_loss = F.cross_entropy(
                logits[:, :-1].reshape(-1, logits.shape[-1]),
                targets[:, 1:].reshape(-1),
                ignore_index=-100,
            )
            loss = ntp_loss + aux_loss

        return logits, loss, []


def router_training_plan_doc() -> str:
    """Plan for MoR router training (Phase D)."""
    return """
    MoR ROUTER TRAINING:
    1. Start from ARCHON SFT v2 final ckpt
    2. Add router (init: uniform softmax over 3 buckets)
    3. Train router only (freeze backbone) for 3K steps on SFT mix
       - LR 1e-4, batch=16, V100, ~6h
    4. Then unfreeze, joint fine-tune backbone+router for 5K steps
       - LR 1e-5, distillation loss vs ARCHON-SFT-v2 outputs
       - V100 ~10h, $7
    5. Eval: GSM8K accuracy + average recursion depth
       Target: depth_avg ≈ 1.8, accuracy >= ARCHON-SFT-v2 baseline
    6. Killer: if router AUC < 0.65 OR depth_avg < 1.3 → abandon

    Combo bonus: ARCHON-MoR + COCONUT (M3) + TTT (M12) = trio inédit,
    soumettre brevet patent-scout.
    """


if __name__ == "__main__":
    cfg = MoRConfig()
    print(f"[M4 MoR] config: {cfg}, buckets: {cfg.bucket_recursions}")
    print(router_training_plan_doc())