import json, math, random
from pathlib import Path
import numpy as np

SEED = 2951
rng = np.random.default_rng(SEED)


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


def entropy(p):
    return float(-(p * np.log(np.maximum(p, 1e-12))).sum(axis=-1).mean())


def math_check(R=5, D=7, N=9):
    # QR gives exactly orthonormal role vectors (up to numerical precision).
    Q, _ = np.linalg.qr(rng.normal(size=(D, R)))
    roles = Q.T                         # [R,D], rows orthonormal
    F = rng.normal(size=(N, R, D))
    O = np.einsum('rd,nre->nde', roles, F)  # [N,D,D]
    recovered = np.einsum('rd,nde->nre', roles, O)
    contraction_error = float(np.max(np.abs(recovered - F)))

    # Query role m and filler from object 3 should retrieve object 3.
    m, target, source = 1, 3, 3
    qf = F[source, m]
    scores = F[:, m] @ qf
    probs = softmax(scores / 0.15)
    top = int(np.argmax(probs))
    out = probs @ F[:, target]
    retrieval_error = float(np.linalg.norm(out - F[source, target]))

    # Wrong-role query is a control: it should not be systematically selective.
    wrong = (m + 1) % R
    wrong_probs = softmax((F[:, wrong] @ qf) / 0.15)
    return {
        'max_role_contraction_abs_error': contraction_error,
        'exact_query_top1': top == source,
        'exact_query_probability': float(probs[source]),
        'exact_rebinding_l2_error': retrieval_error,
        'exact_entropy': entropy(probs[None, :]),
        'mismatched_role_entropy': entropy(wrong_probs[None, :]),
    }


def make_episode(N=16, R=4, D=16, noise=0.08):
    # Each object has independent role fillers. The query is one role filler;
    # the desired answer is another role filler from that same object.
    F = rng.normal(size=(N, R, D)).astype(np.float32)
    src = int(rng.integers(N))
    m, target = 0, 1
    q = F[src, m] + noise * rng.normal(size=D).astype(np.float32)
    return F, q, src, m, target


def structured_trial(F, q, m, target, tau=0.25):
    s = F[:, m] @ q
    a = softmax(s / tau)
    y = a @ F[:, target]
    return int(np.argmax(s)), y, a


def dense_trial(F, q, target, tau=0.25):
    # Control: a standard flattened-object dot product, with the query copied
    # into every role block so it has the same input dimension as an object.
    flat = F.reshape(F.shape[0], -1)
    qflat = np.tile(q, F.shape[1])
    s = flat @ qflat
    a = softmax(s / tau)
    return int(np.argmax(s)), a @ F[:, target], a


def benchmark(episodes=2000):
    stats = {k: [] for k in ['structured_acc','dense_acc','structured_err','dense_err',
                             'structured_entropy','dense_entropy','wrong_entropy']}
    for _ in range(episodes):
        F, q, src, m, target = make_episode()
        si, sy, sa = structured_trial(F, q, m, target)
        di, dy, da = dense_trial(F, q, target)
        wrong = (m + 1) % F.shape[1]
        _, _, wa = structured_trial(F, q, wrong, target)
        stats['structured_acc'].append(si == src)
        stats['dense_acc'].append(di == src)
        stats['structured_err'].append(np.linalg.norm(sy - F[src,target]))
        stats['dense_err'].append(np.linalg.norm(dy - F[src,target]))
        stats['structured_entropy'].append(entropy(sa[None,:]))
        stats['dense_entropy'].append(entropy(da[None,:]))
        stats['wrong_entropy'].append(entropy(wa[None,:]))
    return {k: float(np.mean(v)) for k,v in stats.items()}


def main():
    np.set_printoptions(precision=6, suppress=True)
    result = {'seed': SEED, 'math_check': math_check(), 'benchmark': benchmark()}
    # Repeat the cheap benchmark with a different seed to check direction stability.
    global rng
    rng = np.random.default_rng(SEED + 1)
    result['benchmark_repeat'] = benchmark(episodes=1000)
    out = Path('results.json')
    out.write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    main()