"""CPU mechanism experiment; untrained NumPy decoder, not a paper benchmark."""
import argparse
import json
import platform
import time
from pathlib import Path

import numpy as np


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


def norm(x):
    return (x - x.mean(-1, keepdims=True)) / np.sqrt(x.var(-1, keepdims=True) + 1e-5)


class Decoder:
    """Batch=1, pre-norm causal MHA + tanh FFN, learned absolute positions."""
    def __init__(self, seed=69, layers=2):
        r = np.random.default_rng(seed)
        self.layers, self.h, self.d, self.vocab = layers, 2, 16, 8
        self.emb = r.normal(0, .3, (8, 16))
        self.pos = r.normal(0, .1, (256, 16))
        self.weights = [tuple(r.normal(0, 1 / np.sqrt(a), (a, b))
                             for a, b in [(16, 16)] * 4 + [(16, 32), (32, 16)])
                        for _ in range(layers)]
        self.head = r.normal(0, .25, (16, 8))
        self.reset_counts()

    def reset_counts(self):
        self.calls = self.projected_tokens = self.score_cells = 0

    def forward(self, ids, cache=None):
        cache = [None] * self.layers if cache is None else cache
        old = 0 if cache[0] is None else cache[0][0].shape[1]
        n = len(ids)
        assert n > 0 and old + n <= 256
        x = self.emb[np.asarray(ids)] + self.pos[old:old + n]
        allowed = np.arange(old + n)[None, :] <= (old + np.arange(n))[:, None]
        new = []
        for w, previous in zip(self.weights, cache):
            z = norm(x)
            q, k, v = [(z @ a).reshape(n, self.h, 8).transpose(1, 0, 2) for a in w[:3]]
            if previous is not None:
                k, v = np.concatenate([previous[0], k], 1), np.concatenate([previous[1], v], 1)
            scores = (q @ k.transpose(0, 2, 1)) / np.sqrt(8)
            attn = softmax(np.where(allowed[None], scores, -np.inf))
            x = x + (attn @ v).transpose(1, 0, 2).reshape(n, 16) @ w[3]
            x = x + np.tanh(norm(x) @ w[4]) @ w[5]
            new.append((k, v))
        self.calls += 1
        self.projected_tokens += n * self.layers
        self.score_cells += self.layers * self.h * n * (old + n)
        return norm(x) @ self.head, new


def crop(cache, length):
    # Copy to release rejected suffix allocations; views would retain backing memory.
    return [(k[:, :length].copy(), v[:, :length].copy()) for k, v in cache]


def cache_len(cache):
    return cache[0][0].shape[1]


def assert_cache(model, seq, cache):
    _, reference = model.forward(seq)
    for actual, expected in zip(cache, reference):
        for a, b in zip(actual, expected):
            np.testing.assert_allclose(a, b, rtol=1e-10, atol=1e-10)


def distribution(logits, greedy):
    if greedy:
        return np.eye(logits.shape[-1])[np.argmax(logits, axis=-1)]
    return softmax(logits)  # temperature=1, no top-k/p or repetition penalty


def corrected(p, q):
    residual = np.maximum(p - q, 0)
    assert residual.sum() > 0, "Zero residual cannot occur on a genuine rejection"
    return residual / residual.sum()


def baseline(model, prefix, count, use_cache=True, greedy=True, seed=0):
    seq, cache, r = list(prefix), None, np.random.default_rng(seed)
    for _ in range(count):
        ids = seq if cache is None or not use_cache else seq[-1:]
        logits, new = model.forward(ids, cache if use_cache else None)
        cache = new if use_cache else None
        seq.append(int(r.choice(model.vocab, p=distribution(logits[-1], greedy))))
    return seq


def speculative(target, draft, prefix, count, gamma=4, greedy=False, seed=0, audit=False):
    """Caches retain committed prefix excluding its last token at each round end."""
    seq, rng = list(prefix), np.random.default_rng(seed)
    assert len(seq) >= 2 and count > 0 and gamma >= 1
    _, tc = target.forward(seq[:-1])
    _, dc = draft.forward(seq[:-1])
    stop = len(seq) + count
    stats = dict(rounds=0, proposed=0, examined=0, accepted=0, rejected_rounds=0, full_rounds=0)
    trace = []
    while len(seq) < stop:
        start = len(seq)
        # Reserve one position for correction/bonus; gamma can be zero in final round.
        g = min(gamma, stop - start - 1)
        trial, qs = list(seq), []
        for _ in range(g):
            logits, dc = draft.forward(trial[cache_len(dc):], dc)
            q = distribution(logits[-1], greedy)
            qs.append(q)
            trial.append(int(rng.choice(target.vocab, p=q)))
        logits, tc = target.forward(trial[cache_len(tc):], tc)
        ps = distribution(logits, greedy)
        assert len(ps) == g + 1
        if audit:
            ref, _ = target.forward(trial)
            np.testing.assert_allclose(logits, ref[start - 1:], rtol=1e-10, atol=1e-10)
        accepted = 0
        next_p = ps[-1]
        for i in range(g):
            token = trial[start + i]
            stats['examined'] += 1
            assert qs[i][token] > 0
            if rng.random() < min(1., ps[i, token] / qs[i][token]):
                accepted += 1
            else:
                next_p = corrected(ps[i], qs[i])
                break
        bonus = int(rng.choice(target.vocab, p=next_p))
        seq = seq + trial[start:start + accepted] + [bonus]
        tc = crop(tc, len(seq) - 1)
        # After all accepted, draft may lag one additional token; next call catches up.
        dc = crop(dc, min(cache_len(dc), len(seq) - 1))
        if audit:
            assert_cache(target, seq[:-1], tc)
            assert_cache(draft, seq[:cache_len(dc)], dc)
        stats['rounds'] += 1
        stats['proposed'] += g
        stats['accepted'] += accepted
        stats['rejected_rounds' if accepted < g else 'full_rounds'] += 1
        trace.append(dict(prefix_length=start, proposed=g, accepted=accepted,
                          target_cache_length=cache_len(tc), draft_cache_length=cache_len(dc)))
    assert len(seq) == stop
    return seq, stats, trace


def sampling_audit():
    """Analytic mass conservation plus Monte Carlo; wrong fallback is a negative control."""
    pairs = [([.1, .6, .3], [.7, .2, .1]), ([.2, .3, .5], [.2, .3, .5]),
             ([0, 1, 0], [1, 0, 0]), ([.4, .6, 0], [0, .5, .5])]
    errors = []
    for p, q in pairs:
        p, q = np.array(p), np.array(q)
        accepted_mass = np.minimum(p, q)
        remainder = np.maximum(p - q, 0)
        law = accepted_mass + remainder
        errors.append(float(np.max(np.abs(law - p))))
        np.testing.assert_allclose(law, p, atol=1e-15)
    p, q = np.array(pairs[0][0]), np.array(pairs[0][1])
    rng, n = np.random.default_rng(6900), 200000
    candidates = rng.choice(3, n, p=q)
    keep = rng.random(n) < np.minimum(1, p[candidates] / q[candidates])
    actual, wrong = candidates.copy(), candidates.copy()
    actual[~keep] = rng.choice(3, (~keep).sum(), p=corrected(p, q))
    wrong[~keep] = rng.choice(3, (~keep).sum(), p=p)
    freq = np.bincount(actual, minlength=3) / n
    wrong_freq = np.bincount(wrong, minlength=3) / n
    tv = float(np.abs(freq - p).sum() / 2)
    wrong_tv = float(np.abs(wrong_freq - p).sum() / 2)
    wrong_law = np.minimum(p, q) + (1 - np.minimum(p, q).sum()) * p
    wrong_theoretical_tv = float(np.abs(wrong_law - p).sum() / 2)
    # This pair has exact wrong-law TV=.06; require MC to recover that law.
    np.testing.assert_allclose(wrong_freq, wrong_law, atol=.01)
    assert tv < .01 and wrong_tv > .04
    return dict(draws=n, seed=6900, p=p.tolist(), q=q.tolist(), frequency=freq.tolist(),
                tv=tv, wrong_fallback_frequency=wrong_freq.tolist(), wrong_fallback_tv=wrong_tv,
                wrong_theoretical_tv=wrong_theoretical_tv,
                analytic_max_error=max(errors), analytic_cases=len(pairs))


def main(out):
    out.mkdir(parents=True, exist_ok=True)
    prefix = [0, 1, 3, 2, 4, 1, 5, 6]
    t, d = Decoder(), Decoder(layers=1)
    ids = (prefix * 4)[:29]
    full, full_cache = t.forward(ids)
    cache, parts, start = None, [], 0
    for size in [3, 1, 7, 2, 16]:
        logits, cache = t.forward(ids[start:start + size], cache)
        parts.append(logits)
        start += size
    err = float(np.max(np.abs(full - np.concatenate(parts))))
    np.testing.assert_allclose(full, np.concatenate(parts), rtol=1e-10, atol=1e-10)
    assert_cache(t, ids, cache)
    cache_bytes = sum(a.nbytes for pair in full_cache for a in pair)
    assert cache_bytes == 2 * 2 * 2 * 29 * 8 * 8
    # Detect accidental future visibility by replacing only the final token.
    altered, _ = t.forward(ids[:-1] + [(ids[-1] + 1) % 8])
    np.testing.assert_allclose(full[:-1], altered[:-1], atol=1e-12)
    checks, traces = [], []
    for seed in [69, 70, 71]:
        t, d = Decoder(seed), Decoder(seed, layers=1)
        expected = baseline(t, prefix, 48, use_cache=False)
        assert expected == baseline(t, prefix, 48, use_cache=True)
        for gamma in [1, 4, 8]:
            actual, stats, trace = speculative(t, d, prefix, 48, gamma, True, seed, True)
            assert actual == expected
            checks.append(dict(seed=seed, gamma=gamma, greedy_equal=True, **stats))
        _, stats, trace = speculative(t, d, prefix, 48, 4, False, seed, True)
        traces.append(dict(seed=seed, mode='sampling', stats=stats, rounds=trace))
    # Identity draft exercises full acceptance and bonus; disjoint p/q tested separately.
    t, same = Decoder(), Decoder()
    _, identity, _ = speculative(t, same, prefix, 12, 4, False, 69, True)
    assert identity['rejected_rounds'] == 0
    t, d = Decoder(), Decoder(layers=1)
    jobs = [('full_prefix', lambda: baseline(t, prefix, 48, False)),
            ('kv_cache', lambda: baseline(t, prefix, 48, True)),
            ('speculative_g4', lambda: speculative(t, d, prefix, 48, 4, True, 69))]
    times, costs = {}, {}
    for name, fn in jobs:
        fn()  # warmup, excluded
        samples = []
        for _ in range(7):
            t.reset_counts(); d.reset_counts()
            begin = time.perf_counter()
            fn()
            samples.append((time.perf_counter() - begin) * 1000)
        times[name] = dict(ms=samples, median_ms=float(np.median(samples)))
        costs[name] = dict(target_calls=t.calls, draft_calls=d.calls,
                           target_layer_tokens=t.projected_tokens, target_score_cells=t.score_cells)
    result = dict(environment=dict(python=platform.python_version(), numpy=np.__version__,
                                   platform=platform.platform(), dtype='float64', device='CPU'),
                  config=dict(prefix=prefix, new_tokens=48, vocab=8, dim=16, heads=2,
                              target_layers=2, draft_layers=1, trained=False),
                  kv=dict(max_logit_error=err, shape=list(full_cache[0][0].shape),
                          bytes=cache_bytes, cached_length=29, causal_invariance=True),
                  greedy_checks=checks, sampling_cache_traces=traces, identity_draft=identity,
                  distribution_check=sampling_audit(), timing=times, workload=costs)
    (out / 'results.json').write_text(json.dumps(result, indent=2) + '\n')
    print(json.dumps({k: v for k, v in result.items() if k != 'sampling_cache_traces'}, indent=2))
    print('PASS: cache/block equivalence, causal masking, greedy equality, rollback, residual sampling')


if __name__ == '__main__':
    ap = argparse.ArgumentParser()
    ap.add_argument('--out', type=Path, default=Path(__file__).resolve().parent)
    main(ap.parse_args().out)