"""Complete trained non-language Harmon components with native computations."""
import json
from pathlib import Path
import torch
from torch import nn
from accelerate import init_empty_weights
from safetensors.torch import load_file
from einops import rearrange
from native.vae import AutoencoderKL
from native.mar import mar_base

class HarmonVisualComponents(nn.Module):
    vae_scale=0.2325
    def __init__(self,cfg):
        super().__init__()
        self.vae=AutoencoderKL(embed_dim=16,ch_mult=(1,1,2,2,4))
        mar=dict(cfg['mar']);assert mar.pop('type')=='mar_base'
        self.mar=mar_base(**mar)
        d=cfg['llm']['hidden_size'];s=self.mar.encoder_embed_dim
        self.proj_in=nn.Sequential(nn.Linear(s,d),nn.SiLU(),nn.Linear(d,d))
        self.proj_out=nn.Sequential(nn.Linear(d,d),nn.SiLU(),nn.Linear(d,s))

    def encode(self,x):
        posterior=self.vae.encode(x)
        z=posterior.mode().mul_(self.vae_scale)
        return rearrange(z,'b c (m p) (n q) -> b m n (c p q)',p=self.mar.patch_size,q=self.mar.patch_size)

    def decode(self,z):
        # The original method modifies z in place; preserve caller input here.
        z=z/self.vae_scale
        z=rearrange(z,'b m n (c p q) -> b c (m p) (n q)',p=self.mar.patch_size,q=self.mar.patch_size)
        return self.vae.decode(z)

def load_model(path=None,dtype=torch.float32,device='cpu'):
    root=Path(path or Path(__file__).parent)
    with init_empty_weights():model=HarmonVisualComponents(json.loads((root/'source_config.json').read_text()))
    model.load_state_dict(load_file(root/'model.safetensors'),strict=True,assign=True)
    return model.to(device=device,dtype=dtype).eval()