"""Locator variance experiment on KREF Q1 10-Q pages 23 and 33. Scratch only; no production code touched.
Variants: V0 current prompt; V1 adds a side-by-side rule; V2 = V1 plus a running-header exclusion.
Then V3: the existing vision count probe over the region V0 produced, to see whether it detects two tables."""
import asyncio, json, sys, tempfile, time
from pathlib import Path
from quber.agents import grid_locator as gl
from quber.core.extractors.camelot.acquire import render_pages
S=Path('/tmp/claude-1000/-home-mande-repo-quber/bc2f817a-5c6c-4515-9621-2302b82f7f68/scratchpad')
N=int(sys.argv[1]) if len(sys.argv)>1 else 4
BASE=gl.build_prompt(gl.DEFAULT_ROWS, gl.DEFAULT_COLS)
A="- Two tables stacked vertically are separate ONLY when each has its own\n  title/caption and its own column-header block. Report them separately."
B="Exclude page furniture (running header timestamp, page number, document URL)"
assert A in BASE and B in BASE, "prompt anchors moved"
SIDE=A+"""
- Two tables printed SIDE BY SIDE are separate tables: each has its own
  column-header block and its own row-label column, and a vertical band of
  empty grid columns separates them. Report each with its own column span;
  never report one region spanning both."""
HDR=("Exclude page furniture (running header timestamp, page number, document URL,\n"
     "and the running header block repeated at the top of every page: the company\n"
     "name, the statement title and a page-level units line). A region never\n"
     "extends up into that header block, even when it looks like a title")
VARIANTS={'V0':BASE, 'V1':BASE.replace(A,SIDE), 'V2':BASE.replace(A,SIDE).replace(B,HDR)}
results=[]
async def main():
    for page in (23,33):
        pdf=S/f'kref-p{page}.pdf'
        with tempfile.TemporaryDirectory() as tmp:
            img=render_pages(pdf, 200, Path(tmp))[0]
            for name,prompt in VARIANTS.items():
                gl.build_prompt=lambda r,c,_p=prompt: _p
                loc=gl.PydanticAIGridLocator()
                assert loc.system_prompt==prompt
                for i in range(N):
                    t0=time.time()
                    try:
                        located=await loc.locate(img, pdf, 1)
                        rec={'page':page,'variant':name,'run':i+1,'n':len(located),
                             'tables':[{'title':t.title,'rows':t.grid_rows,'cols':t.grid_cols,'region':[round(x,3) for x in t.region]} for t in located],'secs':round(time.time()-t0,1)}
                    except Exception as e:
                        rec={'page':page,'variant':name,'run':i+1,'error':str(e)[:200]}
                    results.append(rec); print(json.dumps(rec), flush=True)
            # V3: count probe over the widest V0 region
            from quber.core.extractors.set_of_mark.split import vision_count_probe
            from quber.agents.llm_client import PydanticAIClient
            llm=PydanticAIClient()
            v0=[r for r in results if r['page']==page and r['variant']=='V0' and 'tables' in r]
            merged=[t for r in v0 for t in r['tables'] if t['region'][2]-t['region'][0]>0.7]
            if merged:
                region=tuple(merged[0]['region'])
                for i in range(N):
                    c=await vision_count_probe(llm, img, region, 612.0, 792.0, 200)
                    rec={'page':page,'variant':'V3-probe','run':i+1,'region':list(region),'count':c}
                    results.append(rec); print(json.dumps(rec), flush=True)
    json.dump(results, open(S/'som_experiment.json','w'), indent=1)
asyncio.run(main())