from pathlib import Path
import json,torch
from torch import nn
from transformers.configuration_utils import PretrainedConfig
from accelerate import init_empty_weights
from safetensors.torch import load_file
from native.weather.encoder import WeatherEncoder
from native.weather.decoder import WeatherDecoder
from native.weather.projector import IdentityWeatherProjector

def build(config,smoke_geometry=None):
    config=dict(config)
    if smoke_geometry is not None:config["image_size"]=list(smoke_geometry)
    cfg=PretrainedConfig(**config);m=nn.Module();m.model=nn.Module();m.model.modality_router=nn.Module();r=m.model.modality_router
    r.encoders=nn.ModuleDict({"weather":WeatherEncoder(cfg)});r.projectors=nn.ModuleDict({"weather":IdentityWeatherProjector(cfg.qwenvl_dim)});r.decoders=nn.ModuleDict({"weather":WeatherDecoder(cfg,r.encoders["weather"])})
    return m

def load_model(path=None,device="cpu",dtype=torch.float32,smoke_geometry=None):
    root=Path(path or Path(__file__).parent);cfg=json.loads((root/"weather_config.json").read_text())
    with init_empty_weights():m=build(cfg,smoke_geometry)
    m.load_state_dict(load_file(root/"model.safetensors"),strict=True,assign=True)
    return m.to(device=device,dtype=dtype).eval()