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              # WSI 원본 레벨 좌표
    y_wsi: int
    x_selected: int         # 선택 레벨(20x) 좌표
    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"WSI 열기 실패 {wsi_path}: {e}")


def compute_target_level(slide: openslide.OpenSlide, target_mpp: float = 0.5) -> tuple[int, float]:
    """20x 배율(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:
    """HSV saturation 기반 조직 비율. 배경(흰색·낮은 saturation) 제외.
    
    실전에서는 Otsu threshold on saturation 추가 권장.
    """
    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]]:
    """WSI를 224×224 타일로 격자 분할. 배경 자동 제외.
    
    Generator로 스트리밍 → 메모리 부담 최소화.
    """
    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):
            # 원본 레벨 좌표
            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 색 정규화.
    
    target_stain_matrix / target_max_conc가 None이면 첫 타일에서 학습.
    실전에서는 staintools · torchstain 라이브러리 활용.
    """
    try:
        from torchstain import MacenkoNormalizer
    except ImportError:
        raise RuntimeError("torchstain 필요 (pip install torchstain)")
    import torch
    
    normalizer = MacenkoNormalizer(backend="numpy")
    if target_stain_matrix is None:
        # 첫 타일로 target 학습
        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)  # 정규화 실패 시 원본 유지
    
    return normalized, normalizer.HERef, normalizer.maxCRef

import torch
from transformers import AutoImageProcessor, AutoModel


class Virchow2Embedder:
    """Paige.AI Virchow2 (또는 UNI) 병리 파운데이션 임베딩."""

    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()
        # Virchow2 hidden_dim은 모델 config에서 확인 (실행 시)
        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)
            # 표준 CLS token 임베딩 사용. 일부 모델은 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]]:
    """WSI 하나를 스트리밍하며 batch 임베딩 → (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 = [], []
    # 남은 buffer flush
    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,       # cluster 수 (규모별: 100만 → 4k, 1억 → 65k)
    m: int = 32,             # subquantizer 수 (PQ)
    nbits: int = 8,           # 각 subquantizer bit
    training_sample: int = 100_000,
) -> faiss.Index:
    """IVF-PQ: 대용량 압축 인덱스.
    
    hidden_dim=1280 × 100만 × float32 = 5GB → PQ 후 약 40MB 압축.
    """
    dim = embeddings.shape[1]
    quantizer = faiss.IndexFlatIP(dim)
    index = faiss.IndexIVFPQ(quantizer, dim, nlist, m, nbits, faiss.METRIC_INNER_PRODUCT)
    # 학습 표본 (전체가 너무 크면 랜덤 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 학습 시작 (nlist={nlist}, m={m}, sample={sample_size})")
    index.train(train_sample)
    print("학습 완료")
    return index


def build_hnsw_index(
    embeddings: np.ndarray,
    m: int = 32,
    ef_construction: int = 200,
) -> faiss.Index:
    """HNSW: 정확도·속도 우수, 메모리 상대적으로 큼. 100만 이하 규모 권장."""
    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:
    """대용량 임베딩을 청크로 index.add (메모리 관리)."""
    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:
    """MMAP 가능한 인덱스 저장."""
    faiss.write_index(index, str(path))


def load_index_mmap(path: Path) -> faiss.Index:
    """MMAP 로딩 (메모리 절약)."""
    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,      # IVF probe 수
    ef_search: int = 128,    # HNSW efSearch
) -> list[dict]:
    """쿼리 타일 하나 → top-K 유사 타일 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,   # ※ CONCH 이미지 encoder로 만든 인덱스여야 매칭
    metadata: list[TileMetadata],
    k: int = 20,
) -> list[dict]:
    """자연어 쿼리 → 유사 이미지 타일.
    
    ※ 주의: 이 index는 CONCH image encoder 임베딩으로 만들어야 함.
       Virchow2 이미지 인덱스와는 임베딩 공간 다름.
    """
    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 리스트 → FAISS 인덱스 + 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"      타일 {len(meta)}")
        except Exception as e:
            print(f"      실패: {e}")
            continue
    
    if not all_embeddings:
        raise RuntimeError("임베딩 결과 없음")
    concat_emb = np.concatenate(all_embeddings, axis=0)
    print(f"총 타일 {len(concat_emb)}, dim={concat_emb.shape[1]}")
    
    # 규모 판단해서 인덱스 종류 선택
    if len(concat_emb) < 500_000:
        print("HNSW 인덱스 (규모 500k 이하)")
        index = build_hnsw_index(concat_emb)
    else:
        print("IVF-PQ 인덱스 (규모 500k+)")
        index = build_ivf_pq_index(concat_emb)
        add_embeddings_in_chunks(index, [concat_emb])
    
    save_index_mmap(index, work_dir / "wsi_index.faiss")
    # metadata도 저장 (msgpack · jsonl 등)
    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"저장 완료: {work_dir}")
    return index, all_metadata


# 실행 예시 (TCGA-BRCA 100 WSI 시나리오)
# 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:
    """자연어 쿼리 세트로 각 top-5 결과 저장."""
    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)