from pathlib import Path
import json,time,torch
from component import load_model
root=Path(__file__).parent;torch.set_num_threads(4);checks=[];t=time.time()
def check(name,ok,**details):
    assert ok,(name,details)
    checks.append(dict(name=name,passed=True,**details));print(name,details,flush=True)
# Identical learned dimensions/weights; reduced spatial geometry for a bounded CPU test.
m=load_model(root,smoke_geometry=(120,120));r=m.model.modality_router;e=r.encoders["weather"];d=r.decoders["weather"]
check("strict_weather_components",True,tensors=len(m.state_dict()))
check("decoder_encoder_reference",d._encoder_ref[0] is e)
check("decoder_cold_cache",d.predict_from_hidden(torch.zeros(1,400,4096)) is None)
e.inject_meteorological_context(channels=["synthetic_"+str(i) for i in range(70)],indices={"accumid":[]},coords={},buffers={"const":torch.zeros(1,6,120,120)})
values=torch.linspace(-.1,.1,70*120*120).reshape(1,1,70,120,120);ids=torch.zeros(1,400,dtype=torch.long);mask=torch.ones_like(ids)
with torch.inference_mode():
    latent,latent_mask=e(ids,mask,meteo_values=values,times=["2026-10-19T00:00:00"],lead_hours=torch.tensor([6.]));check("weather_encode_native",latent.shape==(1,400,4096) and latent_mask is None and torch.isfinite(latent).all().item(),shape=list(latent.shape))
    check("weather_identity_projector",torch.equal(r.projectors["weather"](latent),latent))
    check("encoder_decoder_cache",set(["condition_embed","patch_embed_post_swin","meteo_values","input_size"]).issubset(e._step_cache))
    # The actual MKB language backbone is deliberately excluded. This tests the
    # native decoder interface using explicit synthetic LLM hidden states.
    prediction=d.predict_from_hidden(torch.zeros_like(latent),input_ids=torch.full_like(ids,e.config.meteo_token_id),__weather_pad_id__=e.config.meteo_token_id);check("weather_decode_native",prediction.shape==(1,70,120,120) and torch.isfinite(prediction).all().item(),shape=list(prediction.shape),dtype=str(prediction.dtype))
    check("native_float32_residual",prediction.dtype==torch.float32)
check("weather_source_bytes_exact",json.loads((root/"extraction.json").read_text())["source_tensor_bytes_equal"])
report=dict(strict_load=True,source_tensor_bytes_equal=True,checks=checks,smoke=dict(passed=True,device="cpu",dtype="float32",geometry_override=[120,120],original_geometry=[721,1440],fields="synthetic standardized inputs; zero static constants",decoder_hidden_states="synthetic zeros, not MKB language outputs"),elapsed_seconds=time.time()-t,limitations=["No 721x1440 full-resolution forecast, ERA5 evaluation, physical-value de-normalization, or multi-day rollout.","Real forecasting additionally requires original ERA5 variable ordering, normalization/static-field data and the MKB language backbone.","No medical SAM3 or biological sequence components included."])
(root/"validation.json").write_text(json.dumps(report,indent=2)+"\n");print("MKB WEATHER VALIDATED",report["elapsed_seconds"],flush=True)