"""
TB Dataset Preprocessor v3 - Per-source splitting to prevent data leakage.
Uses pre-extracted IN-CXR PNGs + other dataset zips.
"""
import zipfile
import hashlib
import shutil
import csv
import io
import os
import sys
import random
import argparse
from pathlib import Path
from collections import Counter
from PIL import Image

SEED = 42
TEST_SPLIT = 0.15
VAL_SPLIT = 0.15


def md5(data):
    return hashlib.md5(data).hexdigest()


def write_image(dest_dir, stem, data):
    ext = ".png"
    try:
        img = Image.open(io.BytesIO(data))
        ext = ".png" if img.format == "PNG" else ".jpg"
    except Exception:
        pass
    path = dest_dir / f"{stem}{ext}"
    dest_dir.mkdir(parents=True, exist_ok=True)
    with open(path, "wb") as f:
        f.write(data)
    return path


def extract_images_from_zip(zf):
    images = []
    for name in zf.namelist():
        if name.endswith("/") or "__MACOSX" in name or ".DS_Store" in name:
            continue
        ext = Path(name).suffix.lower()
        if ext not in (".png", ".jpg", ".jpeg"):
            continue
        images.append((name, zf.read(name)))
    return images


def process_incxr_preprocessed(data_dir):
    """Use pre-extracted IN-CXR PNGs from TB_DATASETS/train|val|test."""
    print("\n" + "=" * 60)
    print("IN-CXR (using pre-extracted PNGs)")
    print("=" * 60)
    images = []
    for split in ["train", "val", "test"]:
        for cls_name, label in [("TB", 1), ("Normal", 0)]:
            cls_dir = data_dir / split / cls_name
            if not cls_dir.exists():
                continue
            for img_path in cls_dir.glob("*"):
                if img_path.suffix.lower() not in (".png", ".jpg", ".jpeg"):
                    continue
                with open(img_path, "rb") as f:
                    data = f.read()
                images.append((md5(data), data, label, "incxr"))
    print(f"  Total: {len(images)} images")
    return images


def process_other_datasets(downloads_dir):
    """Process all other datasets from zip files."""
    print("\n" + "=" * 60)
    print("Other Datasets")
    print("=" * 60)
    all_images = []

    zip_configs = [
        ("belarus.zip", "belarus", None),
        ("DA_DB_tbxpredict.zip", "dadb", None),
        ("DA_DB_archive.zip", "dadb", None),
        ("Mendeley_Dataset.zip", "mendeley", None),
        ("Mendeley_Pakistan_Dataset.zip", "mendeley", None),
        ("Montgomery_archive.zip", "montgomery", None),
        ("shenzhen_archive.zip", "shenzhen", None),
        ("qatar_archive.zip", "qatar", None),
        ("Sakha-TB_Russia.zip", "sakha", None),
        ("TBX11K.zip", "tbx11k", "csv"),
        ("TBX11K_archive.zip", "tbx11k", "csv"),
        ("Shenzhen + Montgomery_archive.zip", "szmc", None),
    ]

    processed = set()

    for zip_name, src_name, mode in zip_configs:
        if src_name in processed:
            continue
        zip_path = downloads_dir / zip_name
        if not zip_path.exists():
            continue

        print(f"\n[{src_name}] {zip_name}")
        zf = zipfile.ZipFile(zip_path)

        if mode == "csv":
            # TBX11K: determine labels by directory structure
            #   imgs/tb/ -> TB (label 1)
            #   imgs/health/ -> Normal (label 0)
            #   imgs/sick/ -> Normal (non-TB disease, label 0)
            #   imgs/test/ -> no labels (skip)
            #   imgs/extra/ -> no labels (skip)
            source_images = []
            label_map = {"tb": 1, "health": 0, "sick": 0}
            for name in zf.namelist():
                if name.endswith("/") or Path(name).suffix.lower() not in (".png", ".jpg", ".jpeg"):
                    continue
                parts = name.replace("\\", "/").split("/")
                if len(parts) < 3:
                    continue
                subdir = parts[2]  # e.g., "tb", "health", "sick", "test", "extra"
                if subdir not in label_map:
                    continue
                data = zf.read(name)
                source_images.append((md5(data), data, label_map[subdir], src_name))
        else:
            source_images = []
            for name, data in extract_images_from_zip(zf):
                fname = Path(name).name.lower()
                label = None
                if src_name == "belarus":
                    label = 1 if "tb" in fname else 0
                elif src_name == "dadb":
                    if fname.startswith("p") or fname.startswith("px"):
                        label = 1
                    elif fname.startswith("n") or fname.startswith("nx"):
                        label = 0
                elif src_name == "mendeley":
                    if "tb" in name or "TB" in name:
                        label = 1
                    elif "normal" in name:
                        label = 0
                elif src_name == "szmc":
                    if name.lower().endswith("_1.png") or name.lower().endswith("_1.jpg"):
                        label = 1
                    elif name.lower().endswith("_0.png") or name.lower().endswith("_0.jpg"):
                        label = 0
                elif src_name in ("montgomery", "shenzhen", "qatar", "sakha"):
                    stem = Path(name).stem.lower()
                    if stem.endswith("_1") or "tb" in stem:
                        label = 1
                    elif stem.endswith("_0") or "normal" in stem:
                        label = 0
                if label is None:
                    continue
                source_images.append((md5(data), data, label, src_name))

        zf.close()
        all_images.extend(source_images)
        tb = sum(1 for _, _, l, _ in source_images if l == 1)
        norm = len(source_images) - tb
        print(f"  -> {len(source_images)} images ({tb} TB, {norm} Normal)")
        processed.add(src_name)

    print(f"\nTotal other datasets: {len(all_images)} images")
    return all_images


def main():
    parser = argparse.ArgumentParser(description="TB Dataset Preprocessor v3")
    parser.add_argument("--downloads", type=str,
                        default=str(Path.home() / "Downloads" / "TB_DATASETS"),
                        help="Directory with TB datasets")
    parser.add_argument("--output", type=str, default="datasets_processed",
                        help="Output directory")
    args = parser.parse_args()

    DOWNLOADS = Path(args.downloads)
    PROCESSED = Path(args.output)

    if not DOWNLOADS.exists():
        print(f"ERROR: Downloads directory not found: {DOWNLOADS}")
        return

    if PROCESSED.exists():
        print(f"Removing existing {PROCESSED}...")
        shutil.rmtree(PROCESSED)

    random.seed(SEED)

    # Step 1: IN-CXR pre-extracted
    incxr = process_incxr_preprocessed(DOWNLOADS)

    # Step 2: Other datasets
    other = process_other_datasets(DOWNLOADS)

    # Step 3: Group by source
    print("\n" + "=" * 60)
    print("PER-SOURCE DEDUP + STRATIFIED SPLIT")
    print("=" * 60)

    by_source = {}
    for h, data, label, src in incxr + other:
        by_source.setdefault(src, []).append((h, data, label))

    def deduplicate(items):
        seen = {}
        for h, data, label in items:
            if h not in seen:
                seen[h] = (data, label)
        result = list(seen.values())
        random.shuffle(result)
        return result

    def split_group(items):
        random.shuffle(items)
        n = len(items)
        n_test = int(n * TEST_SPLIT)
        n_val = int(n * VAL_SPLIT)
        test = items[:n_test]
        val = items[n_test:n_test + n_val]
        train = items[n_test + n_val:]
        return train, val, test

    train_all, val_all, test_all = [], [], []

    for src, items in sorted(by_source.items()):
        deduped = deduplicate(items)
        tb = [(d, l) for d, l in deduped if l == 1]
        norm = [(d, l) for d, l in deduped if l == 0]
        tb_train, tb_val, tb_test = split_group(tb)
        norm_train, norm_val, norm_test = split_group(norm)
        train_all.extend([(d, l, src) for d, l in tb_train + norm_train])
        val_all.extend([(d, l, src) for d, l in tb_val + norm_val])
        test_all.extend([(d, l, src) for d, l in tb_test + norm_test])
        n_tb = sum(1 for _, l in deduped if l == 1)
        n_norm = len(deduped) - n_tb
        print(f"  {src}: {len(deduped)} ({n_tb} TB, {n_norm} Normal) -> train {len(tb_train)+len(norm_train)} | val {len(tb_val)+len(norm_val)} | test {len(tb_test)+len(norm_test)}")

    random.shuffle(train_all)
    random.shuffle(val_all)
    random.shuffle(test_all)

    def count_tb(items):
        return sum(1 for _, l, _ in items if l == 1)

    def count_norm(items):
        return sum(1 for _, l, _ in items if l == 0)

    print(f"\n  Final:")
    print(f"    Train: {len(train_all)} ({count_tb(train_all)} TB, {count_norm(train_all)} Normal)")
    print(f"    Val:   {len(val_all)} ({count_tb(val_all)} TB, {count_norm(val_all)} Normal)")
    print(f"    Test:  {len(test_all)} ({count_tb(test_all)} TB, {count_norm(test_all)} Normal)")

    # Step 4: Write
    print("\nWriting datasets_processed/ ...")
    for split_name, items in [("train", train_all), ("val", val_all), ("test", test_all)]:
        tb_dir = PROCESSED / split_name / "TB"
        norm_dir = PROCESSED / split_name / "Normal"
        for data, label, src in items:
            stem = f"{src}_{md5(data)[:12]}"
            write_image(tb_dir if label == 1 else norm_dir, stem, data)
        print(f"  {split_name}: {count_tb(items)} TB + {count_norm(items)} Normal = {len(items)}")

    grand_total = len(train_all) + len(val_all) + len(test_all)
    total_tb = count_tb(train_all) + count_tb(val_all) + count_tb(test_all)
    total_norm = count_norm(train_all) + count_norm(val_all) + count_norm(test_all)
    print(f"\n{'=' * 60}")
    print(f"DONE! {grand_total} images ({total_tb} TB, {total_norm} Normal)")
    print(f"Output: {PROCESSED.resolve()}")
    print(f"Run: python train_ensemble_v2.py")
    print(f"{'=' * 60}")


if __name__ == "__main__":
    main()