"""Score a sentence against a chest X-ray and save its similarity map.

python inference.py --image cxr.png --text "There is a right pleural effusion."
"""

import argparse
import os
import warnings

import torch
from transformers import AutoImageProcessor, AutoModel, AutoTokenizer

from utils import model_inference, save_similarity_map

warnings.filterwarnings("ignore", category=UserWarning)

REPO_DIR = os.path.dirname(os.path.abspath(__file__))


def load_model(device=None, dtype=torch.float32):
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    tokenizer = AutoTokenizer.from_pretrained(REPO_DIR)
    image_processor = AutoImageProcessor.from_pretrained(REPO_DIR)
    model = AutoModel.from_pretrained(REPO_DIR, trust_remote_code=True, dtype=dtype).to(device).eval()
    return {"tokenizer": tokenizer, "image_processor": image_processor, "model": model}


def main():
    parser = argparse.ArgumentParser(description="Score a sentence against a chest X-ray and save its similarity map.")
    parser.add_argument("--image", default="cxr.png", help="PNG / JPG or DICOM file")
    parser.add_argument("--text", default="There is effusion.")
    parser.add_argument("--output", default="similarity_map.png", help="heat-map overlay")
    args = parser.parse_args()

    models = load_model()
    prob, similarity_map = model_inference(args.image, args.text, **models)
    print(f"{args.text!r}: probability {prob.item():.3f}")
    save_similarity_map(args.image, similarity_map, args.output)
    print(f"similarity map saved to {args.output}")


if __name__ == "__main__":
    main()