"""Opt-in experimental icon labels; labels are evidence, never click commands."""
import io,json
from pathlib import Path
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from safetensors.torch import load_file
from transformers.models.idefics3.configuration_idefics3 import Idefics3VisionConfig
from transformers.models.idefics3.modeling_idefics3 import Idefics3VisionTransformer
from miniswardbower.browser.dom_distiller import DOMDistiller

class VisualIconDistiller(DOMDistiller):
    def __init__(self, directory):
        super().__init__(draw_overlays=False)
        directory=Path(directory)
        self.labels=json.loads((directory/'report.json').read_text())['labels']
        config=Idefics3VisionConfig.from_json_file(directory/'vision_config.json');config._attn_implementation='sdpa'
        self.vision=Idefics3VisionTransformer(config)
        state=load_file(str(directory/'icon_specialist.safetensors'))
        self.vision.load_state_dict({k[7:]:v for k,v in state.items() if k.startswith('vision.')},strict=True)
        self.head=torch.nn.Linear(config.hidden_size,len(self.labels))
        self.head.load_state_dict({k[5:]:v for k,v in state.items() if k.startswith('head.')},strict=True)
        self.vision=self.vision.cuda().bfloat16().eval();self.head=self.head.cuda().eval()
        self.observations=[]

    @torch.inference_mode()
    def classify(self, image):
        a=torch.from_numpy(np.asarray(image.convert('RGB').resize((512,512))).copy()).permute(2,0,1).float()/255
        features=self.vision(pixel_values=((a-.5)/.5)[None].cuda().bfloat16()).last_hidden_state.mean(1).float()
        probs=self.head(F.normalize(features,dim=-1)).softmax(-1)[0]
        return self.labels[int(probs.argmax())],float(probs.max())

    async def extract_pruned_tree(self,page):
        tree=await super().extract_pruned_tree(page)
        # A locator screenshot scrolls offscreen elements into view. Perception must not act.
        viewport=Image.open(io.BytesIO(await page.screenshot(full_page=False,scale='css')))
        count=0
        for elem in tree.elements:
            x,y,w,h=elem.bbox
            if elem.aria_label or elem.tag not in ('button','a') or not (0<w<=160 and 0<h<=100):continue
            if elem.text and any(c.isalpha() for c in elem.text):continue
            if count>=20:break
            if x<0 or y<0 or x+w>viewport.width or y+h>viewport.height:continue
            name,confidence=self.classify(viewport.crop((x,y,x+w,y+h)));count+=1
            self.observations.append(dict(mark=elem.id,label=name,confidence=confidence))
            if name!='other':
                elem.text=(elem.text or '')+f' [visual guess: {name}; confidence {confidence:.2f}]'
        return tree