import glob
import os

import numpy as np
import torch
from huggingface_hub import snapshot_download

from honeybee.processors.radiology import RadiologyProcessor

# Download CT DICOM samples from HuggingFace
data_dir = snapshot_download(
    repo_id="Lab-Rasool/honeybee-samples",
    repo_type="dataset",
    allow_patterns="CT/**",
)

# Find the DICOM series directory
dcm_files = glob.glob(os.path.join(data_dir, "CT", "**", "*.dcm"), recursive=True)
dicom_dir = os.path.dirname(dcm_files[0])
print(f"Found {len(dcm_files)} DICOM files in {dicom_dir}")

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Device: {DEVICE}")

processor = RadiologyProcessor(device=DEVICE)
image, metadata = processor.load_dicom(dicom_dir)

print(f"Volume shape:        {image.shape}")
print(f"Volume dtype:        {image.dtype}")
print(f"Value range:         [{image.min():.1f}, {image.max():.1f}]")
print(f"Modality:            {metadata.modality}")
print(f"Patient ID:          {metadata.patient_id}")
print(f"Pixel spacing:       {metadata.pixel_spacing}")
print(f"Slice thickness:     {metadata.slice_thickness}")

hu_results = processor.verify_hounsfield_units(image, metadata)

print(f"Is HU:          {hu_results['is_hu']}")
print(f"Min value:      {hu_results['min_value']:.1f}")
print(f"Max value:      {hu_results['max_value']:.1f}")
print(f"Mean value:     {hu_results['mean_value']:.1f}")
print(f"Air present:    {hu_results['likely_air_present']}")
print(f"Bone present:   {hu_results['likely_bone_present']}")

# Built-in presets: lung, mediastinum, soft_tissue, bone, abdomen, liver
preset_slices = {
    "lung":        SLICE_HEART,
    "mediastinum": SLICE_CARINA,
    "soft_tissue": SLICE_BASES,
    "bone":        SLICE_APEX,
    "abdomen":     SLICE_LIVER,
    "liver":       SLICE_LIVER,
}

fig, axes = plt.subplots(2, 3, figsize=(15, 10))
for ax, (preset, slice_idx) in zip(axes.flat, preset_slices.items()):
    windowed = processor.apply_window(image[slice_idx], window=preset)
    ax.imshow(windowed, cmap="gray")
    ax.set_title(f"{preset.replace('_', ' ').title()} (slice {slice_idx})")
    ax.axis("off")

tissue_slice = image[SLICE_CARINA]
nlm_denoised = processor.denoise(tissue_slice, method="nlm")
bilateral_denoised = processor.denoise(tissue_slice, method="bilateral")

fig, axes = plt.subplots(1, 3, figsize=(15, 5))
for ax, img, title in zip(
    axes,
    [tissue_slice, nlm_denoised, bilateral_denoised],
    ["Original", "Non-Local Means", "Bilateral"],
):
    windowed = processor.apply_window(img, window="soft_tissue")
    ax.imshow(windowed, cmap="gray")
    ax.set_title(title)
    ax.axis("off")

methods = ["zscore", "minmax", "percentile"]
norm_slice = image[SLICE_HEART]

fig, axes = plt.subplots(1, 3, figsize=(15, 5))
for ax, method in zip(axes, methods):
    normalized = processor.normalize_intensity(norm_slice, method=method)
    ax.imshow(normalized, cmap="gray")
    ax.set_title(f"{method} — range [{normalized.min():.2f}, {normalized.max():.2f}]")
    ax.axis("off")

preprocessed = processor.preprocess(image, metadata)

print(f"Original shape:      {image.shape}")
print(f"Preprocessed shape:  {preprocessed.shape}")
print(f"Original range:      [{image.min():.1f}, {image.max():.1f}]")
print(f"Preprocessed range:  [{preprocessed.min():.3f}, {preprocessed.max():.3f}]")

# Resample to isotropic 1mm spacing
resampled = processor.resample(image, metadata, new_spacing=(1.0, 1.0, 1.0))

print(f"Original shape:    {image.shape}")
print(f"Original spacing:  {metadata.pixel_spacing}")
print(f"Resampled shape:   {resampled.shape}")
print("Target spacing:    (1.0, 1.0, 1.0)")

# Crop to body ROI (remove surrounding air)
body_mask = (image > -500).astype(np.uint8)
cropped = processor.crop_to_roi(image, body_mask)

print(f"Original shape:  {image.shape}")
print(f"Cropped shape:   {cropped.shape}")

lung_mask = processor.segment_lungs(image, spacing=metadata.pixel_spacing)

print(f"Mask shape:     {lung_mask.shape}")
print(f"Mask dtype:     {lung_mask.dtype}")
print(f"Lung voxels:    {lung_mask.sum():,} ({lung_mask.mean():.2%} of volume)")

# Apply mask to isolate lung parenchyma
lung_only = processor.apply_mask(image, lung_mask)

processor_rd = RadiologyProcessor(model="rad-dino", device=DEVICE)
embeddings = processor_rd.generate_embeddings(image, mode="2d", n_slices=16)

print(f"Embeddings shape: {embeddings.shape}")
print(f"Embeddings norm:  {np.linalg.norm(embeddings):.4f}")

# Generate per-slice embeddings
slice_indices = np.linspace(0, image.shape[0] - 1, 16, dtype=int)
slice_embeddings = np.stack(
    [processor_rd.generate_embeddings(image[i], mode="2d") for i in slice_indices]
)
print(f"Per-slice embeddings: {slice_embeddings.shape}")

# Available methods: mean, max, median, std, concat
for method in ["mean", "max", "median", "std", "concat"]:
    agg = processor_rd.aggregate_embeddings(slice_embeddings, method=method)
    print(f"  {method:>8s}: shape={agg.shape}, norm={np.linalg.norm(agg):.2f}")

from honeybee import HoneyBee

honeybee = HoneyBee(config={"radiology": {"model": "rad-dino", "device": DEVICE}})
result = honeybee.process_radiology(
    dicom_path=dicom_dir, preprocess=True, generate_embeddings=True
)

print(f"Result keys:       {list(result.keys())}")
print(f"Image shape:       {result['image'].shape}")
print(f"Metadata modality: {result['metadata'].modality}")
if "embeddings" in result:
    print(f"Embeddings shape:  {result['embeddings'].shape}")

import glob
import os

import numpy as np
import torch
from huggingface_hub import snapshot_download

from honeybee.processors.radiology import RadiologyProcessor

# 1. Download and locate DICOM series
data_dir = snapshot_download(
    repo_id="Lab-Rasool/honeybee-samples",
    repo_type="dataset",
    allow_patterns="CT/**",
)
dcm_files = glob.glob(os.path.join(data_dir, "CT", "**", "*.dcm"), recursive=True)
dicom_dir = os.path.dirname(dcm_files[0])
device = "cuda" if torch.cuda.is_available() else "cpu"

# 2. Load DICOM
processor = RadiologyProcessor(device=device)
image, metadata = processor.load_dicom(dicom_dir)

# 3. Verify HU scale
hu = processor.verify_hounsfield_units(image, metadata)
print(f"HU verified: {hu['is_hu']}")

# 4. Preprocess (denoise + normalize)
preprocessed = processor.preprocess(image, metadata)

# 5. Resample to isotropic spacing
resampled = processor.resample(preprocessed, metadata, new_spacing=(1.0, 1.0, 1.0))

# 6. Segment lungs and apply mask
lung_mask = processor.segment_lungs(image, spacing=metadata.pixel_spacing)
masked = processor.apply_mask(resampled, lung_mask)

# 7. Generate embeddings
processor_rd = RadiologyProcessor(model="rad-dino", device=device)
embeddings = processor_rd.generate_embeddings(masked, mode="2d", n_slices=16)

# 8. Aggregate to volume level
volume_embedding = processor_rd.aggregate_embeddings(
    embeddings.reshape(1, -1) if embeddings.ndim == 1
    else embeddings,
    method="mean",
)
print(f"Volume embedding: {volume_embedding.shape}")