"""前沿架构：手算规模的专家路由、线性注意力与记忆更新；不是语言模型训练。"""

from argparse import ArgumentParser
import hashlib
import json
from pathlib import Path
import platform

import numpy as np


def softmax(values):
    weights = np.exp(values - values.max(axis=-1, keepdims=True))
    return weights / weights.sum(axis=-1, keepdims=True)


def expert(x, first, second):
    return np.maximum(x @ first, 0) @ second


def routed(x, router, first, second, top_k=2, capacity=None):
    """按词元顺序分配；超过容量就丢弃该次分配，不重新归一化。"""
    n_experts = len(first)
    if not 1 <= top_k <= n_experts:
        raise ValueError("top_k 必须位于专家数量范围内")
    if capacity is not None and capacity < 0:
        raise ValueError("容量不能为负")
    ids = np.argsort(-(x @ router), axis=-1, kind="stable")[:, :top_k]
    scores = np.take_along_axis(x @ router, ids, axis=1)
    weights = softmax(scores)
    output = np.zeros_like(x)
    accepted = np.zeros(n_experts, dtype=int)
    requested = np.bincount(ids.ravel(), minlength=n_experts)
    dropped = 0
    for token in range(len(x)):
        for slot, index in enumerate(ids[token]):
            if capacity is not None and accepted[index] >= capacity:
                dropped += 1
                continue
            output[token] += weights[token, slot] * expert(x[token], first[index], second[index])
            accepted[index] += 1
    return output, ids, weights, requested, accepted, dropped


def linear_recurrent(q, k, v):
    """q/k 已是非负特征；写入当前位置后读取，包含当前位置的因果注意力。"""
    state = np.zeros((k.shape[1], v.shape[1]))
    normalizer = np.zeros(k.shape[1])
    outputs = []
    for query, key, value in zip(q, k, v):
        state += np.outer(key, value)
        normalizer += key
        denominator = query @ normalizer
        if denominator <= 0:
            raise ValueError("当前查询没有正的匹配权重")
        outputs.append(query @ state / denominator)
    return np.array(outputs), state, normalizer


def linear_explicit(q, k, v):
    similarity = np.tril(q @ k.T)
    return similarity @ v / similarity.sum(axis=1, keepdims=True)


def delta_memory(k, v, decay, write):
    """S 的形状是 d_k × d_v；门值为手工给定，没有训练投影或门。"""
    state = np.zeros((k.shape[1], v.shape[1]))
    history = []
    for key, value, alpha, beta in zip(k, v, decay, write):
        state = alpha * state
        error = value - key @ state
        state = state + beta * np.outer(key, error)
        history.append(state.copy())
    return np.array(history)


def scalar_state(inputs, decay):
    state, history = 0., []
    for value in inputs:
        state = decay * state + value
        history.append(state)
    return history


def run():
    rng = np.random.default_rng(7)
    x = np.array([[1 + i / 10, i / 10] for i in range(6)])
    router = np.array([[2., 1., -1., -2.], [0., .2, .1, 0.]])
    first = rng.normal(size=(4, 2, 3))
    second = rng.normal(size=(4, 3, 2))
    output, ids, weights, requested, accepted, dropped = routed(x, router, first, second)
    # 全算后遮罩仅作独立执行路径对照，不作为“稀疏省计算”的实现。
    all_outputs = np.stack([expert(x, a, b) for a, b in zip(first, second)], axis=1)
    mask = np.zeros((len(x), len(first)))
    np.put_along_axis(mask, ids, weights, axis=1)
    reference = np.einsum("ne,ned->nd", mask, all_outputs)
    limited, _, _, _, limited_load, limited_drops = routed(
        x, router, first, second, capacity=3)

    q = rng.uniform(.1, 1., size=(7, 3))
    k = rng.uniform(.1, 1., size=(7, 3))
    v = rng.normal(size=(7, 2))
    recurrent, _, _ = linear_recurrent(q, k, v)
    explicit = linear_explicit(q, k, v)
    scores = q @ k.T
    scores[np.triu_indices(len(q), 1)] = -np.inf
    conventional = softmax(scores) @ v  # 单独对照 exp(dot)，未加缩放。

    keys = np.array([[1., 0.], [0., 1.], [1., 0.]])
    values = np.array([[10.], [20.], [30.]])
    _, state, normalizer = linear_recurrent(keys, keys, values)
    queries = np.eye(2)
    averages = (queries @ state) / (queries @ normalizer)[:, None]
    delta = delta_memory(keys, values, [1., 1., 1.], [1., 1., 1.])
    decayed = delta_memory(keys, values, [1., 1., .5], [1., 1., 1.])
    mixed_keys = np.array([[1., 0.], [2 ** -.5, 2 ** -.5]])
    collision = delta_memory(mixed_keys, np.array([[10.], [30.]]), [1., 1.], [1., 1.])
    return {
        "kind": "untrained_synthetic_mechanism_checks",
        "seed": 7,
        "moe": {
            "inputs": x.tolist(), "router": router.tolist(),
            "first": first.tolist(), "second": second.tolist(),
            "ids": ids.tolist(), "weights": weights.tolist(), "output": output.tolist(),
            "requested": requested.tolist(), "accepted": accepted.tolist(),
            "calls": int(accepted.sum()), "dense_calls": int(len(x) * len(first)),
            "total_parameters": int(first.size + second.size + router.size),
            "parameters_used_per_token": int(2 * first[0].size + 2 * second[0].size + router.size),
            "reference_max_error": float(np.max(np.abs(output - reference))),
            "capacity": 3, "limited_load": limited_load.tolist(),
            "dropped_assignments": limited_drops, "limited_output": limited.tolist(),
        },
        "linear": {
            "q": q.tolist(), "k": k.tolist(), "v": v.tolist(),
            "recurrent": recurrent.tolist(), "explicit": explicit.tolist(),
            "max_error": float(np.max(np.abs(recurrent - explicit))),
            "softmax_max_difference": float(np.max(np.abs(recurrent - conventional))),
        },
        "memory": {
            "keys": keys.tolist(), "values": values.tolist(),
            "average_read": averages.ravel().tolist(),
            "delta_read": (queries @ delta[-1]).ravel().tolist(),
            "decayed_read": (queries @ decayed[-1]).ravel().tolist(),
            "collision_read": float(collision[-1, 0, 0]),
            "impulse_traces": {str(a): scalar_state([1.] + [0.] * 50, a) for a in [.8, .99]},
            "cache_counts": [{"tokens": n, "kv_scalars": n * 3,
                              "normalized_linear_scalars": int(state.size + normalizer.size)}
                             for n in [32, 128]],
        },
    }


def main():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--out", type=Path, help="将完整机制算例保存为新 JSON 文件")
    args = parser.parse_args()
    report = run()
    report["environment"] = {"python": platform.python_version(), "numpy": np.__version__,
                             "code_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest()}
    if args.out:
        args.out.parent.mkdir(parents=True, exist_ok=True)
        with args.out.open("x", encoding="utf-8") as file:
            json.dump(report, file, ensure_ascii=False, indent=2, allow_nan=False)
            file.write("\n")
    moe, linear, memory = report["moe"], report["linear"], report["memory"]
    print(f'专家请求数：{moe["requested"]}；实际调用 {moe["calls"]} 次，全算需 {moe["dense_calls"]} 次。')
    print(f'总参数 {moe["total_parameters"]}；每个词元使用 {moe["parameters_used_per_token"]} 个；稀疏与遮罩路径一致：{moe["reference_max_error"] < 1e-12}')
    print(f'容量为 3 时接收：{moe["limited_load"]}；丢弃 {moe["dropped_assignments"]} 次分配。')
    print(f'线性注意力两种计算一致：{linear["max_error"] < 1e-12}；与 softmax 对照最大差：{linear["softmax_max_difference"]:.6f}')
    print(f'读取 A/B：平均 {memory["average_read"]}；delta {memory["delta_read"]}；末步衰减 {memory["decayed_read"]}')
    print(f'非正交键更新后再读 A：{memory["collision_read"]:.6f}（原值 10）。')
    print(f'50 步后旧脉冲：0.8 门 {memory["impulse_traces"]["0.8"][-1]:.6f}；0.99 门 {memory["impulse_traces"]["0.99"][-1]:.6f}')
    print("范围：固定输入与未训练参数的机制检查，没有测量语言能力或端到端速度。")


if __name__ == "__main__":
    main()

1. 不能。不同词元可能选不同专家；权重仍需存在于某处，设备布局与搬运另有成本。
2. 调用计数没有包含路由、调度、通信与硬件执行效率，也没有进行计时。
3. A 的两次值为 10 与 30，匹配权重相同，平均为 20。
4. 旧值为 10，差额为 20，写一半后得到 20；B 仍为 20。
5. 衰减作用于整份旧状态。只在 A 方向写回差额，不会补回 B 丢失的部分。
6. 不能。保存全部 KV 的层仍随长度增加；应按不同层的状态形式求和。
7. 容量 6 足以接收当前全部请求，应恢复无容量限制时的输出。改变路由后要重新核对选择和负载，不能只看总次数。
8. 新键先读到 $10/\sqrt{2}$，差额为 $30-10/\sqrt{2}$。A 方向的新值为 $10+(30-10/\sqrt{2})/\sqrt{2}=5+15\sqrt{2}$，约 26.213203。它否定了这个更新规则能对任意非正交键都无干扰改写的说法。