"""LongCat-Next dNaViT/RVQ codec and base pixel decoder. Full AR backbone excluded."""
from pathlib import Path
import json,torch
from torch import nn
from accelerate import init_empty_weights
from safetensors.torch import load_file
from native.configuration_longcat_next import LongcatNextConfig
from native.visual import LongcatNextVisualTokenizer,VisionTransformerDecoder

def build(config,decoder=False,decoder_state=None):
    cfg=LongcatNextConfig(**config);m=nn.Module();m.model=nn.Module();m.model.visual_tokenizer=LongcatNextVisualTokenizer(cfg)
    if decoder:
        from diffusers import AutoencoderKL
        m.image_decoder=VisionTransformerDecoder(cfg.visual_config.visual_decoder_config.image_decoder_config)
        m.image_refiner=nn.Module();m.image_refiner.vae=AutoencoderKL.from_config(config["visual_config"]["visual_decoder_config"]["vae_config"])
        w=decoder_state["image_refiner.cond_proj.weight"];m.image_refiner.cond_proj=nn.Linear(w.shape[1],w.shape[0],bias=True)
    return m

def load_model(path=None,device="cpu",decoder=True):
    root=Path(path or Path(__file__).parent);c=json.loads((root/"source_config.json").read_text());s=load_file(root/"model.safetensors")
    aliases=json.loads((root/"aliases.json").read_text())
    for alias,canonical in aliases.items():s[alias]=s[canonical]
    ds=load_file(root/"decoder/model.safetensors") if decoder else None
    if ds:s.update(ds)
    with init_empty_weights():m=build(c,decoder,ds)
    m.load_state_dict(s,strict=True,assign=True);m=m.to(device=device).eval()
    # Preserve F32 codebooks and BF16 input weights. Base decoder and optional VAE
    # use F32 on CPU so quantized F32 features do not hit BF16 linear weights.
    if decoder:m.image_decoder.float();m.image_refiner.float()
    return m

def patchify_pixels(pixels):
    """NCHW RGB in [0,1]; no resize. Spatial dimensions must be divisible by 28."""
    if pixels.ndim!=4 or pixels.shape[1]!=3 or any(v%28 for v in pixels.shape[-2:]):raise ValueError("Expected NCHW RGB with dimensions divisible by 28")
    if pixels.min()<0 or pixels.max()>1:raise ValueError("Expected pixels in [0,1]")
    mean=pixels.new_tensor([.48145466,.4578275,.40821073])[None,:,None,None];std=pixels.new_tensor([.26862954,.26130258,.27577711])[None,:,None,None]
    x=(pixels-mean)/std;patches=[];grids=[]
    for image in x:
        _,h,w=image.shape;gh,gw=h//14,w//14
        p=image.unsqueeze(0).expand(2,-1,-1,-1).reshape(1,2,3,gh//2,2,14,gw//2,2,14).permute(0,3,6,4,7,2,1,5,8).reshape(gh*gw,-1)
        patches.append(p);grids.append([1,gh,gw])
    return torch.cat(patches),torch.tensor(grids,device=pixels.device,dtype=torch.long)

def encode_pixels(model,pixels):
    patches,grid=patchify_pixels(pixels);ids=model.model.visual_tokenizer.encode(patches,grid)
    return ids,grid

def code_features(model,ids):
    if ids.shape[-1]!=8 or ids.min()<0 or ids.max()>=16384:raise ValueError("Expected 8 residual-stage IDs in [0,16383]")
    books=model.model.visual_tokenizer.visual_bridge_model.quantizer.quantize.codebooks
    return sum(books[i].embed[ids[...,i].long()] for i in range(8))

def decode_codes(model,ids,grid_thw):
    features=code_features(model,ids).reshape(-1,3584).to(next(model.image_decoder.parameters()))
    decoded=model.image_decoder(features,grid_thw)["images"]
    return decoded # native CLIP-normalized RGB; first frame of each temporal patch


def load_refiner(path=None,device="cpu"):
    """Load the complete native flow Transformer, refiner VAE and conditioning projection."""
    from native.image_refiner import ImageRefinerContainer
    root=Path(path or Path(__file__).parent);cfg=LongcatNextConfig(**json.loads((root/"source_config.json").read_text())).visual_config.visual_decoder_config
    with init_empty_weights():m=ImageRefinerContainer(cfg)
    state={k.removeprefix("image_refiner."):v for k,v in load_file(root/"decoder/model.safetensors").items() if k.startswith("image_refiner.")}
    state.update({k.removeprefix("image_refiner."):v for k,v in load_file(root/"refiner/model.safetensors").items()})
    m.load_state_dict(state,strict=True,assign=True)
    return m.to(device=device).eval()

def refine_codes(codec,refiner,ids,grid_thw,seed=1019,num_inference_steps=28,output_type="pt"):
    """Native flow refinement of the base reconstruction. Requires both bundles loaded."""
    from native.image_refiner import RefinerPipeline,RefinerImageProcessor,de_transform
    from native.refiner_modules import FlowMatchEulerDiscreteScheduler
    device=next(refiner.parameters()).device;cfg=json.loads((Path(__file__).parent/"source_config.json").read_text())["visual_config"]["visual_decoder_config"]
    scheduler=FlowMatchEulerDiscreteScheduler(**cfg["scheduler_config"])
    pipeline=RefinerPipeline(vae=refiner.vae,transformer=refiner.base_transformer,cond_proj=refiner.cond_proj,scheduler=scheduler)
    raw=decode_codes(codec,ids,grid_thw);images=[]
    for image in raw:
        x=de_transform(image,mean=[.48145466,.4578275,.40821073],std=[.26862954,.26130258,.27577711],rescale_factor=1/255)
        images.append(RefinerImageProcessor.normalize(x).squeeze(0).to(device))
    features=code_features(codec,ids)
    if features.ndim==2:features=features.unsqueeze(0)
    result=pipeline(encoder_hidden_states=features.to(device=device,dtype=next(refiner.cond_proj.parameters()).dtype),grid_thw_list=[tuple(map(int,g)) for g in grid_thw],image=images,generator=torch.Generator(device=device).manual_seed(seed),num_inference_steps=num_inference_steps,output_type=output_type,return_dict=True,enable_processor_bar=False)
    return result.images