#!/usr/bin/env python3
# TreeXcode — reference decoder
# Copyright 2026 Aleksandr Sitar <support@treexcode.com>
# SPDX-License-Identifier: Apache-2.0
# Licensed under the Apache License, Version 2.0 (see the LICENSE file or
# http://www.apache.org/licenses/LICENSE-2.0). Distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND.
"""
TreeXcode reference decoder.

Usage:
    python treexcode_decode.py code.png

Input:  an image containing a TreeXcode (a clean render or a photo).
Output: the decoded content plus a breakdown -- version, type, encoding,
        error-correction level and whether errors were repaired.

Pipeline (see https://treexcode.com/math.html):
    colour-independent foreground -> connected components -> central node
    -> four diamond fiducials -> homography (perspective) OR centre-only
    -> lattice size N -> ray classifier -> Reed-Solomon + CRC32 -> unpack.

Requires: pillow, numpy.  Must sit next to treexcode_encode.py (shared layout code).
"""
import sys, math, itertools, zlib
import numpy as np
from PIL import Image
import treexcode_encode as tx

TYPE_NAMES=['text','url','treex-short','email','phone','sms','geo','wifi','vcard','event','payment','binary']
ENC_NAMES=['UTF-8','Numeric','Alphanumeric','Latin-1']
AN=tx.AN; ECC_NSYM=tx.ECC_NSYM; GREEN=tx.GREEN; TYPES4=tx.TYPES4

# ---------------------------------------------------------------- GF(256) + RS decode
EXP,LOG=tx.EXP,tx.LOG
def gmul(a,b): return 0 if a==0 or b==0 else EXP[LOG[a]+LOG[b]]
def gpow(a,p): return EXP[(LOG[a]*p)%255]
def ginv(a): return EXP[255-LOG[a]]
def gdiv(a,b): return 0 if a==0 else EXP[(LOG[a]+255-LOG[b])%255]
def pscale(p,s): return [gmul(c,s) for c in p]
def padd(p,q):
    r=[0]*max(len(p),len(q))
    for i in range(len(p)): r[i+len(r)-len(p)]=p[i]
    for i in range(len(q)): r[i+len(r)-len(q)]^=q[i]
    return r
def pmul(p,q):
    r=[0]*(len(p)+len(q)-1)
    for j in range(len(q)):
        for i in range(len(p)): r[i+j]^=gmul(p[i],q[j])
    return r
def peval(p,x):
    y=p[0]
    for i in range(1,len(p)): y=gmul(y,x)^p[i]
    return y
def rs_synd(msg,nsym): return [0]+[peval(msg,gpow(2,i)) for i in range(nsym)]
def rs_errloc(synd,nsym):
    err=[1]; old=[1]
    for i in range(nsym):
        delta=synd[i+1]
        for j in range(1,len(err)): delta^=gmul(err[-(j+1)],synd[i+1-j])
        old=old+[0]
        if delta:
            if len(old)>len(err):
                nl=pscale(old,delta); old=pscale(err,ginv(delta)); err=nl
            err=padd(err,pscale(old,delta))
    while err and err[0]==0: err.pop(0)
    return err
def rs_finderr(errloc,n):
    errs=len(errloc)-1; pos=[peval(errloc,gpow(2,i)) for i in range(n)]
    p=[n-1-i for i,v in enumerate(pos) if v==0]
    if len(p)!=errs: raise ValueError('too many errors')
    return p
def rs_correct(msg,synd,pos):
    coef=[len(msg)-1-p for p in pos]
    eloc=[1]
    for i in coef: eloc=pmul(eloc,padd([1],[gpow(2,i),0]))
    eev=pmul(list(reversed(synd)),eloc); eev=eev[len(eev)-len(eloc):]; eev=list(reversed(eev))
    X=[gpow(2,c) for c in coef]; E=[0]*len(msg)
    for i in range(len(X)):
        Xi=X[i]; Xinv=ginv(Xi); prime=1
        for j in range(len(X)):
            if j!=i: prime=gmul(prime,1^gmul(Xinv,X[j]))
        y=peval(list(reversed(eev)),Xinv); y=gmul(Xi,y)
        if prime==0: raise ValueError('div0')
        E[pos[i]]=gdiv(y,prime)
    return padd(msg,E)
def rs_decode(msg,nsym):
    synd=rs_synd(msg,nsym)
    if max(synd)==0: return msg[:-nsym]
    errloc=list(reversed(rs_errloc(synd,nsym))); pos=rs_finderr(errloc,len(msg))
    cor=rs_correct(msg,synd,pos)
    if max(rs_synd(cor,nsym))!=0: raise ValueError('correction failed')
    return cor[:-nsym]

# ---------------------------------------------------------------- unpack encodings
def bytes_to_bits(bs):
    out=[]
    for b in bs:
        for i in range(7,-1,-1): out.append((b>>i)&1)
    return out
def packed_byte_len(L,mode):
    if mode==1: return math.ceil((L//3*10+[0,4,7][L%3])/8)
    if mode==2: return math.ceil((L//2*11+(L%2)*6)/8)
    return L
def unpack_content(bs,L,mode):
    if mode==1:
        bits=bytes_to_bits(bs); p=[0]; out=''
        def rd(n):
            v=0
            for _ in range(n): v=(v<<1)|bits[p[0]]; p[0]+=1
            return v
        g=L//3
        for _ in range(g): out+=str(rd(10)).zfill(3)
        r=L-g*3
        if r==2: out+=str(rd(7)).zfill(2)
        elif r==1: out+=str(rd(4))
        return out
    if mode==2:
        bits=bytes_to_bits(bs); p=[0]; out=''
        def rd(n):
            v=0
            for _ in range(n): v=(v<<1)|bits[p[0]]; p[0]+=1
            return v
        pr=L//2
        for _ in range(pr): v=rd(11); out+=AN[v//45]+AN[v%45]
        if L%2: out+=AN[rd(6)]
        return out
    if mode==3: return ''.join(chr(b) for b in bs)
    return bytes(bs).decode('utf-8')

def unpack(syms):
    by=[]
    for i in range(0,len(syms)-3,4): by.append((syms[i]<<6)|(syms[i+1]<<4)|(syms[i+2]<<2)|syms[i+3])
    if len(by)<8: return None
    b0,b1=by[0],by[1]; ver=(b0>>6)&3; typ=(b0>>2)&15; mode=(b1>>5)&7; level=b1&15
    if level>=len(ECC_NSYM): return None
    nsym=ECC_NSYM[level]; L=(by[2]<<8)|by[3]; PB=packed_byte_len(L,mode)
    D=max(1,251-nsym); K=max(1,math.ceil(PB/D))          # число RS-блоков
    pos=4; content=[]; corrected=0
    for bi in range(K):
        clen=0 if PB==0 else min(D, PB-bi*D); blockLen=clen+4+nsym
        if pos+blockLen>len(by): return None
        block=by[pos:pos+blockLen]; pos+=blockLen
        if nsym>0:
            try:
                if max(rs_synd(block,nsym))!=0: corrected=1
                block=rs_decode(block,nsym)
            except Exception: return None
        chunk=block[:clen]; cb=block[clen:clen+4]
        crc=(cb[0]<<24)|(cb[1]<<16)|(cb[2]<<8)|cb[3]
        if (zlib.crc32(bytes(chunk))&0xffffffff)!=crc: return None
        content+=chunk
    try: text=unpack_content(content,L,mode)
    except Exception: return None
    if not text: return None    # пустой/тривиальный результат (напр. чистый лист) — не считаем кодом
    return dict(version=ver,type=typ,mode=mode,ecc=level,corrected=corrected,content=text)

# ---------------------------------------------------------------- image -> foreground
def foreground(a):
    H,W=a.shape[:2]; g=a.astype(np.float64).mean(axis=2)
    win=max(9,(H//12)|1); win=min(win,41)|1; r=win//2
    ii=np.zeros((H+1,W+1)); ii[1:,1:]=np.cumsum(np.cumsum(g,0),1)
    loc=np.empty((H,W)); X0=np.clip(np.arange(W)-r,0,W); X1=np.clip(np.arange(W)+r+1,0,W)
    for yy in range(H):
        y0=max(0,yy-r); y1=min(H,yy+r+1)
        loc[yy]=(ii[y1,X1]-ii[y0,X1]-ii[y1,X0]+ii[y0,X0])/((y1-y0)*(X1-X0))
    diff=g-loc; thr=max(12,1.2*diff.std())
    return np.abs(diff)>thr

def fill_holes(mask):
    # local-threshold turns filled shapes into rings; refill enclosed background so
    # the node, diamonds and leaves become solid blobs again (thin glyphs are unaffected).
    from collections import deque
    H,W=mask.shape; bg=np.zeros((H,W),bool); q=deque()
    for x in range(W):
        if not mask[0,x]: bg[0,x]=True; q.append((0,x))
        if not mask[H-1,x]: bg[H-1,x]=True; q.append((H-1,x))
    for y in range(H):
        if not mask[y,0]: bg[y,0]=True; q.append((y,0))
        if not mask[y,W-1]: bg[y,W-1]=True; q.append((y,W-1))
    while q:
        y,x=q.popleft()
        for dy,dx in ((1,0),(-1,0),(0,1),(0,-1)):
            ny,nx=y+dy,x+dx
            if 0<=ny<H and 0<=nx<W and not mask[ny,nx] and not bg[ny,nx]: bg[ny,nx]=True; q.append((ny,nx))
    return ~bg

def components(mask):
    from collections import deque
    H,W=mask.shape; lab=-np.ones((H,W),int); comps=[]; cur=0
    for sy in range(H):
        for sx in range(W):
            if mask[sy,sx] and lab[sy,sx]<0:
                q=deque([(sy,sx)]); lab[sy,sx]=cur; sxx=syy=n=0; a0=b0=sx; a1=b1=sy
                while q:
                    y,x=q.popleft(); sxx+=x; syy+=y; n+=1
                    a0=min(a0,x); a1=max(a1,x); b0=min(b0,y); b1=max(b1,y)
                    for dy,dx in ((1,0),(-1,0),(0,1),(0,-1)):
                        ny,nx=y+dy,x+dx
                        if 0<=ny<H and 0<=nx<W and mask[ny,nx] and lab[ny,nx]<0:
                            lab[ny,nx]=cur; q.append((ny,nx))
                w=a1-a0+1; h=b1-b0+1
                comps.append(dict(area=n,cx=sxx/n,cy=syy/n,w=w,h=h,sol=n/(w*h))); cur+=1
    return comps

# ---------------------------------------------------------------- geometry
def gauss(A,b):
    n=len(b); M=[list(A[i])+[b[i]] for i in range(n)]
    for col in range(n):
        piv=max(range(col,n),key=lambda r:abs(M[r][col]))
        if abs(M[piv][col])<1e-12: return None
        M[col],M[piv]=M[piv],M[col]; d=M[col][col]
        for r in range(n):
            if r!=col and M[r][col]:
                f=M[r][col]/d
                for c in range(col,n+1): M[r][c]-=f*M[col][c]
    return [M[i][n]/M[i][i] for i in range(n)]
def homography(src,dst):
    A=[];b=[]
    for (x,y),(u,v) in zip(src,dst):
        A.append([x,y,1,0,0,0,-u*x,-u*y]); b.append(u)
        A.append([0,0,0,x,y,1,-v*x,-v*y]); b.append(v)
    h=gauss(A,b)
    return None if h is None else [[h[0],h[1],h[2]],[h[3],h[4],h[5]],[h[6],h[7],1.0]]
def mv(H,x,y):
    px=H[0][0]*x+H[0][1]*y+H[0][2]; py=H[1][0]*x+H[1][1]*y+H[1][2]; pw=H[2][0]*x+H[2][1]*y+H[2][2]
    return px/pw,py/pw
def inv3(H):
    a,b,c=H[0]; d,e,f=H[1]; g,h,i=H[2]
    A=e*i-f*h; B=-(d*i-f*g); C=d*h-e*g; det=a*A+b*B+c*C
    return [[A/det,-(b*i-c*h)/det,(b*f-c*e)/det],[B/det,(a*i-c*g)/det,-(a*f-c*d)/det],[C/det,-(a*h-b*g)/det,(a*e-b*d)/det]]
def infer_N(pts):
    if len(pts)<8: return None
    nn=[]
    for i in range(len(pts)):
        best=1e9
        for j in range(len(pts)):
            if i!=j:
                d=math.hypot(pts[i][0]-pts[j][0],pts[i][1]-pts[j][1])
                if 1e-4<d<best: best=d
        if best<0.15: nn.append(best)
    if len(nn)<8: return None
    lo,hi=min(nn),max(nn)
    if hi<=lo: return None
    cnt=[0]*80
    for d in nn: cnt[min(79,int((d-lo)/(hi-lo)*80))]+=1
    k=cnt.index(max(cnt)); pitch=lo+(k+0.5)*(hi-lo)/80
    return round(1/pitch) if pitch>0 else None
def refine_N(src,cc,Nest):
    best=Nest; berr=1e18
    for N in range(max(24,Nest-6),Nest+7):
        H=homography(src,[[0.5,2.0/N],[0.5,1-2.0/N],[2.0/N,0.5],[1-2.0/N,0.5]])
        if H is None: continue
        e=0.0
        for cx,cy in cc:
            u,v=mv(H,cx,cy); gx=u*N-0.5; gy=v*N-0.5; e+=abs(gx-round(gx))+abs(gy-round(gy))
        e/=len(cc)
        if e<berr: berr=e; best=N
    return best
def pixel_pitch(pts):
    nn=[]
    for i in range(len(pts)):
        best=1e9
        for j in range(len(pts)):
            if i!=j:
                d=math.hypot(pts[i][0]-pts[j][0],pts[i][1]-pts[j][1])
                if 1<d<best: best=d
        if best<1e8: nn.append(best)
    if len(nn)<8: return None
    lo=min(nn); hi=max(nn)*0.6+min(nn)*0.4
    if hi<=lo: hi=lo+1
    cnt=[0]*60
    for d in nn:
        if d<=hi: cnt[min(59,int((d-lo)/(hi-lo)*60))]+=1
    k=cnt.index(max(cnt)); return lo+(k+0.5)*(hi-lo)/60
def refine_CSpx(thin,O,CS0):
    def snap(v): f=v-math.floor(v); return min(f,abs(f-0.5),1-f)
    best=CS0; berr=1e18
    for k in range(61):
        s=CS0*0.85+CS0*0.30*k/60; e=sum(snap((cx-O[0])/s)+snap((cy-O[1])/s) for cx,cy in thin)
        if e<berr: berr=e; best=s
    return best
def find_repers(solid,node):
    O=(node['cx'],node['cy']); na=node['area']; nr=math.sqrt(na)
    cand=[c for c in solid if c is not node and 0.30<c['sol']<0.80 and 0.04*na<c['area']<1.4*na
          and math.hypot(c['cx']-O[0],c['cy']-O[1])>1.6*nr]
    if len(cand)<4: return None
    cand.sort(key=lambda c:-math.hypot(c['cx']-O[0],c['cy']-O[1])); cand=cand[:16]
    best=None; bsc=1e18
    for quad in itertools.combinations(range(len(cand)),4):
        P=[cand[i] for i in quad]; cs=[(p['cx'],p['cy']) for p in P]
        Cx=sum(p[0] for p in cs)/4; Cy=sum(p[1] for p in cs)/4
        ordr=sorted(range(4),key=lambda k:math.atan2(cs[k][1]-Cy,cs[k][0]-Cx)); Po=[cs[k] for k in ordr]
        d1=math.hypot((Po[0][0]+Po[2][0])/2-Cx,(Po[0][1]+Po[2][1])/2-Cy)+math.hypot((Po[1][0]+Po[3][0])/2-Cx,(Po[1][1]+Po[3][1])/2-Cy)
        v1=(Po[2][0]-Po[0][0],Po[2][1]-Po[0][1]); v2=(Po[3][0]-Po[1][0],Po[3][1]-Po[1][1])
        perp=abs(v1[0]*v2[0]+v1[1]*v2[1])/(math.hypot(*v1)*math.hypot(*v2)+1e-6)
        radii=[math.hypot(p[0]-Cx,p[1]-Cy) for p in cs]; rmean=sum(radii)/4+1e-6
        rvar=sum(abs(r-rmean) for r in radii)/rmean
        areas=[p['area'] for p in P]; amean=sum(areas)/4+1e-6; avar=sum(abs(x-amean) for x in areas)/amean
        centO=math.hypot(Cx-O[0],Cy-O[1])/rmean
        score=d1/rmean+2.0*perp+1.2*rvar+0.8*avar+1.5*centO
        if score<bsc: bsc=score; best=P
    if best is None or bsc>1.3: return None
    P=best; T=min(P,key=lambda c:c['cy']); Bt=max(P,key=lambda c:c['cy']); L=min(P,key=lambda c:c['cx']); R=max(P,key=lambda c:c['cx'])
    return [[T['cx'],T['cy']],[Bt['cx'],Bt['cy']],[L['cx'],L['cy']],[R['cx'],R['cy']]]

def sample_classify(fg,Hi,N,x,y,CSpx):
    H,W=fg.shape; n=max(8,CSpx); half=(n-1)/2.0; hr=n/2.0
    m=np.zeros((n,n),bool)
    for jj in range(n):
        for ii in range(n):
            u=(x+(ii+0.5)/n)/N; v=(y+(jj+0.5)/n)/N
            px,py=mv(Hi,u,v); ix=int(round(px)); iy=int(round(py))
            if 0<=ix<W and 0<=iy<H: m[jj,ii]=fg[iy,ix]
    def spoke(ux,uy):
        c=0
        for k in range(6):
            s=0.30+(0.82-0.30)*k/5; ix=int(round(half+ux*s*hr)); iy=int(round(half+uy*s*hr))
            if 0<=ix<n and 0<=iy<n and m[iy,ix]: c+=1
        return c/6>0.5
    r2=1/math.sqrt(2)
    E=spoke(1,0);Ws=spoke(-1,0);NE=spoke(r2,-r2);NW=spoke(-r2,-r2);SE=spoke(r2,r2);SW=spoke(-r2,r2)
    horiz=E or Ws; diag=NE or NW or SE or SW
    if not horiz and not diag: return 'empty'
    if horiz and not diag: return 'full_h' if(E and Ws) else 'half_h'
    if diag and not horiz: return 'full_d' if((NE and SW) or(NW and SE)) else 'half_d'
    return 'full_d'

def _read(fg,Hi,N):
    _,greens,_=tx.gen_tree(N); greens.sort(key=lambda p:(p[1],p[0]))
    a0,b0=mv(Hi,0.5,0.5); a1,b1=mv(Hi,0.5+1.0/N,0.5); CSpx=max(6,int(round(math.hypot(a1-a0,b1-b0))))
    syms=[]
    for x,y in greens:
        t=sample_classify(fg,Hi,N,x,y,CSpx); syms.append(TYPES4.index(t) if t in TYPES4 else 0)
    return syms

def decode(a):
    raw=foreground(a); H,W=raw.shape
    # fast path: чистый рендер — код это центрированный квадрат min(W,H) (учитывает поля caption сверху/снизу)
    if min(W,H) >= 0.8*max(W,H):
        S=min(W,H); ox=(W-S)/2.0; oy=(H-S)/2.0
        for N in range(48,121):
            r=unpack(_read(raw,[[S,0,ox],[0,S,oy],[0,0,1]],N))
            if r is not None: r['N']=N; r['mode_geom']='frame'; return r
    comps=[c for c in components(raw) if c['area']>=8]
    round_c=[c for c in comps if 0.45<c['w']/max(1,c['h'])<2.2]
    if not round_c: return None
    node=max(round_c,key=lambda c:c['area'])          # central node = largest roundish structure
    O=(node['cx'],node['cy']); nr=math.sqrt(node['area']); fg=raw
    solid=[c for c in comps if c['sol']>0.35 and 0.4<c['w']/max(1,c['h'])<2.6]
    thin=[(c['cx'],c['cy']) for c in comps if c['sol']<0.85 and 3<c['area']<0.30*node['area']]
    src=find_repers(solid,node)
    if src is not None:
        dst=[[0.5,0],[0.5,1],[0,0.5],[1,0.5]]; N=None
        for _ in range(4):
            Hc=homography(src,dst)
            if Hc is None: break
            Nn=infer_N([mv(Hc,cx,cy) for cx,cy in thin])
            if not Nn or Nn<24: break
            N=Nn; dst=[[0.5,2.0/N],[0.5,1-2.0/N],[2.0/N,0.5],[1-2.0/N,0.5]]
        if N and N>=24:
            N=refine_N(src,thin,N)
            Hi=inv3(homography(src,[[0.5,2.0/N],[0.5,1-2.0/N],[2.0/N,0.5],[1-2.0/N,0.5]]))
            syms=_read(fg,Hi,N); r=unpack(syms)
            if r is not None: r['N']=N; r['mode_geom']='homography'; return r
    if len(thin)>=8:
        CSpx=pixel_pitch(thin)
        if CSpx and CSpx>3:
            CSpx=refine_CSpx(thin,O,CSpx)
            for N in range(24,122):
                nodeYc=round(N*0.6); s=N*CSpx
                Hi=[[s,0,O[0]-N/2*CSpx],[0,s,O[1]-(nodeYc+0.5)*CSpx],[0,0,1]]
                r=unpack(_read(fg,Hi,N))
                if r is not None: r['N']=N; r['mode_geom']='centre'; return r
    return None

def main():
    if len(sys.argv)<2:
        print('usage: python treexcode_decode.py <image>'); sys.exit(1)
    img=Image.open(sys.argv[1]).convert('RGB'); a=np.array(img)
    r=decode(a)
    if r is None:
        print('No TreeXcode could be decoded from this image.'); sys.exit(2)
    print('Content : %s' % r['content'])
    print('Type    : %s' % TYPE_NAMES[r['type']] if r['type']<len(TYPE_NAMES) else r['type'])
    print('Version : %d' % r['version'])
    print('Encoding: %s' % ENC_NAMES[r['mode']])
    print('ECC     : %s%s' % (['none','low','medium','high'][r['ecc']], '  (errors repaired)' if r['corrected'] else ''))
    print('Size    : %dx%d  (%s)' % (r['N'], r['N'], r['mode_geom']))

if __name__=='__main__':
    main()