import json, math, random
from itertools import product
import numpy as np

SEED = 3040
random.seed(SEED); np.random.seed(SEED)


def majority(x):
    return 1 if sum(x) >= 0 else -1


def mutual_info(d, mapping):
    q = {}
    for (l, r), v in d.items():
        x = mapping(l); q[x, r] = q.get((x, r), 0.0) + v
    px, py = {}, {}
    for (x, y), v in q.items():
        px[x] = px.get(x, 0) + v; py[y] = py.get(y, 0) + v
    return sum(v * math.log2(v / (px[x] * py[y])) for (x, y), v in q.items())


def exact_check(p=.15):
    # Z is a latent remote spin; L is a three-spin local block and R is a
    # distant noisy observation. This is the paper's majority failure mode.
    d = {}
    for z in (-1, 1):
        for l in product((-1, 1), repeat=3):
            for r in (-1, 1):
                v = .5
                v *= np.prod([(1-p) if x == z else p for x in l])
                v *= (1-p) if r == z else p
                d[l, r] = d.get((l, r), 0) + v
    full = mutual_info(d, lambda x: x)
    maj = mutual_info(d, majority)
    # Deterministic pooling gives I(L:R|M)=I(L:R)-I(M:R).
    post = {}
    for l in ((1, 1, 1), (1, 1, -1)):
        a = (1-p)**sum(x == 1 for x in l) * p**sum(x == -1 for x in l)
        b = p**sum(x == 1 for x in l) * (1-p)**sum(x == -1 for x in l)
        post[str(l)] = a / (a+b)
    # Unequal positional weights are injective over the eight binary blocks.
    code = lambda x: round(sum(w*y for w, y in zip((.7, .2, .1), x)), 8)
    return {'p': p, 'I(L:R)': full, 'I(majority(L):R)': maj,
            'discarded_conditional_information_bits': full-maj,
            'I(injective_pool(L):R)': mutual_info(d, code),
            'P(Z=+1|L)': post}


def neural_demo(p=.15, n=16000):
    import torch
    torch.manual_seed(SEED)
    requested = 'cuda' if torch.cuda.is_available() else 'cpu'
    try:
        dev = torch.device(requested)
        z = torch.where(torch.rand(n, device=dev) < .5, -torch.ones(n, device=dev), torch.ones(n, device=dev))
        l = torch.where(torch.rand(n, 3, device=dev) < p, -z[:, None], z[:, None]).float()
        r = torch.where(torch.rand(n, device=dev) < p, -z, z).float()
        cut = n * 3 // 4
        tr, te = torch.arange(cut, device=dev), torch.arange(cut, n, device=dev)

        def train_pool(kind):
            if kind == 'mean':
                s = l.mean(1, keepdim=True); params = []
            elif kind == 'majority':
                s = torch.sign(l.sum(1)).reshape(-1, 1); s[s == 0] = 1; params = []
            else:
                # Slightly asymmetric initialization breaks the permutation
                # symmetry; unlike majority, the scalar can retain all 3 bits.
                logits = torch.nn.Parameter(torch.log(torch.tensor([.7,.2,.1], device=dev)))
                params = [logits]
                s = None
            head = torch.nn.Sequential(torch.nn.Linear(1, 24), torch.nn.Tanh(),
                                       torch.nn.Linear(24, 1)).to(dev)
            params += list(head.parameters())
            opt = torch.optim.Adam(params, lr=.025)
            y = (r + 1) / 2
            for _ in range(900):
                if kind == 'learned':
                    a = torch.softmax(logits, 0); s = (l*a).sum(1, keepdim=True)
                loss = torch.nn.functional.binary_cross_entropy_with_logits(head(s[tr]).squeeze(1), y[tr])
                opt.zero_grad(); loss.backward(); opt.step()
            with torch.no_grad():
                if kind == 'learned':
                    a = torch.softmax(logits, 0); s = (l*a).sum(1, keepdim=True)
                pred = head(s[te]).squeeze(1)
                nll = torch.nn.functional.binary_cross_entropy_with_logits(pred, y[te]).item()
                acc = ((pred > 0) == (r[te] > 0)).float().mean().item()
                result = {'heldout_remote_accuracy': acc, 'heldout_remote_nll': nll}
                if kind == 'learned': result['weights'] = a.detach().cpu().numpy().tolist()
                return result

        # q_phi has the microscopic source block; q_psi has only the pooled S.
        # Their held-out NLL difference is the empirical information gap.
        aux = torch.nn.Sequential(torch.nn.Linear(3, 24), torch.nn.Tanh(), torch.nn.Linear(24,1)).to(dev)
        opt = torch.optim.Adam(aux.parameters(), lr=.025); y=(r+1)/2
        for _ in range(900):
            loss=torch.nn.functional.binary_cross_entropy_with_logits(aux(l[tr]).squeeze(1), y[tr])
            opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            nll_full=torch.nn.functional.binary_cross_entropy_with_logits(aux(l[te]).squeeze(1),y[te]).item()
        out = {k: train_pool(k) for k in ('mean','majority','learned')}
        out['full_source_predictor_nll'] = nll_full
        out['mean_predictor_gap_nll'] = out['mean']['heldout_remote_nll'] - nll_full
        out['majority_predictor_gap_nll'] = out['majority']['heldout_remote_nll'] - nll_full
        out['learned_predictor_gap_nll'] = out['learned']['heldout_remote_nll'] - nll_full
        out['device'] = str(dev)
        return out
    except Exception as e:
        return {'device': 'cpu-fallback', 'error': repr(e)}


if __name__ == '__main__':
    print(json.dumps({'exact_check': exact_check(), 'neural_demo': neural_demo()}, indent=2))