from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from zipfile import ZipFile
import torch
from safetensors.torch import load_file
from transformers import AutoConfig,AutoProcessor
from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5Model,Qwen3_5VisionModel

from inference import Engine,Stream


class VisualModel(torch.nn.Module):
    get_rope_index=Qwen3_5Model.get_rope_index
    get_vision_position_ids=Qwen3_5Model.get_vision_position_ids
    get_placeholder_mask=Qwen3_5Model.get_placeholder_mask

    def __init__(self,text,config,visual):
        super().__init__()
        self.text,self.config,self.visual=text,config,visual

    def get_input_embeddings(self):
        return self.text.get_input_embeddings()

    def forward(self,input_ids,attention_mask,pixel_values,image_grid_thw):
        embeds=self.get_input_embeddings()(input_ids)
        features=self.visual(pixel_values.to(dtype=embeds.dtype),grid_thw=image_grid_thw,return_dict=True).pooler_output
        masks=self.get_placeholder_mask(input_ids,embeds,image_features=features)
        embeds=embeds.masked_scatter(masks[0],features.to(embeds))
        rope,deltas=self.get_rope_index(input_ids,mm_token_type_ids=(input_ids==self.config.image_token_id).long(),
            image_grid_thw=image_grid_thw,attention_mask=attention_mask)
        physical=(attention_mask.long().cumsum(-1)-1).clamp_min(0)
        positions=torch.cat((physical.unsqueeze(0),rope),dim=0)
        output=self.text(inputs_embeds=embeds,attention_mask=attention_mask,position_ids=positions,use_cache=True,logits_to_keep=1)
        lengths=attention_mask.sum(-1)
        spatial=lengths+deltas.flatten()
        output['vision_decode_positions']=torch.stack((lengths,spatial,spatial,spatial))[:,:,None]
        return output


class VisualContext:
    def __init__(self,model,stream,batch):
        self.model,self.stream=model,stream
        self.positions=torch.empty((4,2*batch+1,1),dtype=torch.long,device='cuda')

    def __call__(self,ids,mask,indices):
        selected=[self.stream.rows[i]['vision_inputs'] for i in indices.tolist()]
        output=self.model(input_ids=ids,attention_mask=mask,
            pixel_values=torch.cat([v['pixel_values'] for v in selected]).to(ids.device),
            image_grid_thw=torch.cat([v['image_grid_thw'] for v in selected]).to(ids.device))
        self.positions.index_copy_(1,indices,output['vision_decode_positions'])
        return output

    def decode_positions(self,indices):
        return self.positions.index_select(1,indices)


def generate_visual(args,claim,write,report):
    engine=Engine(args.model,args.batch,args.profile)
    config=AutoConfig.from_pretrained(args.model/'vision',local_files_only=True)
    config.vision_config._attn_implementation='flash_attention_2'
    visual=Qwen3_5VisionModel(config.vision_config)
    visual.load_state_dict(load_file(str(args.model/'vision/model.safetensors')))
    model=VisualModel(engine.model,config,visual).cuda().to(dtype=torch.bfloat16).eval().requires_grad_(False)
    processor=AutoProcessor.from_pretrained(args.model/'vision',local_files_only=True)
    processor.tokenizer.padding_side='left'
    processor.image_processor.size={'longest_edge':4194304,'shortest_edge':65536}
    media=ZipFile(args.model/'evaluation_data/vision/media.zip')
    pool=ThreadPoolExecutor(20)
    def image(row):
        from PIL import Image
        with media.open(row['image_path']) as source, Image.open(source) as image:
            return processor.image_processor(images=image.convert('RGB'),return_tensors='pt')
    def fetch(count):
        rows=claim(count)
        if not rows:return []
        images=list(pool.map(image,rows))
        texts=[]
        for row,pixels in zip(rows,images):
            question=((row.get('hint') or '')+'\n'+row['question']).strip()
            choices={k:v for k,v in (row.get('choices') or {}).items() if v is not None}
            row['choices']=choices
            if choices:
                question+='\n'+'\n'.join(f'{k}. {v}' for k,v in choices.items())
                question+='\nGive only the option letter as your final answer.'
            elif row['source'] in ('mme','pope'):
                question+='\nGive only Yes or No as your final answer.'
            else:question+='\nGive a short answer without explanation as your final answer.'
            text=processor.apply_chat_template([{'role':'user','content':[{'type':'image'},{'type':'text','text':question}]}],
                tokenize=False,add_generation_prompt=True,enable_thinking=True)
            count=int(pixels['image_grid_thw'].prod())//processor.image_processor.merge_size**2
            texts.append(text.replace(processor.image_token,processor.image_token*count,1))
        ids=processor.tokenizer(texts,add_special_tokens=False,padding=False)['input_ids']
        return [dict(row,input_ids=token,vision_inputs=pixels,max_tokens=args.max_new_tokens)
                for row,token,pixels in zip(rows,ids,images)]
    first=fetch(args.batch)
    if not first:return {}
    stream=Stream(engine,fetch,write,report)
    stream.prefill_chunk_size=lambda width:4
    context=VisualContext(model,stream,args.batch)
    try:
        return engine.run_stream(stream,first,cap=args.max_new_tokens,temperature=1,top_k=20,top_p=.95,
            presence_penalty=1.5,seed=args.seed,mechanical=False,prefill_context=context,
            eos_ids=(processor.tokenizer.eos_token_id,))
    finally:
        pool.shutdown(wait=True)
        media.close()