"""Session 20. Source-pixel counts with explicitly manual plaintext alignment.
Run beside lp_core.py and the source images. Requires NumPy, SciPy, Pillow, matplotlib.
No downloads, image restoration, OCR or cipher-key search.
"""
from pathlib import Path
import json,hashlib
from collections import Counter
import numpy as np
from PIL import Image
from scipy import ndimage
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from lp_core import R,LET,pages,vals,SEQS
ROOT=Path(__file__).resolve().parent
MAP={1:0,5:3,6:4,7:5,8:6,9:7,10:8,11:9,12:10,13:11,16:14,73:71,74:72}
# Each entry: page-local rune offset, expected visual dot count, contextual role.
# Prefix marks can share a rune offset with the preceding sentence's end mark.
ALIGN={1:[(6,11,'after warning heading'),(184,11,'end of warning')],
5:[(10,11,'after wisdom heading'),(81,11,'after KNOW THIS, before square')],
6:[(5,11,'after koan heading')],7:[],8:[],
9:[(103,11,'end of koan'),(115,11,'after instruction heading')],
10:[(16,11,'after divinity heading'),(78,10,'before consumption heading'),(88,10,'after consumption heading')],
11:[(97,10,'end of consumption subsection'),(97,10,'before preservation heading'),(108,10,'after preservation heading')],
12:[(25,10,'end of preservation subsection'),(25,10,'before adherence heading'),(34,10,'after adherence heading'),(145,10,'end of adherence subsection')],
13:[(44,11,'end of divinity discussion'),(54,11,'after wisdom heading'),(129,11,'end of wisdom block'),(141,11,'after instruction heading'),(169,11,'end of instruction block')],
16:[(12,11,'after instruction heading'),(82,11,'end of instruction block'),(89,11,'after KNOW THIS, before square')],
73:[(5,13,'after AN END heading'),(45,10,'before hexadecimal digest'),(85,13,'end of page text')],
74:[(7,13,'after PARABLE heading'),(95,23,'end of parable / final page')]}

def detect(n,dom):
 a=np.asarray(Image.open(ROOT/f'{n:02}.jpg').convert('RGB')).astype(int)
 mask=(a[:,:,0]>100)&(a[:,:,0]-a[:,:,1]>dom)&(a[:,:,0]-a[:,:,2]>dom)
 lab,_=ndimage.label(mask);dots=[]
 for k,bb in enumerate(ndimage.find_objects(lab),1):
  if bb is None:continue
  sy,sx=bb;area=int(np.count_nonzero(lab[bb]==k))
  if 15<=area<=200 and sx.stop-sx.start<=18 and sy.stop-sy.start<=18:dots.append([sx.start,sy.start,sx.stop,sy.stop,area])
 joined=[]
 for d in sorted(dots,key=lambda d:(d[1],d[0])):
  near=[g for g in joined if any(abs(d[0]-e[0])<28 and abs(d[1]-e[1])<28 for e in g)]
  if not near:joined.append([d])
  else:
   g=near[0];g.append(d)
   for other in near[1:]:g.extend(other);joined.remove(other)
 out=[]
 for g in joined:
  if len(g)>=3:out.append(dict(dots=len(g),bbox=[min(d[0] for d in g),min(d[1] for d in g),max(d[2] for d in g),max(d[3] for d in g)],dot_boxes=g))
 return sorted(out,key=lambda g:(g['bbox'][1],g['bbox'][0]))

def plaintext(n):
 s=pages[MAP[n]];v=vals(s)
 if n==1:v=(28-v)%29
 if 6<=n<=9:v=(2-v)%29
 if n==73:
  ix=np.arange(len(v));ix[57:]-=1;d=(v-SEQS['prime_minus_one'][ix])%29;d[56]=v[56];v=d
 it=iter(v)
 return ''.join(LET[int(next(it))] if c in R else c for c in s)

def main():
 records=[]
 for n in MAP:
  p=ROOT/f'{n:02}.jpg';gs=detect(n,65);checks={str(d):[g['dots'] for g in detect(n,d)] for d in [40,65,90]}
  expected=[a[1] for a in ALIGN[n]]
  assert all(v==expected for v in checks.values()),(n,checks,expected)
  for g,(off,count,role) in zip(gs,ALIGN[n]):g.update(page_rune_offset=off,role=role,alignment='manual visual comparison with reproduced plaintext')
  records.append(dict(file=p.name,sha256=hashlib.sha256(p.read_bytes()).hexdigest(),size=Image.open(p).size,plaintext_with_source_delimiters=plaintext(n),groups=gs,threshold_checks=checks))
 result=dict(scope='13 selected solved pages, not the complete solved corpus',counts=dict(Counter(g['dots'] for p in records for g in p['groups'])),records=records)
 (ROOT/'solved_punctuation_results.json').write_text(json.dumps(result,ensure_ascii=False,indent=2)+'\n')
 fig,axes=plt.subplots(5,6,figsize=(15,12),layout='constrained')
 for ax,(p,g) in zip(axes.flat,[(p,g) for p in records for g in p['groups']]):
  x,y,xx,yy=g['bbox'];im=Image.open(ROOT/p['file']);ax.imshow(np.asarray(im)[y-6:yy+6,x-6:xx+6],interpolation='nearest');ax.axis('off');ax.set_title(f"{p['file']} / rune {g['page_rune_offset']}\n{g['dots']} dots",fontsize=10)
 fig.suptitle('Solved-page punctuation: native source-pixel samples, 30 marks',fontsize=15)
 fig.savefig(ROOT/'solved_punctuation_samples.png',dpi=140);plt.close(fig)
 print(json.dumps(dict(pages=len(records),marks=sum(result['counts'].values()),counts=result['counts'],threshold_checks=39)))
if __name__=='__main__':main()