"""
Pulmo — end-to-end two-stage inference.

    raw CT volume  ->  Stage 1 (find nodule centres)
                   ->  Stage 2 (per-candidate diagnosis + concept explanation)

`analyze_scan` returns one dict per detected nodule with its location, the
detection / malignancy probabilities, the segmentation mask of the central
slice, and the concept-bottleneck explanation of the malignancy decision.

Requires: torch, numpy, scipy, huggingface_hub
    pip install torch numpy scipy huggingface_hub
"""

import numpy as np
import torch
from huggingface_hub import hf_hub_download

from modeling import (
    load_stage1, load_stage2, find_candidates,
    crop_stage2_input, explain_malignancy,
)

REPO_ID = "ariyul/Pulmo"
STAGE1_FILE = "stage1_detector_v2.pth"
STAGE2_FILE = "student_2p5d_best.pth"


def load_pipeline(device=None):
    """Download both checkpoints from the Hub and return (stage1, stage2, device)."""
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    s1_path = hf_hub_download(REPO_ID, STAGE1_FILE)
    s2_path = hf_hub_download(REPO_ID, STAGE2_FILE)
    stage1 = load_stage1(s1_path, device=device)
    stage2 = load_stage2(s2_path, device=device)
    return stage1, stage2, device


@torch.no_grad()
def analyze_scan(volume, spacing, stage1, stage2, device="cpu",
                 peak_thresh=0.1, det_thresh=0.5, max_candidates=None):
    """Run the full Pulmo pipeline on one CT volume.

    Args:
        volume       : (Z, Y, X) numpy array of raw HU values.
        spacing      : (sz, sy, sx) voxel spacing in mm, [z, y, x] order.
        stage1/stage2: loaded models (see load_pipeline).
        peak_thresh  : Stage-1 heatmap threshold (lower -> more candidates;
                       Stage 2 filters false positives).
        det_thresh   : keep a candidate only if Stage-2 detection prob >= this.
        max_candidates: optionally cap how many Stage-1 candidates to characterise.

    Returns:
        list of dicts (one per kept nodule), sorted by malignancy probability:
          {
            'location_voxel'   : (z, y, x),
            'detection_prob'   : float,
            'malignancy_prob'  : float,
            'prediction'       : 'MALIGNANT' | 'BENIGN',
            'segmentation'     : (64, 64) float mask of the central slice,
            'concepts'         : {name: value, ...},
            'top_reasons'      : [(concept, value, contribution), ...],  # most malignancy-driving first
          }
    """
    cands = find_candidates(stage1, volume, spacing, device=device,
                            peak_thresh=peak_thresh)
    if max_candidates is not None:
        cands = cands[:max_candidates]

    findings = []
    for (z, y, x) in cands:
        xin = crop_stage2_input(volume, (z, y, x)).to(device)
        out = stage2(xin)
        det_p = torch.softmax(out["detection"][0], 0)[1].item()
        if det_p < det_thresh:                 # Stage 2 rejects this candidate
            continue
        mal_p = torch.softmax(out["malignancy"][0], 0)[1].item()
        seg = torch.sigmoid(out["segmentation"][0, 0]).cpu().numpy()
        concepts = out["concepts"][0].cpu().numpy()
        from modeling import CONCEPT_NAMES
        findings.append({
            "location_voxel": (int(z), int(y), int(x)),
            "detection_prob": det_p,
            "malignancy_prob": mal_p,
            "prediction": "MALIGNANT" if mal_p >= 0.5 else "BENIGN",
            "segmentation": seg,
            "concepts": {n: float(v) for n, v in zip(CONCEPT_NAMES, concepts)},
            "top_reasons": explain_malignancy(stage2, out)[:3],
        })

    findings.sort(key=lambda f: -f["malignancy_prob"])
    return findings


def main():
    stage1, stage2, device = load_pipeline()
    print(f"Pipeline loaded on {device}.")

    # --- Replace this with a real CT volume + its voxel spacing ---
    # volume: (Z, Y, X) raw HU; spacing: (sz, sy, sx) mm in [z, y, x] order.
    dummy_volume = np.random.randint(-1000, 200, size=(120, 512, 512)).astype(np.int16)
    dummy_spacing = (1.25, 0.7, 0.7)

    findings = analyze_scan(dummy_volume, dummy_spacing, stage1, stage2, device=device)
    print(f"\n{len(findings)} nodule(s) kept after Stage-2 filtering:\n")
    for i, f in enumerate(findings, 1):
        z, y, x = f["location_voxel"]
        print(f"[{i}] voxel (z={z}, y={y}, x={x})  "
              f"det={f['detection_prob']:.2f}  "
              f"malignancy={f['malignancy_prob']:.2f}  -> {f['prediction']}")
        print("    reasons:", ", ".join(
            f"{name}({contrib:+.2f})" for name, _val, contrib in f["top_reasons"]))


if __name__ == "__main__":
    main()