#!/usr/bin/env python3
"""Small complete coded-consensus executions under a specified static attack.

The tests implement every Phase King round, finite quantization, and the
external pairwise-majority decoder. They condition on a good four-member
committee; they are not an experiment establishing the sampling theorem.
Sub-bound energy scalings are stress tests, not predicted safe protocols.
N0=1. Byzantine internal values are injected at the decoded-message
interface, an overapproximation used by the deterministic agreement proof.
All honest transmissions use AWGN. Fixed seeds make the checks reproducible.
"""
from __future__ import annotations
import json
import math
from pathlib import Path
import numpy as np
from verify_communication import make_code

ROOT=Path(__file__).resolve().parents[1]
SEED=20261028+327
CSTAR=9/64
ETA=.25
N0=1.

def bits(value: int,L: int)->np.ndarray:
    return ((value >> np.arange(L))&1).astype(int)

def message(a: np.ndarray)->int:
    return int(np.dot(a.astype(int),1<<np.arange(len(a))))

def quantize(y: np.ndarray,amplitude: float,radius: float)->np.ndarray:
    step=2*ETA*amplitude
    J=math.ceil(radius/step)
    return np.clip(np.rint(y/step),-J,J).astype(np.int64)

class Experiment:
    def __init__(self,L: int,rng: np.random.Generator):
        self.L=L; self.rng=rng; self.n=7; self.m=4; self.t=1
        self.fault=3; self.honest=np.array([0,1,2]); self.eps=.125
        self.P=(self.t+1)*(self.m-1)*(2*self.m+1)
        self.generators={}
        self.generators['internal'],self.code=make_code(2*L,rng)
        self.generators['external'],self.outer=make_code(L,rng)
        self.candidates=np.array([bits(i,L) for i in range(1<<L)])
        self.packet_energy=(2*L*math.log(2)+math.log(6*self.P/self.eps))/CSTAR
        self.r=(self.m+1)//2-self.t
        self.outer_energy=(L*math.log(2)+math.log(6*(self.n-self.m)/self.eps)
                           +math.log(math.comb(self.m,self.r)))/(CSTAR*self.r)
        self.internal_bad=False

    def channel(self,msg:int,scale:float)->int:
        energy=self.packet_energy*scale; amp=math.sqrt(energy/(16*self.L))
        radius=amp+math.sqrt(math.log(4*(16*self.L)/(self.eps/(3*self.P))))
        y=amp*self.code[msg]+self.rng.normal(0,math.sqrt(.5),self.code.shape[1])
        q=quantize(y,amp,radius)
        return int(np.argmax(self.code@q))

    def transmit(self,kind:int,s:int,r:int,state:np.ndarray,proposals:np.ndarray,
                 candidates:np.ndarray,phase:int,scale:float)->np.ndarray:
        if s==self.fault:
            # Equivocation: receiver-dependent valid-looking packets. It is only
            # one concrete attack family, not an adversarial maximization.
            if kind==1:
                out=np.array([(r+phase+j)%3 for j in range(self.L)])
            else:out=np.array([(r+phase+j+kind)%2 for j in range(self.L)])
            return out
        raw=proposals[s] if kind==1 else (state[s] if kind==0 else candidates[s])
        if s==r:return raw.copy()
        packet=np.zeros(2*self.L,dtype=int)
        if kind==1:
            packet[0::2]=raw&1;packet[1::2]=raw>>1
        else:packet[:self.L]=raw
        mi=message(packet); received=self.channel(mi,scale)
        if received!=mi:self.internal_bad=True
        decoded=bits(received,2*self.L)
        if kind==1:
            vals=decoded[0::2]+2*decoded[1::2]
            return np.minimum(vals,2)
        return decoded[:self.L]

    def run(self,scale:float,trial:int)->dict:
        L=self.L
        self.fault=trial % self.m
        self.honest=np.array([i for i in range(self.m) if i!=self.fault])
        X=self.rng.integers(0,2,size=(self.n,L))
        # Include all-unanimous and coordinatewise-unanimous input cases.
        if trial%5==0:X[:]=0
        elif trial%5==1:X[:]=1
        elif L>1 and trial%5==2:X[:,0]=trial%2
        state=X[:self.m].copy();self.internal_bad=False
        for phase in range(self.t+1):
            proposed=np.full((self.m,L),2,dtype=int)
            cand=np.zeros((self.m,L),dtype=int)
            for r in self.honest:
                seen=np.array([self.transmit(0,s,int(r),state,proposed,cand,phase,scale)
                               for s in range(self.m)])
                for value in [0,1]:proposed[r,(seen==value).sum(axis=0)>=self.m-self.t]=value
            grade=np.zeros((self.m,L),dtype=int)
            for r in self.honest:
                seen=np.array([self.transmit(1,s,int(r),state,proposed,cand,phase,scale)
                               for s in range(self.m)])
                # Conflicting grade-1 candidates can occur on a channel-error
                # execution; fix a deterministic smallest-value tie break.
                for ell in range(L):
                    counts=np.array([(seen[:,ell]==v).sum() for v in [0,1]])
                    if counts.max()>=self.m-self.t:
                        grade[r,ell]=2;cand[r,ell]=int(np.argmax(counts))
                    elif counts.max()>=self.t+1:
                        grade[r,ell]=1;cand[r,ell]=int(np.argmax(counts))
            king=phase
            updated=state.copy()
            for r in self.honest:
                kvalue=self.transmit(2,king,int(r),state,proposed,cand,phase,scale)
                updated[r]=np.where(grade[r]==2,cand[r],kvalue)
            state=updated
        committee_ok=bool(np.all(state[self.honest]==state[self.honest[0]]))
        if not self.internal_bad:assert committee_ok
        outputs=np.zeros((self.n,L),dtype=int);outputs[:self.m]=state
        outer_energy=scale*self.outer_energy; amp=math.sqrt(outer_energy/(8*L))
        radius=amp+math.sqrt(math.log(12*(8*L)*self.m*(self.n-self.m)/self.eps))
        external_bad=False;certificate_bad=False
        D=message(state[self.honest[0]]);candidate_count=len(self.outer)
        for rec in range(self.m,self.n):
            y=np.zeros((self.m,8*L))
            for s in range(self.m):
                if s==self.fault:
                    rival=(D+1+rec) % candidate_count
                    y[s]=9*amp*self.outer[rival]
                else:
                    y[s]=amp*self.outer[message(state[s])]+self.rng.normal(0,math.sqrt(.5),8*L)
            q=quantize(y,amp,radius)
            scores=q@self.outer.T
            winner=None
            for u in range(candidate_count):
                success=True
                for v in range(candidate_count):
                    if u==v:continue
                    votes=(scores[:,u]>scores[:,v])|((scores[:,u]==scores[:,v])&(u<v))
                    if int(votes.sum())<=self.m/2:success=False;break
                if success:
                    assert winner is None;winner=u
            output=0 if winner is None else winner
            outputs[rec]=self.candidates[output]
            if committee_ok and output!=D:
                external_bad=True; found=False
                for v in range(candidate_count):
                    if v==D:continue
                    favorable=(scores[self.honest,D]>scores[self.honest,v])|(
                        (scores[self.honest,D]==scores[self.honest,v])&(D<v))
                    if int((~favorable).sum())>=self.r:found=True
                assert found
                certificate_bad=True
        honest_all=[i for i in range(self.n) if i!=self.fault]
        vals=outputs[honest_all]
        failure=not bool(np.all(vals==vals[0]))
        for ell in range(L):
            if np.all(X[honest_all,ell]==X[honest_all[0],ell]):
                failure|=not bool(np.all(vals[:,ell]==X[honest_all[0],ell]))
        assert not failure or self.internal_bad or external_bad
        return dict(failed=int(failure),internal_error=int(self.internal_bad),
                    external_disagreement=int(external_bad),
                    honest_subset_certificate=int(certificate_bad))

def main()->None:
    rng=np.random.default_rng(SEED);records=[];generators={}
    for L in [1,2,3]:
        experiment=Experiment(L,rng)
        generators[str(L)]={key:arr.tolist() for key,arr in experiment.generators.items()}
        for scale in [.01,.03,.1,1.]:
            count=150;totals=dict(failed=0,internal_error=0,external_disagreement=0,
                                 honest_subset_certificate=0)
            for trial in range(count):
                result=experiment.run(scale,trial)
                for key,value in result.items():totals[key]+=value
            records.append(dict(n=7,committee_size=4,static_faults=1,L=L,
                                epsilon=.125,energy_scale=scale,trials=count,**totals))
    report=dict(seed=SEED,total_executions=sum(r['trials'] for r in records),
                meaning='Complete small digital protocol executions on a fixed good committee; '
                        'Byzantine decoded-message injection overapproximates Byzantine channels; '
                        'sub-bound scales are stress tests, not theoretical guarantees.',
                invariant_violations=0,records=records,generators=generators)
    (ROOT/'results').mkdir(exist_ok=True)
    (ROOT/'results'/'protocol.json').write_text(json.dumps(report,indent=2),encoding='utf-8')
    print(json.dumps({k:v for k,v in report.items() if k!='generators'},indent=2))

if __name__=='__main__':main()