from pathlib import Path
from dataclasses import dataclass, asdict
from typing import Iterator

import openslide
import numpy as np
from PIL import Image


@dataclass
class TileMetadata:
    wsi_id: str
    slide_index: int
    level: int
    x_wsi: int              # Original WSI level coordinates
    y_wsi: int
    x_selected: int         # Selected level (20x) coordinates
    y_selected: int
    size: int
    tissue_ratio: float


def open_wsi(wsi_path: Path) -> openslide.OpenSlide:
    try:
        return openslide.OpenSlide(str(wsi_path))
    except openslide.OpenSlideError as e:
        raise RuntimeError(f"Failed to open WSI {wsi_path}: {e}")


def compute_target_level(slide: openslide.OpenSlide, target_mpp: float = 0.5) -> tuple[int, float]:
    """Selects the level closest to 20x magnification (0.5 mpp)."""
    mpp_x = float(slide.properties.get(openslide.PROPERTY_NAME_MPP_X, 0.25))
    downsample = target_mpp / mpp_x
    level = slide.get_best_level_for_downsample(downsample)
    return level, slide.level_downsamples[level]


def compute_tissue_ratio(tile: np.ndarray) -> float:
    """Tissue ratio based on HSV saturation. Excludes background (white, low saturation).

    It is recommended to add Otsu thresholding on saturation in practice.
    """
    from skimage.color import rgb2hsv
    hsv = rgb2hsv(tile)
    saturation = hsv[..., 1]
    tissue_mask = saturation > 0.05
    return float(tissue_mask.mean())


def tile_wsi(
    wsi_path: Path,
    wsi_id: str,
    slide_index: int,
    tile_size: int = 224,
    target_mpp: float = 0.5,
    tissue_threshold: float = 0.2,
    output_dir: Path | None = None,
    save_tiles: bool = False,
) -> Iterator[tuple[TileMetadata, np.ndarray]]:
    """Divides the WSI into 224×224 tiles in a grid, automatically excluding background.

    Streaming with a generator minimizes memory usage.
    """
    slide = open_wsi(wsi_path)
    level, level_downsample = compute_target_level(slide, target_mpp)
    width, height = slide.level_dimensions[level]
    print(f"[{wsi_id}] level={level}, dims={width}x{height}, mpp≈{target_mpp}")

    for y in range(0, height - tile_size, tile_size):
        for x in range(0, width - tile_size, tile_size):
            # Original level coordinates
            x_wsi = int(x * level_downsample)
            y_wsi = int(y * level_downsample)
            try:
                tile_pil = slide.read_region((x_wsi, y_wsi), level, (tile_size, tile_size)).convert("RGB")
            except openslide.OpenSlideError:
                continue
            tile_np = np.asarray(tile_pil)
            tissue_ratio = compute_tissue_ratio(tile_np)
            if tissue_ratio < tissue_threshold:
                continue
            meta = TileMetadata(
                wsi_id=wsi_id, slide_index=slide_index, level=level,
                x_wsi=x_wsi, y_wsi=y_wsi,
                x_selected=x, y_selected=y,
                size=tile_size, tissue_ratio=tissue_ratio,
            )
            if save_tiles and output_dir:
                out = output_dir / f"{wsi_id}_x{x}_y{y}.png"
                out.parent.mkdir(parents=True, exist_ok=True)
                tile_pil.save(out, format="PNG")
            yield meta, tile_np
    slide.close()

def macenko_normalize(
    tiles: list[np.ndarray],
    target_stain_matrix: np.ndarray | None = None,
    target_max_conc: np.ndarray | None = None,
    alpha: float = 1.0,
    beta: float = 0.15,
) -> tuple[list[np.ndarray], np.ndarray, np.ndarray]:
    """Macenko color normalization.

    If target_stain_matrix / target_max_conc is None, it learns from the first tile.
    In practice, use staintools or torchstain libraries.
    """
    try:
        from torchstain import MacenkoNormalizer
    except ImportError:
        raise RuntimeError("torchstain is required (pip install torchstain)")
    import torch

    normalizer = MacenkoNormalizer(backend="numpy")
    if target_stain_matrix is None:
        # Learn the target from the first tile
        normalizer.fit(np.transpose(tiles[0], (2, 0, 1)))
    else:
        normalizer.HERef = target_stain_matrix
        normalizer.maxCRef = target_max_conc

    normalized = []
    for tile in tiles:
        try:
            tile_t = np.transpose(tile, (2, 0, 1))
            norm_t, _, _ = normalizer.normalize(tile_t, stains=False)
            normalized.append(np.transpose(np.asarray(norm_t), (1, 2, 0)))
        except Exception:
            normalized.append(tile)  # Keep the original if normalization fails
    return normalized, normalizer.HERef, normalizer.maxCRef

import torch
from transformers import AutoImageProcessor, AutoModel


class Virchow2Embedder:
    """Paige.AI Virchow2 (or UNI) pathology foundation embedding."""

    MODEL_ID = "paige-ai/Virchow2"

    def __init__(self, device: str = "cuda", torch_dtype: torch.dtype = torch.float16):
        self.device = device
        self.processor = AutoImageProcessor.from_pretrained(self.MODEL_ID)
        self.model = AutoModel.from_pretrained(
            self.MODEL_ID,
            torch_dtype=torch_dtype,
        ).to(device).eval()
        # Check the Virchow2 hidden_dim in the model config (at runtime)
        self.hidden_dim = getattr(self.model.config, "hidden_size", 1280)

    @torch.no_grad()
    def embed_batch(self, tiles: list[np.ndarray], batch_size: int = 32) -> np.ndarray:
        """(N, 224, 224, 3) uint8 → (N, hidden_dim) fp16."""
        if not tiles:
            return np.zeros((0, self.hidden_dim), dtype=np.float16)
        embeddings = []
        for i in range(0, len(tiles), batch_size):
            chunk = tiles[i:i + batch_size]
            inputs = self.processor(images=chunk, return_tensors="pt").to(self.device)
            outputs = self.model(**inputs)
            # Use the standard CLS token embedding. Some models recommend mean pooling.
            cls_emb = outputs.last_hidden_state[:, 0, :]
            embeddings.append(cls_emb.cpu().numpy())
        return np.concatenate(embeddings, axis=0).astype(np.float16)


def stream_and_embed_wsi(
    wsi_path: Path,
    wsi_id: str,
    slide_index: int,
    embedder: Virchow2Embedder,
    normalize_color: bool = True,
    batch_size: int = 32,
) -> tuple[np.ndarray, list[TileMetadata]]:
    """Streams one WSI and extracts batch embeddings → (N, dim), metadata."""
    tile_buffer: list[np.ndarray] = []
    meta_buffer: list[TileMetadata] = []
    all_embeddings: list[np.ndarray] = []
    all_metadata: list[TileMetadata] = []

    for meta, tile in tile_wsi(wsi_path, wsi_id, slide_index):
        tile_buffer.append(tile)
        meta_buffer.append(meta)
        if len(tile_buffer) >= batch_size:
            if normalize_color:
                tile_buffer, _, _ = macenko_normalize(tile_buffer)
            emb = embedder.embed_batch(tile_buffer, batch_size=batch_size)
            all_embeddings.append(emb)
            all_metadata.extend(meta_buffer)
            tile_buffer, meta_buffer = [], []
    # Flush the remaining buffer
    if tile_buffer:
        if normalize_color:
            tile_buffer, _, _ = macenko_normalize(tile_buffer)
        emb = embedder.embed_batch(tile_buffer, batch_size=batch_size)
        all_embeddings.append(emb)
        all_metadata.extend(meta_buffer)

    if not all_embeddings:
        return np.zeros((0, embedder.hidden_dim), dtype=np.float16), []
    return np.concatenate(all_embeddings, axis=0), all_metadata

import faiss


def build_ivf_pq_index(
    embeddings: np.ndarray,
    nlist: int = 4096,       # Number of clusters (scale-dependent: 1 million → 4k, 100 million → 65k)
    m: int = 32,             # Number of subquantizers (PQ)
    nbits: int = 8,           # Bits per subquantizer
    training_sample: int = 100_000,
) -> faiss.Index:
    """IVF-PQ: Large-scale compression index.

    hidden_dim=1280 × 1 million × float32 = 5GB → Approximately 40MB compressed with PQ.
    """
    dim = embeddings.shape[1]
    quantizer = faiss.IndexFlatIP(dim)
    index = faiss.IndexIVFPQ(quantizer, dim, nlist, m, nbits, faiss.METRIC_INNER_PRODUCT)
    # Training sample (if the entire dataset is too large, use a random 100k)
    sample_size = min(training_sample, len(embeddings))
    sample_idx = np.random.choice(len(embeddings), sample_size, replace=False)
    train_sample = embeddings[sample_idx].astype(np.float32)
    faiss.normalize_L2(train_sample)
    print(f"IVF-PQ training starts (nlist={nlist}, m={m}, sample={sample_size})")
    index.train(train_sample)
    print("Training complete")
    return index


def build_hnsw_index(
    embeddings: np.ndarray,
    m: int = 32,
    ef_construction: int = 200,
) -> faiss.Index:
    """HNSW: Excellent accuracy and speed, relatively large memory. Recommended for scales below 1 million.

    """
    dim = embeddings.shape[1]
    index = faiss.IndexHNSWFlat(dim, m, faiss.METRIC_INNER_PRODUCT)
    index.hnsw.efConstruction = ef_construction
    embs = embeddings.astype(np.float32)
    faiss.normalize_L2(embs)
    index.add(embs)
    return index


def add_embeddings_in_chunks(
    index: faiss.Index,
    embeddings_iter: Iterator[np.ndarray],
    chunk_size: int = 100_000,
) -> None:
    """Add large embeddings in chunks to the index (memory management)."""
    for chunk in embeddings_iter:
        chunk_f32 = chunk.astype(np.float32)
        faiss.normalize_L2(chunk_f32)
        index.add(chunk_f32)


def save_index_mmap(index: faiss.Index, path: Path) -> None:
    """Save index that can be accessed with MMAP."""
    faiss.write_index(index, str(path))


def load_index_mmap(path: Path) -> faiss.Index:
    """Load with MMAP (memory saving)."""
    return faiss.read_index(str(path), faiss.IO_FLAG_MMAP | faiss.IO_FLAG_READ_ONLY)

def search_by_image(
    query_tile: np.ndarray,
    embedder: Virchow2Embedder,
    index: faiss.Index,
    metadata: list[TileMetadata],
    k: int = 20,
    nprobe: int = 128,      # Number of IVF probes
    ef_search: int = 128,    # HNSW efSearch
) -> list[dict]:
    """One query tile → top-K similar tiles metadata."""
    query_emb = embedder.embed_batch([query_tile]).astype(np.float32)
    faiss.normalize_L2(query_emb)
    if hasattr(index, "nprobe"):
        index.nprobe = nprobe
    if hasattr(index, "hnsw"):
        index.hnsw.efSearch = ef_search
    distances, indices = index.search(query_emb, k)
    results = []
    for dist, idx in zip(distances[0], indices[0]):
        if idx < 0 or idx >= len(metadata):
            continue
        m = metadata[idx]
        results.append({
            **asdict(m),
            "similarity": float(dist),
        })
    return results

from transformers import CLIPModel, CLIPProcessor


class CONCHTextEncoder:
    """CONCH text encoder wrapping."""

    MODEL_ID = "MahmoodLab/CONCH"

    def __init__(self, device: str = "cuda"):
        self.device = device
        self.processor = CLIPProcessor.from_pretrained(self.MODEL_ID)
        self.model = CLIPModel.from_pretrained(self.MODEL_ID).to(device).eval()

    @torch.no_grad()
    def encode_text(self, text: str) -> np.ndarray:
        inputs = self.processor(text=[text], return_tensors="pt", padding=True).to(self.device)
        text_emb = self.model.get_text_features(**inputs)
        return text_emb.cpu().numpy()


def search_by_text(
    query_text: str,
    text_encoder: CONCHTextEncoder,
    index: faiss.Index,   # Must be the CONCH image encoder index for matching
    metadata: list[TileMetadata],
    k: int = 20,
) -> list[dict]:
    """Natural language query → similar image tiles.

    Note: This index must be created using the CONCH image encoder embeddings.
    It is different from the Virchow2 image index in terms of embedding space.
    """
    query_emb = text_encoder.encode_text(query_text).astype(np.float32)
    faiss.normalize_L2(query_emb)
    distances, indices = index.search(query_emb, k)
    return [
        {**asdict(metadata[i]), "similarity": float(d)}
        for d, i in zip(distances[0], indices[0]) if 0 <= i < len(metadata)
    ]

def full_pipeline(
    wsi_paths: list[Path],
    work_dir: Path,
    device: str = "cuda",
) -> tuple[faiss.Index, list[TileMetadata]]:
    """WSI list → FAISS index + metadata."""
    work_dir.mkdir(parents=True, exist_ok=True)
    embedder = Virchow2Embedder(device=device)
    all_metadata: list[TileMetadata] = []
    all_embeddings: list[np.ndarray] = []

    for i, wsi_path in enumerate(wsi_paths):
        wsi_id = wsi_path.stem
        print(f"[{i+1}/{len(wsi_paths)}] {wsi_id}")
        try:
            emb, meta = stream_and_embed_wsi(
                wsi_path, wsi_id=wsi_id, slide_index=i,
                embedder=embedder, normalize_color=True,
            )
            all_embeddings.append(emb)
            all_metadata.extend(meta)
            print(f"      Tiles {len(meta)}")
        except Exception as e:
            print(f"      Failed: {e}")
            continue

    if not all_embeddings:
        raise RuntimeError("No embedding results")
    concat_emb = np.concatenate(all_embeddings, axis=0)
    print(f"Total tiles {len(concat_emb)}, dim={concat_emb.shape[1]}")

    # Determine the index type based on scale
    if len(concat_emb) < 500_000:
        print("HNSW index (scale below 500k)")
        index = build_hnsw_index(concat_emb)
    else:
        print("IVF-PQ index (scale 500k+)")
        index = build_ivf_pq_index(concat_emb)
        add_embeddings_in_chunks(index, [concat_emb])

    save_index_mmap(index, work_dir / "wsi_index.faiss")
    # Save metadata as well (msgpack, jsonl, etc.)
    import json
    with open(work_dir / "metadata.jsonl", "w") as f:
        for m in all_metadata:
            f.write(json.dumps(asdict(m)) + "\n")
    print(f"Saved to: {work_dir}")
    return index, all_metadata


# Example execution (TCGA-BRCA 100 WSI scenario)
# tcga_wsi_dir = Path("/data/tcga_brca_svs")
# wsi_files = sorted(tcga_wsi_dir.glob("*.svs"))[:100]
# index, meta = full_pipeline(wsi_files, work_dir=Path("./tcga_brca_output"))

EXAMPLE_QUERIES = [
    "invasive ductal carcinoma with dense lymphocytic infiltrate at tumor margin",
    "high-grade tumor with necrotic center and prominent nucleoli",
    "tubular carcinoma with well-formed glandular structures",
    "adjacent normal breast tissue with lobular architecture preserved",
    "medullary carcinoma with syncytial growth pattern and TIL",
    "ductal carcinoma in situ (DCIS) with cribriform architecture",
    "stromal desmoplasia with fibroblast proliferation surrounding tumor nests",
]


def demo_text_queries(
    text_encoder: CONCHTextEncoder,
    conch_index: faiss.Index,
    metadata: list[TileMetadata],
    output_dir: Path,
) -> None:
    """For each query in the set, save the top 5 results."""
    output_dir.mkdir(parents=True, exist_ok=True)
    import json
    for query in EXAMPLE_QUERIES:
        results = search_by_text(query, text_encoder, conch_index, metadata, k=5)
        out_path = output_dir / f"query_{hash(query) % 10**6}.json"
        with open(out_path, "w") as f:
            json.dump({"query": query, "top_5": results}, f, ensure_ascii=False, indent=2)

## Performance, Cost, and Known Failure Cases

### Performance Reference (Citing Public Benchmarks)

| Model | Benchmark | Linear Probing AUC | Retrieval R@10 | Source |
|------|------|:------------------:|:-------------:|------|
| ImageNet ResNet-50 | Camelyon16 | 0.83 | 0.32 | Legacy baseline |
| CTransPath | Camelyon16 | 0.91 | 0.51 | Wang et al., MedIA 2022 |
| UNI | Camelyon16 + BRACS | 0.95+ | 0.68 | Chen et al., Nat Med 2024 [1] |
| Virchow | Multi-cohort | 0.94~0.97 | 0.70 | Vorontsov et al., Nat Med 2024 [2] |
| Virchow2 | Multi-cohort | 0.96~0.98 (estimated) | 0.72 (estimated) | Paige.AI 2024 [2] |
| PLIP | Zero-shot classification | 0.87 | 0.45 (text query) | Huang et al., Nat Med 2023 [3] |
| CONCH | Zero-shot | 0.90 | 0.60 (text query) | Lu et al., Nat Med 2024 [4] |
| GigaPath (WSI-level) | Multi-task | 0.88~0.94 | — | Xu et al., Nature 2024 [5] |

### Estimated Cost for Learner Reproduction

- API Cost: 0 (completely local).
- Assuming a 16-32GB VRAM server GPU, indexing approximately 100 TCGA-BRCA WSI embeddings takes about 2-4 hours.
- FAISS IVF-PQ index training: Training on a 1 million vector sample takes 10-20 minutes (CPU).
- Natural language query response: Less than 100ms (index mmap).

### 5 Known Failure Cases (Collected from Community and Papers)

1. **Domain Shift from Different Hospital Data (Scanner/Staining Differences)**
   Symptom: A model trained and indexed on Hospital A data shows a significant drop in retrieval performance on Hospital B WSI (R@10 0.72 → 0.42).
   Cause: Differences in H&E staining protocols, scanner vendors (Aperio, Hamamatsu, Leica), and color correction. Reinhard normalization alone is insufficient.
   Mitigation: (a) Macenko/Vahadane color normalization (Step 2 in this section), (b) Fine-tuning on data from multiple hospitals, (c) Using stain augmentation, (d) Domain-adversarial training, (e) CONCH and UNI are relatively robust to domain shift due to training on diverse hospital data.
   Source: Ciompi et al. "The importance of stain normalization in colorectal tissue classification" ISBI 2017; UNI benchmark discussion [1].

2. **WSI File Format Compatibility (SVS, NDPI, MRXS, isyntax)**
   Symptom: OpenSlide cannot read specific vendor files (e.g., Philips isyntax).
   Cause: OpenSlide has limited support for some commercial formats.
   Mitigation: (a) For isyntax, use the Philips SDK or separate conversion tools (bfconvert, isyntax-cli), (b) For MRXS, use the latest version of OpenSlide, (c) Use TIAToolbox as an alternative (OpenSlide + extensions), (d) Convert to DICOM WSI (DICOM-WG 26 standard).
   Source: OpenSlide GitHub Issues [6]; TIAToolbox docs [7].

3. **FAISS IVF-PQ Training Data Bias**
   Symptom: If the PQ codebook is trained with an overrepresentation of specific tissue types, the accuracy of searching for other tissues will significantly decrease.
   Cause: Insufficient randomness in the training samples or bias towards specific cases.
   Mitigation: (a) Stratify the training samples across multiple hospitals and tissue types, (b) Set `nprobe` to a large value (256+) to ensure accuracy, (c) Re-train the index when adding new tissue types, (d) For datasets smaller than 500k, prioritize HNSW.
   Source: FAISS Wiki "IVF training pitfalls" [8].

4. **CONCH Text Query and Image Index Matching Failure (Encoder Space Discrepancy)**
   Symptom: When querying with CONCH text embeddings on an image index created with Virchow2, the results are meaningless.
   Cause: Virchow2 and CONCH use different training methods and have different embedding spaces. No contrastive alignment is performed.
   Mitigation: (a) For natural language queries, create a separate index using the CONCH image encoder, (b) Maintain a dual index (Virchow2 images, CONCH images + text), (c) Evaluate the trade-off between storage/search cost and accuracy.
   Source: CONCH GitHub usage examples [9]; Basic principles of contrastive learning.

5. **Large TCGA WSI Download and Storage Burden**
   Symptom: The total size of 1000+ TCGA-BRCA WSI is several TB. Learners face time and disk space constraints when downloading locally.
   Cause: Each TCGA WSI averages 500MB to 2GB.
   Mitigation: (a) Use a TCGA subset (e.g., 100 slides instead of the full dataset), (b) Use the GDC (Genomic Data Commons) API to stream processing and store only the index, deleting the originals, (c) Process directly from cloud storage (S3, GCS), (d) Utilize academic cloud consortia (e.g., NCI Cloud Resource).
   Source: TCGA / GDC official documentation [10].

## Expansion Ideas

- **Slide-level embedding:** Aggregate tile embeddings into slide-level embeddings using attention pooling (e.g., ABMIL, GigaPath) for slide-level search and classification. Similar thinking as Section 12 Geneformer approach.
- **Multi-modal fusion:** Use late fusion to re-rank search results by combining WSI with clinical data (age, stage, mutation).
- **Survival prediction:** Collect outcomes from similar cases found through search and estimate the prognosis for new cases.
- **Federated search:** Share only embeddings of data from multiple hospitals without exporting the original data, enabling federated similarity search indexing.
- **QuPath plugin:** Allow pathologists to select a region in QuPath, automatically search for similar cases, and display the results.
- **Natural language re-rank:** Re-rank the top-100 image search candidates based on their relevance to the natural language query.

## Next Section

- Section 03 `cellpose-sam-segmentation`: Combine the WSI tiles from this section with individual cell segmentation and cell-level embeddings.
- Section 10 `med-llm-reproduction`: A framework for reproducing the CONCH and PLIP benchmarks.
- Section 12 `single-cell-perturbation`: Integrate pathology images with single-cell data (linking to spatial transcriptomics).
- Section 14 `bio-mcp-agent`: Expose pathology search as an MCP tool, allowing autonomous execution of tasks like "Find similar cases to this case."

## References

1. Chen RJ, Ding T, Lu MY, et al. "Towards a general-purpose foundation model for computational pathology (UNI)." Nature Medicine 2024. `https://www.nature.com/articles/s41591-024-02857-3`
2. Vorontsov E, Bozkurt A, Casson A, et al. "A foundation model for clinical-grade computational pathology (Virchow)." Nature Medicine 2024. `https://www.nature.com/articles/s41591-024-03141-0` / Virchow2 follow-up release information: `https://huggingface.co/paige-ai/Virchow2`
3. Huang Z, Bianchi F, Yuksekgonul M, et al. "A visual-language foundation model for pathology image analysis using medical Twitter (PLIP)." Nature Medicine 2023. `https://www.nature.com/articles/s41591-023-02504-3`
4. Lu MY, Chen B, Williamson DFK, et al. "A visual-language foundation model for computational pathology (CONCH)." Nature Medicine 2024. `https://www.nature.com/articles/s41591-024-02856-4`
5. Xu H, Usuyama N, Bagga J, et al. "A whole-slide foundation model for digital pathology from real-world data (GigaPath)." Nature 2024. `https://www.nature.com/articles/s41586-024-07441-w`
6. OpenSlide GitHub Issues: `https://github.com/openslide/openslide/issues`
7. TIAToolbox documentation: `https://tia-toolbox.readthedocs.io/`
8. FAISS Wiki (IVF-PQ tuning): `https://github.com/facebookresearch/faiss/wiki`
9. HuggingFace MahmoodLab CONCH: `https://huggingface.co/MahmoodLab/CONCH`
10. TCGA / GDC data portal: `https://portal.gdc.cancer.gov/`
11. Camelyon16 challenge dataset: `https://camelyon16.grand-challenge.org/`
12. QuPath (open pathology viewer): `https://qupath.github.io/`
13. Macenko M et al. "A method for normalizing histology slides for quantitative analysis." ISBI 2009.
14. Vahadane A et al. "Structure-Preserving Color Normalization and Sparse Stain Separation for Histological Images." IEEE TMI 2016.
15. torchstain (Macenko implementation): `https://github.com/EIDOSLAB/torchstain`
16. FAISS GitHub: `https://github.com/facebookresearch/faiss`
17. HuggingFace paige-ai/Virchow2: `https://huggingface.co/paige-ai/Virchow2`
18. HuggingFace MahmoodLab UNI: `https://huggingface.co/MahmoodLab/UNI`
19. DICOM WSI (WG 26 standard): `https://dicom.nema.org/`
20. staintools (alternative color normalization): `https://github.com/Peter554/StainTools`