"""Deterministic feature extraction and paired-resolution reconstruction."""
import argparse, importlib.util
from pathlib import Path
import numpy as np, torch, yaml
ROOT=Path(__file__).resolve().parents[1]
def load(cfg):
    s=importlib.util.spec_from_file_location("scalemae",ROOT/"model/scalemae.py"); m=importlib.util.module_from_spec(s); s.loader.exec_module(m); return m.ScaleMAE(**cfg["model"])
def run(model,a,device,ratio):
    with torch.no_grad(): return model(torch.from_numpy(a["images"]).to(device),torch.from_numpy(a["gsd"]).to(device),target=torch.from_numpy(a["targets"]).to(device),mask_ratio=ratio)
def main():
    p=argparse.ArgumentParser(); p.add_argument("--config",type=Path,default=ROOT/"conf/config.yaml"); p.add_argument("--data",type=Path); p.add_argument("--train-data",type=Path); p.add_argument("--checkpoint",type=Path); p.add_argument("--output-dir",type=Path); p.add_argument("--device",choices=("auto","cpu","cuda")); a=p.parse_args(); cfg=yaml.safe_load(a.config.read_text()); ck=a.checkpoint or ROOT/cfg["paths"]["checkpoint"]
    if not ck.exists(): raise FileNotFoundError("Run training before inference")
    requested=a.device or cfg["runtime"]["device"]
    if requested == "cuda" and not torch.cuda.is_available(): raise RuntimeError("CUDA requested but unavailable")
    device=torch.device("cuda" if torch.cuda.is_available() and requested!="cpu" else "cpu"); torch.manual_seed(cfg["seed"]); model=load(cfg); model.load_state_dict(torch.load(ck,map_location=device,weights_only=False)["model"]); model.to(device).eval(); train=np.load(a.train_data or ROOT/cfg["data"]["root"]/"train.npz"); test=np.load(a.data or ROOT/cfg["data"]["root"]/"test.npz"); tr=run(model,train,device,0.0); te=run(model,test,device,cfg["model"]["mask_ratio"])
    out=a.output_dir or ROOT/cfg["paths"]["inference_dir"]; out.mkdir(parents=True,exist_ok=True); np.savez_compressed(out/"reconstruction.npz",target=test["targets"],prediction=te["reconstruction"].cpu().numpy(),low_target=te["low_target"].cpu().numpy(),low_prediction=te["low_reconstruction"].cpu().numpy(),high_target=te["high_target"].cpu().numpy(),high_prediction=te["high_reconstruction"].cpu().numpy(),mask=te["mask"].cpu().numpy(),gsd=test["gsd"],labels=test["labels"],test_features=te["features"].cpu().numpy(),train_features=tr["features"].cpu().numpy(),train_labels=train["labels"],train_gsd=train["gsd"]); print("inference=",out/"reconstruction.npz")
if __name__=="__main__": main()