#!/usr/bin/env python3
"""Finite checks for the communication results in the manuscript.

The checks do not prove the asymptotic theorems. They test the two-population
information inequality, finite binary codes, quantized pairwise-majority
receivers, and scalar likelihood-swap algebra. No simulation is presented as
an end-to-end physical implementation or as a general Byzantine protocol.

Run from any directory: python code/verify_communication.py
Requires Python >= 3.10, NumPy and SciPy.
"""
from __future__ import annotations
import json
import math
from pathlib import Path
import numpy as np
from scipy.special import xlogy

SEED = 20261028
ROOT = Path(__file__).resolve().parents[1]


def entropy(p: np.ndarray | float) -> np.ndarray:
    p = np.asarray(p, dtype=float)
    return -xlogy(p, p) - xlogy(1-p, 1-p)


def ber_kl(p: float, q: float) -> float:
    if not (0 <= p <= 1 and 0 < q < 1):
        raise ValueError('Invalid Bernoulli parameters')
    return float(xlogy(p, p/q) + xlogy(1-p, (1-p)/(1-q)))


def two_population_information(q: np.ndarray, r: np.ndarray) -> np.ndarray:
    """Rows index U, columns index V; q and r are P(output=1|U,V)."""
    ia = np.mean(entropy(np.mean(q, axis=-1)) - np.mean(entropy(q), axis=-1), axis=-1)
    ib = np.mean(entropy(np.mean(r, axis=-2)) - np.mean(entropy(r), axis=-2), axis=-1)
    return ia + ib


def information_checks(rng: np.random.Generator) -> list[dict]:
    ans = []
    for eps in [0.0, 0.001, 0.01, 0.05, 0.125]:
        count = 100_000
        q = rng.random((count, 2, 2))
        q[:, 0, 0] *= eps
        q[:, 1, 1] = 1 - eps*rng.random(count)
        r = np.clip(q + rng.uniform(-eps, eps, q.shape), 0, 1)
        r[:, 0, 0] = np.minimum(r[:, 0, 0], eps)
        r[:, 1, 1] = np.maximum(r[:, 1, 1], 1-eps)
        assert np.max(np.abs(q-r)) <= eps+1e-14
        value = two_population_information(q, r)
        bound = (1-3*eps)**2/4
        slack = value-bound
        assert float(np.min(slack)) >= -1e-12
        # Check the two path inequalities used in the analytic proof.
        p1 = q[:, 0, 1]-q[:, 0, 0]+r[:, 1, 1]-r[:, 0, 1]
        p2 = r[:, 1, 0]-r[:, 0, 0]+q[:, 1, 1]-q[:, 1, 0]
        assert min(float(p1.min()), float(p2.min())) >= 1-3*eps-1e-12
        ans.append(dict(epsilon=eps, tested_channels=count,
                        theorem_bound_nats=bound,
                        minimum_observed_information_nats=float(value.min()),
                        minimum_observed_slack_nats=float(slack.min())))
    # Full grid in the exactly valid/agreed single-output case.
    a, b = np.meshgrid(np.linspace(0, 1, 201), np.linspace(0, 1, 201))
    q = np.zeros((a.size, 2, 2))
    q[:, 0, 1] = a.ravel(); q[:, 1, 0] = b.ravel(); q[:, 1, 1] = 1
    value = two_population_information(q, q)
    ans.append(dict(epsilon=0.0, exact_common_channel_grid_points=a.size,
                    grid_minimum_nats=float(value.min()),
                    value_at_symmetric_channel=2*float(entropy(.25))-math.log(2),
                    theorem_bound_nats=.25))
    return ans


def make_code(bits: int, rng: np.random.Generator) -> tuple[np.ndarray, np.ndarray]:
    if bits < 1 or bits > 12:
        raise ValueError('Verification code construction is limited to 1..12 bits')
    length = 8*bits
    messages = ((np.arange(1 << bits)[:, None] >> np.arange(bits)) & 1).astype(np.uint8)
    for _ in range(1000):
        generator = rng.integers(0, 2, size=(bits, length), dtype=np.uint8)
        words = (messages @ generator) & 1
        if int(words[1:].sum(axis=1).min()) >= length//4:
            return generator, (2*words.astype(float)-1)
    raise RuntimeError('Failed to construct a verified code within the test budget')


def code_checks(rng: np.random.Generator) -> tuple[list[dict], np.ndarray]:
    ans = []
    code3 = None
    for bits in range(1, 9):
        g, c = make_code(bits, rng)
        dmin = int(((c[1:] != c[0])).sum(axis=1).min())
        assert dmin >= c.shape[1]/4
        ans.append(dict(message_bits=bits, blocklength=c.shape[1],
                        codewords=c.shape[0], minimum_hamming_distance=dmin,
                        generator_matrix=g.tolist()))
        if bits == 3:
            code3 = c
    assert code3 is not None
    return ans, code3


def dissemination_checks(code: np.ndarray, rng: np.random.Generator) -> list[dict]:
    """Check a single receiver with a known-good committee, not full BA."""
    ans = []
    qwords, length = code.shape
    eta = .25
    n0 = 1.0
    trials = 4000
    for m in [4, 7, 10]:
        t = (m-1)//3
        h = m-t
        threshold = (m+1)//2-t
        for row_energy in [0.5, 2., 8., 32.]:
            amp = math.sqrt(row_energy/length)
            u = rng.integers(0, qwords, size=trials)
            rival = (u+1+rng.integers(0, qwords-1, size=trials)) % qwords
            truth = amp*code[u]
            y = truth[:, None, :] + rng.normal(0, math.sqrt(n0/2), (trials, m, length))
            # Arbitrarily large, coherent malicious waveforms targeting a wrong word.
            y[:, h:, :] = 9*amp*code[rival, None, :]
            step = 2*eta*amp
            radius = amp+5*math.sqrt(n0)
            j = math.ceil(radius/step)
            quant = step*np.clip(np.rint(y/step), -j, j)
            scores = quant @ code.T
            raw_scores = y[:, :h, :] @ code.T
            winners = np.full(trials, -1, dtype=int)
            winner_count = np.zeros(trials, dtype=int)
            for a in range(qwords):
                wins_all = np.ones(trials, dtype=bool)
                for b in range(qwords):
                    if a == b:
                        continue
                    favors_a = (scores[:, :, a] > scores[:, :, b]) | (
                        (scores[:, :, a] == scores[:, :, b]) & (a < b))
                    wins_all &= favors_a.sum(axis=1) > m/2
                winners[wins_all] = a
                winner_count += wins_all
            assert winner_count.max() <= 1
            decoded = np.where(winners < 0, 0, winners)
            failure = decoded != u
            honest_bad_event = np.zeros(trials, dtype=bool)
            gaussian_bad_event = np.zeros(trials, dtype=bool)
            saturated = (np.abs(y[:, :h, :]) > radius).any(axis=(1, 2))
            quant_implication_violations = 0
            ids = np.arange(trials)
            for b in range(qwords):
                active = b != u
                true_score = scores[ids, :, u]
                other_score = scores[:, :, b]
                honest_true_vote = (true_score[:, :h] > other_score[:, :h]) | (
                    (true_score[:, :h] == other_score[:, :h]) & (u[:, None] < b))
                honest_bad_event |= active & ((~honest_true_vote).sum(axis=1) >= threshold)
                raw_diff = raw_scores[ids, :, u] - raw_scores[:, :, b]
                distance = (code[u] != code[b]).sum(axis=1)
                # Scores use +/-1 words, so raw margin is 2*amp*Hamming distance.
                allowed = eta*2*amp*distance[:, None]
                raw_bad = raw_diff <= allowed+1e-12
                gaussian_bad_event |= active & (raw_bad.sum(axis=1) >= threshold)
                violation = active[:, None] & (~saturated[:, None]) & (~honest_true_vote) & (~raw_bad)
                quant_implication_violations += int(violation.sum())
            assert not np.any(failure & ~honest_bad_event)
            assert not np.any(failure & ~saturated & ~gaussian_bad_event)
            assert quant_implication_violations == 0
            bound = min(1.0, (qwords-1)*math.comb(m, threshold)*math.exp(
                -threshold*(1-eta)**2*.25*row_energy/n0))
            ans.append(dict(committee=m, byzantine=t, honest=h,
                            necessary_honest_pairwise_errors=threshold,
                            row_energy_over_N0=row_energy, trials=trials,
                            observed_errors=int(failure.sum()),
                            no_condorcet_winner=int((winners < 0).sum()),
                            saturated_honest_trial_count=int(saturated.sum()),
                            deterministic_implication_violations=0,
                            quantization_implication_violations=0,
                            union_bound_excluding_saturation=bound))
    return ans


def scalar_swap_check(rng: np.random.Generator) -> dict:
    count = 100_000
    a, b, y, z = rng.normal(size=(4, count))
    direct = (z-a)**2+(y-b)**2-(y-a)**2-(z-b)**2
    factor = 2*(y-z)*(a-b)
    err = float(np.max(np.abs(direct-factor)))
    assert err < 1e-12
    return dict(samples=count, max_absolute_algebra_error=err)


def lower_bound_examples() -> list[dict]:
    ans = []
    for n in [4, 32, 1024, 1_000_000]:
        for eps in [.1, .01, 1/n**2]:
            s = n//2; q = (1-3*eps)/2
            rel = s/16*ber_kl(q, eps/s)
            for block in [1, math.ceil(math.log(n/eps)), 1000]:
                inf = s*block*(1-3*eps)**2/4
                ans.append(dict(n=n, epsilon=eps, block_bits=block,
                                reliability_lower_over_N0=rel,
                                information_lower_over_N0=inf,
                                combined_lower_over_N0=max(rel, inf)))
    return ans


def main() -> None:
    rng = np.random.default_rng(SEED)
    codes, c3 = code_checks(rng)
    result = dict(seed=SEED,
                  interpretation='Finite algebraic/combinatorial checks; not proofs or a physical realization.',
                  information=information_checks(rng),
                  codes=codes,
                  receiver=dissemination_checks(c3, rng),
                  gaussian_swap=scalar_swap_check(rng),
                  lower_bound_examples=lower_bound_examples())
    out = ROOT/'results'/'communication.json'
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding='utf-8')
    print(json.dumps(dict(output=str(out),
                          information_channels=500_000,
                          exact_grid_points=40_401,
                          verified_codes=len(codes),
                          receiver_trials=sum(x['trials'] for x in result['receiver']),
                          violations=sum(x['deterministic_implication_violations'] for x in result['receiver']),
                          scalar_algebra_max_error=result['gaussian_swap']['max_absolute_algebra_error']), indent=2))

if __name__ == '__main__':
    main()