"""NumPy CCSDS 131.0-B-5 section 6 turbo codec, including custom 3200/6400.

Positive LLR means bit 0. The custom exact-half-rate profile omits the eight
termination symbols; standard-size terminated mode is provided separately.
See SPACE_TURBO_SPEC.md for precise bit ordering and deviations.
"""
from functools import lru_cache
import numpy as np

STANDARD_SIZES=(1784,3568,7136,8920)

def bits(value,length=None):
    x=np.asarray(value)
    if x.ndim!=1 or not np.isin(x,[0,1]).all() or (length is not None and len(x)!=length):
        raise ValueError(f'Expected {length or "a one-dimensional array of"} binary bits')
    return x.astype(np.uint8)

@lru_cache(maxsize=8)
def interleaver(k=3200):
    """Zero-based read addresses; encoder b receives u[pi]. Sec. 6.3g."""
    if not isinstance(k,(int,np.integer)) or isinstance(k,bool) or k<8 or k%8:
        raise ValueError('k must be a positive multiple of 8')
    k1,k2=8,k//8;s=np.arange(k,dtype=np.int64);m=s%2;i=s//(2*k2);j=s//2-i*k2
    t=(19*i+1)%(k1//2);q=t%8
    prime=np.asarray([31,37,43,47,53,59,61,67],dtype=np.int64)
    c=(prime[q]*j+21*m)%k2
    pi=2*(t+c*(k1//2)+1)-m-1
    if not np.array_equal(np.sort(pi),np.arange(k)):
        raise ValueError('Document interleaver extension is not a permutation for this k')
    pi.flags.writeable=False
    return pi

@lru_cache(maxsize=1)
def trellis():
    nxt=np.empty((16,2),dtype=np.int64);par=np.empty((16,2),dtype=np.uint8)
    incoming=np.empty((16,2),dtype=np.int64);ins=np.empty((16,2),dtype=np.int64);counts=np.zeros(16,dtype=int)
    for s in range(16):
        d1,d3,d4=s&1,(s>>2)&1,(s>>3)&1
        for u in range(2):
            w=u^d3^d4;dest=((s<<1)&15)|w
            nxt[s,u]=dest;par[s,u]=w^d1^d3^d4
            incoming[dest,counts[dest]]=s;ins[dest,counts[dest]]=u;counts[dest]+=1
    return nxt,par,incoming,ins

def component_encode(u,terminated=True):
    """Figure 6-2: w=u xor d3 xor d4, parity=w xor d1 xor d3 xor d4."""
    u=bits(u);n=len(u)+4*bool(terminated);out=np.empty((n,2),dtype=np.uint8)
    nxt,par,_,_=trellis();s=0
    for t in range(n):
        b=int(u[t]) if t<len(u) else ((s>>2)^(s>>3))&1
        out[t]=(b,par[s,b]);s=int(nxt[s,b])
    if terminated and s:raise RuntimeError('Failed trellis termination')
    return out

def bcjr(systematic,parity,apriori,terminated=False,algorithm='logmap'):
    """Exact log-MAP or max-log-MAP SISO for one 16-state constituent."""
    nxt,par,inc,inp=trellis();n=len(systematic)
    sign=np.asarray([1.,-1.]);psign=1-2*par.astype(np.float64)
    gamma=.5*((systematic+apriori)[:,None,None]*sign+parity[:,None,None]*psign)
    alpha=np.full((n+1,16),-np.inf);beta=np.full((n+1,16),-np.inf);alpha[0,0]=0
    if terminated:beta[n,0]=0
    else:beta[n]=0
    combine=np.logaddexp if algorithm=='logmap' else np.maximum
    for t in range(n):
        v=alpha[t,inc]+gamma[t,inc,inp];alpha[t+1]=combine(v[:,0],v[:,1]);alpha[t+1]-=np.max(alpha[t+1])
    for t in range(n-1,-1,-1):
        v=gamma[t]+beta[t+1,nxt];beta[t]=combine(v[:,0],v[:,1]);beta[t]-=np.max(beta[t])
    joint=alpha[:-1,:,None]+gamma+beta[1:,nxt]
    reduce=np.logaddexp.reduce if algorithm=='logmap' else np.max
    return reduce(joint[:,:,0],axis=1)-reduce(joint[:,:,1],axis=1)

@lru_cache(maxsize=4)
def randomizer(n=6400):
    """Section 10 degree-17 PN, seed 11000111000111000 (left to right)."""
    if n<1:raise ValueError('n must be positive')
    # Rightmost register is the output; x17+x14+1 feedback from rightmost and stage 14.
    reg=[int(x) for x in '11000111000111000'];seq=np.empty(n,dtype=np.uint8)
    for i in range(n):
        seq[i]=reg[-1]
        feedback=reg[-1]^reg[2]
        reg=[feedback]+reg[:-1]
    seq.flags.writeable=False
    return seq

@lru_cache(maxsize=4)
def channel_permutation(n=6400):
    """Explicit custom physical dispersal, independent from the turbo interleaver."""
    # Affine bijection, portable across NumPy versions; gcd(173,6400)=1.
    if np.gcd(173,n)!=1:raise ValueError('Physical permutation requires gcd(173,n)=1')
    p=(173*np.arange(n,dtype=np.int64)+29)%n;p.flags.writeable=False;return p

class TurboCode:
    def __init__(self,info_bits=3200,terminated=False,transport=True):
        if not isinstance(info_bits,(int,np.integer)) or isinstance(info_bits,bool):raise ValueError('Invalid information length')
        self.info_bits=int(info_bits);self.terminated=bool(terminated);self.pi=interleaver(self.info_bits)
        self.code_bits=2*(self.info_bits+4*self.terminated);self.transport=bool(transport)
        self.physical=channel_permutation(self.code_bits) if transport else np.arange(self.code_bits)
        self.pn=randomizer(self.code_bits) if transport else np.zeros(self.code_bits,dtype=np.uint8)
    @property
    def rate(self):return self.info_bits/self.code_bits
    def encode(self,u):
        u=bits(u,self.info_bits);a=component_encode(u,self.terminated);b=component_encode(u[self.pi],self.terminated)
        out=np.empty((len(a),2),dtype=np.uint8);out[:,0]=a[:,0];out[::2,1]=a[::2,1];out[1::2,1]=b[1::2,1]
        return (out.reshape(-1)^self.pn)[self.physical]
    def encode_vectors(self,u):
        if self.code_bits%640:raise ValueError('Coded length must be divisible by 640 for 10-frame generator')
        return self.encode(u).reshape(-1,10,64)
    def decode_llr(self,llr,iterations=12,algorithm='logmap',return_details=False,early_stop=True):
        llr=np.asarray(llr,dtype=np.float64)
        if llr.size!=self.code_bits or not np.isfinite(llr).all():raise ValueError('Wrong LLR length or nonfinite LLR')
        if not isinstance(iterations,(int,np.integer)) or isinstance(iterations,bool) or iterations<1:raise ValueError('iterations must be positive')
        if algorithm not in ('logmap','maxlog'):raise ValueError('algorithm must be logmap or maxlog')
        ordered=np.empty(self.code_bits);ordered[self.physical]=llr.reshape(-1)
        ordered*=1-2*self.pn.astype(float);ordered=np.clip(ordered,-40,40).reshape(-1,2)
        k=self.info_bits;n=len(ordered);sa=ordered[:,0];sb=np.zeros(n);sb[:k]=sa[:k][self.pi]
        pa=np.zeros(n);pb=np.zeros(n);pa[::2]=ordered[::2,1];pb[1::2]=ordered[1::2,1]
        prior=np.zeros(n);previous=None;stable=0;changes=[]
        for iteration in range(1,iterations+1):
            post_a=bcjr(sa,pa,prior,self.terminated,algorithm)
            ext_a=np.clip(post_a[:k]-sa[:k]-prior[:k],-40,40)
            prior_b=np.zeros(n);prior_b[:k]=ext_a[self.pi]
            post_b=bcjr(sb,pb,prior_b,self.terminated,algorithm)
            ext_b=np.clip(post_b[:k]-sb[:k]-prior_b[:k],-40,40)
            prior[:k][self.pi]=ext_b
            post=np.empty(k);post[self.pi]=post_b[:k]
            hard=(post<0).astype(np.uint8)
            changed=k if previous is None else int(np.count_nonzero(hard!=previous));changes.append(changed)
            stable=stable+1 if changed==0 else 0;previous=hard.copy()
            if early_stop and iteration>=4 and stable>=2:break
        detail={'bits':hard,'posterior_llr':post,'iterations':iteration,'changed_bits_per_iteration':changes,
                'converged_stable':stable>=2,'note':'Decision stability is not a CRC or proof of successful recovery.'}
        return detail if return_details else hard
    def decode_probabilities(self,p,**kwargs):
        p=np.asarray(p,dtype=float)
        if not np.isfinite(p).all() or np.any((p<0)|(p>1)):raise ValueError('Invalid probability')
        p=np.clip(p,1e-7,1-1e-7)
        return self.decode_llr(np.log1p(-p)-np.log(p),**kwargs)