"""
Pulmo inference examples.

Two ways to use Pulmo:

  A) FULL SCAN (recommended) — both stages, raw volume in, findings out.
     See `analyze_scan.py` (Stage 1 finds candidates, Stage 2 characterises them).

  B) SINGLE PATCH — Stage 2 only, when you already have a candidate location
     (e.g. you supply your own detector or use LUNA16 candidates.csv).
     That is what this script demonstrates.

Requires: torch, numpy, huggingface_hub  (scipy too for the full pipeline)
    pip install torch numpy scipy huggingface_hub
"""

import numpy as np
import torch
from huggingface_hub import hf_hub_download

from modeling import load_stage2, crop_stage2_input, explain_malignancy, CONCEPT_NAMES

REPO_ID = "ariyul/Pulmo"


def main():
    device = "cuda" if torch.cuda.is_available() else "cpu"

    ckpt_path = hf_hub_download(repo_id=REPO_ID, filename="student_2p5d_best.pth")
    model = load_stage2(ckpt_path, device=device)

    # --- Replace this with a real 64^3 patch cropped around a candidate ---
    # (Z, Y, X) raw HU. If you only have a candidate (z, y, x) in a full volume,
    # use crop_stage2_input(volume, (z, y, x)) instead of building the patch here.
    dummy_patch = np.random.randint(-1000, 400, size=(64, 64, 64)).astype(np.int16)
    x = crop_stage2_input(dummy_patch, (32, 32, 32)).to(device)   # (1, 7, 64, 64)

    with torch.no_grad():
        out = model(x)

    det_p = torch.softmax(out["detection"][0], 0)[1].item()
    mal_p = torch.softmax(out["malignancy"][0], 0)[1].item()
    seg = torch.sigmoid(out["segmentation"][0, 0]).cpu().numpy()

    print(f"Nodule probability     : {det_p:.3f}")
    print(f"Malignancy probability : {mal_p:.3f}  -> {'MALIGNANT' if mal_p >= 0.5 else 'BENIGN'}")
    print(f"Segmented voxels (>0.5): {(seg > 0.5).sum()}")

    print("\nPer-concept contribution to the malignancy decision:")
    for name, value, contrib in explain_malignancy(model, out):
        print(f"  {name:18s}  value={value:+.2f}  contribution={contrib:+.3f}")

    print("\nFor full-scan inference (Stage 1 + Stage 2), see analyze_scan.py.")


if __name__ == "__main__":
    main()