#!/usr/bin/env python3
"""Finite causal-path checks for the reception-history interchange identity.

The test machines are finite-state interactive networks, NOT consensus
protocols. Algebra, local output interchange, and the involution are checked
pathwise. Monte Carlo means are diagnostic, not proofs of the identity.
"""
from __future__ import annotations
import json
import platform
from pathlib import Path
import numpy as np
import scipy

ROOT = Path(__file__).resolve().parents[1]
SEED = 20261028 + 91
N0 = 1.7
SAMPLES = 4000


def schedule(n: int) -> list[tuple[int, int]]:
    return [(s, r) for _ in range(2) for s in range(n)
            for r in range(n) if r != s]


def machine(inputs: np.ndarray, private: np.ndarray,
            slots: list[tuple[int, int]], *,
            rng: np.random.Generator | None = None,
            observations: np.ndarray | None = None
            ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Generate or replay a deterministic finite-state causal machine."""
    batches, n = private.shape
    state = (private + 3 * inputs[None, :]) % 19
    amplitudes = np.empty((batches, len(slots)), dtype=float)
    if observations is None:
        if rng is None:
            raise ValueError('Generation requires an RNG')
        values = np.empty_like(amplitudes)
    else:
        if observations.shape != amplitudes.shape:
            raise ValueError('Observation shape does not match schedule')
        values = observations.copy()
    for t, (s, r) in enumerate(slots):
        amplitudes[:, t] = (0.45 * (2 * inputs[s] - 1)
                            + 0.13 * ((state[:, s] % 7) - 3)
                            + 0.09 * np.sin(private[:, s] + t))
        if observations is None:
            values[:, t] = (amplitudes[:, t]
                            + rng.normal(0, np.sqrt(N0 / 2), batches))
        quantized = np.clip(np.rint(2 * values[:, t]), -9, 9).astype(int)
        state[:, r] = (3 * state[:, r] + quantized + t + inputs[r]) % 19
    outputs = (state + inputs[None, :]) % 2
    return amplitudes, values, outputs


def run() -> dict:
    rng = np.random.default_rng(SEED)
    reports = []
    for n in (4, 7, 10):
        slots = schedule(n)
        x0 = np.arange(n) % 2
        x1 = x0.copy()
        x1[n // 2:] = 1 - x1[n // 2:]
        private = rng.integers(0, 23, (SAMPLES, n))
        a0, y0, o0 = machine(x0, private, slots, rng=rng)
        a1, y1, o1 = machine(x1, private, slots, rng=rng)
        for i in range(min(3, n // 2)):
            incoming = np.array([r == i for s, r in slots])
            incident = np.array([s == i or r == i for s, r in slots])
            z0, z1 = y0.copy(), y1.copy()
            z0[:, incoming] = y1[:, incoming]
            z1[:, incoming] = y0[:, incoming]
            b0, _, p0 = machine(x0, private, slots, observations=z0)
            b1, _, p1 = machine(x1, private, slots, observations=z1)
            out0, out1 = o0.copy(), o1.copy()
            out0[:, i], out1[:, i] = o1[:, i], o0[:, i]
            assert np.array_equal(p0, out0) and np.array_equal(p1, out1)
            # Log(p(y)/p(T_i y)): both densities replay their own causal encoders.
            logratio = np.sum(((z0-b0)**2 + (z1-b1)**2
                               - (y0-a0)**2 - (y1-a1)**2) / N0, axis=1)
            edge_identity = np.sum(2 * (y0[:, incident]-y1[:, incident])
                                    * (a0[:, incident]-a1[:, incident]) / N0,
                                    axis=1)
            residual = float(np.max(np.abs(logratio-edge_identity)))
            assert residual < 2e-10
            target = np.sum(2 * (a0[:, incident]-a1[:, incident])**2 / N0,
                            axis=1)
            fluctuation = logratio-target
            se = float(np.std(fluctuation, ddof=1)/np.sqrt(SAMPLES))
            # A Monte Carlo standardized difference is reported, not asserted.
            reports.append(dict(n=n, node=i, samples=SAMPLES, uses=len(slots),
                                max_pathwise_algebra_error=residual,
                                mean_log_likelihood_ratio=float(np.mean(logratio)),
                                mean_signal_difference_cost=float(np.mean(target)),
                                mean_difference_standard_error=se,
                                mean_difference_z=float(np.mean(fluctuation)/se)))
    return dict(seed=SEED, N0=N0,
                scope='Finite-state test machines; not consensus simulations or proofs.',
                paths=reports,
                environment=dict(python=platform.python_version(),
                                 numpy=np.__version__, scipy=scipy.__version__))


if __name__ == '__main__':
    result = run()
    out = ROOT / 'results' / 'interchange.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(tests=len(result['paths']),
                         pairs_per_test=SAMPLES,
                         max_pathwise_error=max(t['max_pathwise_algebra_error'] for t in result['paths']),
                         max_abs_mean_z=max(abs(t['mean_difference_z']) for t in result['paths'])),indent=2))