#!/usr/bin/env python3
"""Inference with the trained BaseNet -> Brainmu generation-expert LoRA."""
from __future__ import annotations
import argparse, json, math, os, random, sys
from pathlib import Path
import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from accelerate import init_empty_weights, load_checkpoint_and_dispatch
from PIL import Image
from safetensors import safe_open
from project_config import load_config

class LoRALinear(nn.Module):
    def __init__(self, base: nn.Linear, rank: int, alpha: float, dropout: float):
        super().__init__(); self.base=base; self.base.requires_grad_(False); self.rank=rank; self.alpha=alpha; self.scaling=alpha/rank
        self.dropout=nn.Dropout(dropout) if dropout>0 else nn.Identity()
        self.lora_A=nn.Parameter(torch.empty(rank,base.in_features,device=base.weight.device,dtype=torch.float32),requires_grad=False)
        self.lora_B=nn.Parameter(torch.empty(base.out_features,rank,device=base.weight.device,dtype=torch.float32),requires_grad=False)
    def forward(self,x):
        return self.base(x)+F.linear(F.linear(self.dropout(x),self.lora_A),self.lora_B)*self.scaling

def resolve_parent(root,dotted):
    parts=dotted.split('.'); parent=root
    for part in parts[:-1]: parent=parent[int(part)] if part.isdigit() else getattr(parent,part)
    return parent,parts[-1]

def inject_lora(model,spec,adapter_path):
    modules={}
    for name in spec['target_modules']:
        parent,leaf=resolve_parent(model,name); base=getattr(parent,leaf)
        if not isinstance(base,nn.Linear): raise TypeError(f'{name} is {type(base)}, expected nn.Linear')
        wrapped=LoRALinear(base,int(spec['rank']),float(spec['alpha']),float(spec['dropout'])); setattr(parent,leaf,wrapped); modules[name]=wrapped
    with safe_open(adapter_path,framework='pt',device='cpu') as f:
        keys=set(f.keys()); expected={f'{n}.lora_A' for n in modules}|{f'{n}.lora_B' for n in modules}
        if keys!=expected: raise RuntimeError(f'adapter key mismatch missing={len(expected-keys)} extra={len(keys-expected)}')
        with torch.no_grad():
            for name,m in modules.items():
                m.lora_A.copy_(f.get_tensor(f'{name}.lora_A').to(m.lora_A.device)); m.lora_B.copy_(f.get_tensor(f'{name}.lora_B').to(m.lora_B.device))
    return modules

def set_seed(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)

def ssim_gray(a,b):
    c1,c2=.01**2,.03**2; ma=cv2.GaussianBlur(a,(11,11),1.5); mb=cv2.GaussianBlur(b,(11,11),1.5)
    sa=cv2.GaussianBlur(a*a,(11,11),1.5)-ma*ma; sb=cv2.GaussianBlur(b*b,(11,11),1.5)-mb*mb; sab=cv2.GaussianBlur(a*b,(11,11),1.5)-ma*mb
    return float((((2*ma*mb+c1)*(2*sab+c2))/((ma*ma+mb*mb+c1)*(sa+sb+c2))).mean())

def metrics(pred,gt):
    p=np.asarray(pred.convert('L'),np.float32)/255.; t=np.asarray(gt.convert('L').resize(pred.size),np.float32)/255.; mse=float(np.mean((p-t)**2))
    return {'psnr_db':float(-10*np.log10(max(mse,1e-12))),'ssim':ssim_gray(p,t)}

def files(path):
    if path.is_file(): return [path]
    return sorted(p for p in path.iterdir() if p.suffix.lower() in {'.png','.jpg','.jpeg','.webp'})

def main():
    base=Path(__file__).resolve().parents[1]
    settings=load_config()['generation']['inference']
    ap=argparse.ArgumentParser()
    ap.add_argument('--input',type=Path,required=True,help='BaseNet PNG or directory of BaseNet PNGs')
    ap.add_argument('--output-dir',type=Path,required=True)
    ap.add_argument('--gt',type=Path,help='optional GT PNG or directory with matching filenames')
    ap.add_argument('--brainmu-repo',type=Path,default=base/'vendor/Brainmu')
    ap.add_argument('--model-path',type=Path,default=base/'checkpoints/Brainmu')
    ap.add_argument('--adapter',type=Path,default=base/'checkpoints/lora/adapter.safetensors')
    ap.add_argument('--adapter-config',type=Path,default=base/'checkpoints/lora/adapter_config.json')
    ap.add_argument('--prompt',default=settings['prompt'])
    ap.add_argument('--steps',type=int,default=settings['steps']); ap.add_argument('--seed',type=int,default=settings['seed']); ap.add_argument('--seed-offset',type=int,default=0); ap.add_argument('--seed-mode',choices=['sequential','base-prefix'],default='sequential',help='base-prefix uses the leading numeric clip id so all temporal spans of one clip share a seed')
    a=ap.parse_args()
    a.input=a.input.resolve(); a.output_dir=a.output_dir.resolve(); a.brainmu_repo=a.brainmu_repo.resolve(); a.model_path=a.model_path.resolve(); a.adapter=a.adapter.resolve(); a.adapter_config=a.adapter_config.resolve()
    if a.gt: a.gt=a.gt.resolve()
    a.output_dir.mkdir(parents=True,exist_ok=True)
    inputs=files(a.input)
    if not inputs: raise RuntimeError(f'no images found: {a.input}')
    sys.path.insert(0,str(a.brainmu_repo.resolve())); os.chdir(a.brainmu_repo)
    from data.data_utils import add_special_tokens
    from data.transforms import ImageTransform
    from inferencer import InterleaveInferencer
    from modeling.autoencoder import load_ae
    from modeling.brainmu import Brainmu,BrainmuConfig,Qwen2Config,Qwen2ForCausalLM,SiglipVisionConfig,SiglipVisionModel
    from modeling.qwen2 import Qwen2Tokenizer
    model_path=a.model_path.resolve(); device=torch.device('cuda:0'); torch.cuda.set_device(device); torch.backends.cuda.matmul.allow_tf32=True
    llm=Qwen2Config.from_json_file(str(model_path/'llm_config.json')); llm.qk_norm=True; llm.tie_word_embeddings=False; llm.layer_module='Qwen2MoTDecoderLayer'
    vit=SiglipVisionConfig.from_json_file(str(model_path/'vit_config.json')); vit.rope=False; vit.num_hidden_layers-=1
    vae,vae_cfg=load_ae(local_path=str(model_path/'ae.safetensors'))
    cfg=BrainmuConfig(visual_gen=True,visual_und=True,llm_config=llm,vit_config=vit,vae_config=vae_cfg,vit_max_num_patch_per_side=70,connector_act='gelu_pytorch_tanh',latent_patch_size=2,max_latent_size=64,timestep_shift=1.0)
    with init_empty_weights():
        lm=Qwen2ForCausalLM(llm); vm=SiglipVisionModel(vit); model=Brainmu(lm,vm,cfg); model.vit_model.vision_model.embeddings.convert_conv2d_to_linear(vit,meta=True)
    print('LOAD_BASE_BEGIN',flush=True)
    model=load_checkpoint_and_dispatch(model,checkpoint=str(model_path/'ema.safetensors'),device_map={'':0},dtype=torch.bfloat16,force_hooks=True)
    model.requires_grad_(False).eval(); vae=vae.to(device=device,dtype=torch.float32).eval().requires_grad_(False)
    oe,od=vae.encode,vae.decode; vae.encode=lambda x:oe(x.to(device=device,dtype=torch.float32)); vae.decode=lambda z:od(z.to(device=device,dtype=torch.float32))
    spec=json.load(open(a.adapter_config)); adapters=inject_lora(model,spec,a.adapter); model.eval(); print(f'LORA_LOADED modules={len(adapters)} adapter={a.adapter}',flush=True)
    tok=Qwen2Tokenizer.from_pretrained(str(model_path)); tok,new_ids,_=add_special_tokens(tok)
    infer=InterleaveInferencer(model,vae,tok,ImageTransform(400,256,16),ImageTransform(392,252,14),new_ids)
    rows=[]
    for i,p in enumerate(inputs):
        seed_index = i if a.seed_mode == 'sequential' else int(p.stem.split('_',1)[0])
        seed=a.seed+a.seed_offset+seed_index; set_seed(seed); inp=Image.open(p).convert('RGB')
        with torch.no_grad():
            pred=infer.interleave_inference([inp,a.prompt],think=False,understanding_output=False,cfg_text_scale=settings['cfg_text_scale'],cfg_img_scale=settings['cfg_img_scale'],cfg_interval=settings['cfg_interval'],timestep_shift=settings['timestep_shift'],num_timesteps=a.steps,cfg_renorm_min=settings['cfg_renorm_min'],cfg_renorm_type=settings['cfg_renorm_type'])[-1]
        out=a.output_dir/(p.stem+'.png'); pred.save(out); row={'input':str(p),'output':str(out),'seed':seed,'prompt':a.prompt}
        if a.gt:
            gp=a.gt if a.gt.is_file() else a.gt/(p.stem+'.png'); row.update(metrics(pred,Image.open(gp).convert('RGB'))); row['gt']=str(gp)
        rows.append(row); print(json.dumps(row),flush=True)
    result={'count':len(rows),'items':rows}
    if a.gt:
        result['mean_psnr_db']=float(np.mean([x['psnr_db'] for x in rows])); result['mean_ssim']=float(np.mean([x['ssim'] for x in rows]))
    (a.output_dir/'metrics.json').write_text(json.dumps(result,indent=2)+'\n'); print('INFERENCE_COMPLETE '+json.dumps({k:v for k,v in result.items() if k!='items'}),flush=True)
if __name__=='__main__': main()