"""Load FastFill (Qwen3-8B + LoRA adapter + geometry heads) and lay out one room.

    pip install torch  # for your CUDA / platform first
    pip install -r requirements.txt
    python predict.py example_condition.json                  # Qwen/Qwen3-8B from the Hugging Face Hub
    python predict.py example_condition.json --qwen /path/to/Qwen3-8B --device cuda

Prints the layout as JSON: for every requested object id, ``target_size_local_m`` [depth, width, height] (m),
``bottom_center_m`` [x, y, z] (m, room frame) and ``yaw_rad`` (the direction its front faces, counter-clockwise from +x).
"""
import argparse
import json
import sys
from dataclasses import replace
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE / "code"))

import torch  # noqa: E402

from fastfill.v2.batch import load_tokenizer  # noqa: E402
from fastfill.v2.evaluate import predict_layout  # noqa: E402
from fastfill.v2.io import load_checkpoint_config  # noqa: E402
from fastfill.v2.model import ModelConfig, StructuredFastFillModel, _qwen_backbone  # noqa: E402

CHECKPOINT = HERE / "checkpoint"


def load(qwen: str = "Qwen/Qwen3-8B", device: str = "cuda"):
    """(model, tokenizer, max_length). ``qwen``: the base model, a Hub id or a local directory."""
    config = ModelConfig(**json.loads((CHECKPOINT / "model_config.json").read_text()))
    config = replace(config, backbone=qwen, local_files_only=Path(qwen).is_dir())
    from peft import PeftModel
    backbone = _qwen_backbone(replace(config, lora_rank=0))  # the frozen base
    backbone = PeftModel.from_pretrained(backbone, CHECKPOINT / "backbone")  # FastFill's LoRA adapter
    model = StructuredFastFillModel(config, backbone)
    state = torch.load(CHECKPOINT / "geometry_model.pt", map_location="cpu", weights_only=True)
    missing = model.load_state_dict(state, strict=False)
    if missing.unexpected_keys or any(not k.startswith("backbone.") for k in missing.missing_keys):
        raise ValueError("geometry heads do not match the model configuration")
    tokenizer = load_tokenizer(str(HERE / "tokenizer"), local_files_only=True)
    max_length = load_checkpoint_config(CHECKPOINT)["max_length"] or 8192
    return model.to(device).eval(), tokenizer, max_length


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("condition", help="a fastfill.v2 condition JSON (room, objects, constraints)")
    ap.add_argument("--qwen", default="Qwen/Qwen3-8B", help="base model: Hub id or local Qwen3-8B directory")
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    ap.add_argument("--decode", default="spread", choices=("spread", "argmax"), help="position grid decoding")
    a = ap.parse_args()
    model, tokenizer, max_length = load(a.qwen, a.device)
    condition = json.loads(Path(a.condition).read_text())
    with torch.no_grad():
        layout = predict_layout(model, tokenizer, condition, max_length=max_length, device=a.device,
                                grid_decode=a.decode)
    print(json.dumps(layout, indent=1))


if __name__ == "__main__":
    main()