"""Score a sentence against a chest CT volume and save its similarity map as NIfTI.

python inference.py --volume volume.nii.gz --text "There is a lung nodule."
"""

import argparse
import os
import warnings

import torch
from transformers import AutoModel, AutoTokenizer

from utils import (
    load_ct_volume,
    model_inference,
    save_similarity_map,
    similarity_scores_to_volume,
)

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)
    model = AutoModel.from_pretrained(REPO_DIR, trust_remote_code=True, dtype=dtype).to(device).eval()
    return {"tokenizer": tokenizer, "model": model}


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--volume", default="volume.nii.gz", help="NIfTI CT volume or DICOM series directory")
    parser.add_argument("--text", default="There is a lung nodule.")
    parser.add_argument("--output", default="similarity_map.nii.gz", help="similarity map on the input CT grid (NIfTI)")
    parser.add_argument(
        "--raw_spacing",
        type=float,
        nargs=3,
        default=None,
        help="voxel spacing along the reoriented I, P, L axes (e.g. 1.5 0.7 0.7) if the input affine is unreliable",
    )
    args = parser.parse_args()

    models = load_model()
    pixel_values, geometry = load_ct_volume(args.volume, raw_spacing=args.raw_spacing, return_geometry=True)

    prob, scores = model_inference(pixel_values, args.text, **models)
    print(f"{args.text!r}: probability {prob.item():.3f}")
    similarity_map = similarity_scores_to_volume(scores, models["model"], geometry)
    save_similarity_map(similarity_map, geometry["affine"], args.output)
    print(f"similarity map {tuple(similarity_map.shape)} saved to {args.output}")


if __name__ == "__main__":
    main()