"""Load the complete InternVL-U image-generation decoder with native weights."""
import json
import sys
from pathlib import Path

import torch
from safetensors.torch import load_file


class ImageDiT:
    def __init__(self, native, root):
        self.native = native
        self.root = root

    @torch.inference_mode()
    def predict(self, latents, timestep, condition, *, condition_mask=None,
                condition_image_mask=None, condition_grid=None, reference_latents=None):
        if latents.ndim != 4 or latents.shape[1] != 16 or any(x % 2 for x in latents.shape[-2:]):
            raise ValueError('Expected even spatial dimensions and [B,16,H,W] native VAE latents')
        if condition.ndim != 3 or condition.shape[0] != latents.shape[0] or condition.shape[2] != 4096:
            raise ValueError('Expected original InternVL-U [B,L,4096] conditioning features')
        if condition_mask is None:
            condition_mask = torch.ones(condition.shape[:2], dtype=torch.bool, device=condition.device)
        if condition_image_mask is None:
            condition_image_mask = torch.zeros_like(condition_mask, dtype=torch.bool)
        if condition_grid is None:
            if condition_image_mask.any():
                raise ValueError('Visual conditioning tokens require their native image grid')
            condition_grid = [torch.empty((0, 3), dtype=torch.long, device=condition.device) for _ in range(len(latents))]
        timestep = torch.as_tensor(timestep, device=latents.device, dtype=latents.dtype)
        if timestep.ndim == 0:
            timestep = timestep.expand(len(latents))
        return self.native.decoder(
            hidden_states=latents, encoder_hidden_states=condition, timestep=timestep,
            encoder_attention_mask=condition_mask, encoder_image_token_mask=condition_image_mask,
            image_fhw_cond=condition_grid, conditional_input=reference_latents, return_dict=False,
        )[0]

    def make_scheduler(self):
        from diffusers import DPMSolverMultistepScheduler
        return DPMSolverMultistepScheduler.from_config(json.loads((self.root / 'scheduler_config.json').read_text()))


def load_component(root=None, *, dtype=torch.float32, device='cpu'):
    root = Path(root or Path(__file__).resolve().parent)
    sys.path.insert(0, str(root))
    from intern_native.configuration_internvlu_generation_decoder import InternVLUGenerationDecoderConfig
    from intern_native.modeling_internvlu_generation_decoder import InternVLUGenerationDecoder
    config = InternVLUGenerationDecoderConfig(**json.loads((root / 'config.json').read_text()))
    with torch.device('meta'):
        native = InternVLUGenerationDecoder(config)
    state = load_file(root / 'model.safetensors')
    native.load_state_dict(state, strict=True, assign=True)
    del state
    # This native module stores deterministic RoPE frequencies as a plain
    # tensor attribute. They are absent from state_dict and remain on meta
    # after assign=True, so reconstruct them with the upstream constructor.
    rope = native.decoder.pos_embed
    native.decoder.pos_embed = type(rope)(theta=rope.theta, axes_dim=rope.axes_dim, scale_rope=rope.scale_rope)
    native.to(device=device, dtype=dtype).eval()
    return ImageDiT(native, root)