"""Sparse experts, recurrent state, and user-scoped application memory. Requires torch."""
from __future__ import annotations
from collections import OrderedDict
import torch
from torch import nn


class MoE(nn.Module):
    """Pre: independent rows (N,D), 1<=k<=experts. Post: sparse weighted FFN output."""
    def __init__(self, width=8, experts=4, k=2):
        super().__init__()
        if not 1 <= k <= experts:
            raise ValueError("k must be between one and the number of experts")
        self.k, self.count = k, experts
        self.router = nn.Linear(width, experts)
        self.experts = nn.ModuleList([
            nn.Sequential(nn.Linear(width, 16), nn.ReLU(), nn.Linear(16, width))
            for _ in range(experts)
        ])

    def forward(self, inputs):
        probabilities = self.router(inputs).softmax(-1)
        weights, selected = probabilities.topk(self.k, dim=-1)
        weights = weights/weights.sum(-1, keepdim=True)
        result = torch.zeros_like(inputs)
        # Invariant: result contains contributions of earlier experts only.
        # Bound: remaining experts; each receives only its selected rows.
        for index, expert in enumerate(self.experts):
            rows, slots = (selected == index).nonzero(as_tuple=True)
            if rows.numel():
                contribution = weights[rows, slots, None] * expert(inputs[rows])
                result = result.index_add(0, rows, contribution)
        return result, probabilities, selected


def balancing_loss(probabilities, selected, experts):
    """Post: Switch-style auxiliary statistic; uniform routing gives 1, not a universal minimum."""
    fraction = torch.bincount(selected.reshape(-1), minlength=experts).float()/selected.numel()
    return experts * (fraction*probabilities.mean(0)).sum()


def recurrence(values, *, decay=.5, state=0.):
    """Pre: 0<decay<1, finite inputs. Post: outputs/state of h_t=decay*h_(t-1)+x_t."""
    if not 0 < decay < 1:
        raise ValueError("decay must be strictly between zero and one")
    output = []
    # Invariant: state summarizes the processed prefix. Bound: remaining input values.
    for value in values:
        state = decay*state + value
        output.append(state)
    return output, state


class MemoryBank:
    """An application-level LRU store per user; writes are explicit, not learned neural gates."""
    def __init__(self, capacity=2):
        if capacity < 1:
            raise ValueError("capacity must be positive")
        self.capacity, self.users = capacity, {}

    def write(self, user, key, value):
        store = self.users.setdefault(user, OrderedDict())
        store[key] = value
        store.move_to_end(key)
        if len(store) > self.capacity:
            store.popitem(last=False)

    def read(self, user, key):
        store = self.users.get(user, {})
        if key not in store:
            return None
        store.move_to_end(key)
        return store[key]

    def reset(self, user):
        self.users.pop(user, None)


def train(*, steps=120):
    """Post: MoE trained to approximate 2*x on a reproducible synthetic batch."""
    if steps <= 0:
        raise ValueError("steps must be positive")
    torch.set_num_threads(1)
    torch.manual_seed(7)
    inputs = torch.randn(128, 8)
    target = 2*inputs
    model = MoE()
    optimizer = torch.optim.Adam(model.parameters(), lr=.01)
    history = []
    # Invariant: history records each completed sparse update. Bound: remaining steps.
    for _ in range(steps):
        optimizer.zero_grad(set_to_none=True)
        output, probabilities, selected = model(inputs)
        fit = (output-target).square().mean()
        loss = fit + .01*balancing_loss(probabilities, selected, model.count)
        loss.backward()
        optimizer.step()
        history.append(fit.item())
    return model, history


def demo():
    model, history = train()
    print("MoE initial/final regression MSE:", history[0], history[-1])
    print("Scalar state for [1,0,0,2]:", recurrence([1, 0, 0, 2]))
    bank = MemoryBank()
    bank.write("Ada", "language", "Python")
    print("Ada's value:", bank.read("Ada", "language"))
    print("Another user's value:", bank.read("Grace", "language"))
    p = torch.tensor([[.51, .49]]*9 + [[.01, .99]])
    print("Valid top-1 routing statistic below 1:", balancing_loss(p, p.argmax(1)[:, None], 2).item())
    print("These primitives are educational; the recurrence is not a full Mamba implementation.")


def self_test():
    outputs, state = recurrence([1, 0, 0, 2])
    check(outputs, [1., .5, .25, 2.125])
    first, state = recurrence([1, 0])
    second, state = recurrence([0, 2], state=state)
    check(first+second, outputs)
    bank = MemoryBank(capacity=1)
    bank.write("a", "x", 1)
    bank.write("a", "y", 2)
    check(bank.read("a", "x"), None)
    check(bank.read("b", "y"), None)
    bank.reset("a")
    check(bank.read("a", "y"), None)
    p = torch.tensor([[.51, .49]]*9 + [[.01, .99]])
    check(balancing_loss(p, p.argmax(1)[:, None], 2).item() < 1, True)
    model, history = train()
    check(history[-1] < history[0]/10, True)
    check(model.router.weight.grad is not None, True)
    return 8

def check(actual, expected):
    """Require an expected result; this check also runs under python -O."""
    if actual != expected:
        raise AssertionError(f"expected {expected!r}, got {actual!r}")


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args()
    if args.self_test:
        print(f"{self_test()} checks passed")
    else:
        demo()