#!/usr/bin/env python3
"""
End-to-End Medical Image Dataset Preprocessing Pipeline

This script combines three steps into one:
1. Convert raw BiomedParse format to HF-compatible format
2. Pack dataset and extract template features  
3. Precompute template matches

Usage:
    python end_to_end_preprocessing.py \
        --dataset_name ACDC \
        --template_percent 5 \
        --raw_data_dir /projects/p32685/medical_data/ACDC \
        --output_base_dir /projects/p32685/RL_dataset \
        --device cuda:0

Output:
    /projects/p32685/RL_dataset/acdc/
    /projects/p32685/RL_dataset/acdc_5percent_processed/
        ├── backup_entries/  (with precomputed template matches)
        ├── hf_dataset/
        └── template/
"""

import os
import sys
import json
import argparse
import subprocess
import random
import glob
from datetime import datetime
from pathlib import Path
from typing import List, Dict, Any, Tuple

# Add LISA-main to path
SCRIPT_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(SCRIPT_DIR))

# Add VESSA to path for TemplateManager
# Try local template_manager first, then fallback to VESSA path if exists
VESSA_PATH = "/home/wfm8324/VESSA/train_vessa"
if os.path.exists(VESSA_PATH) and VESSA_PATH not in sys.path:
    sys.path.insert(0, VESSA_PATH)
# template_manager.py is in the same directory, so it will be found automatically

import numpy as np
from PIL import Image

# Ensure logs directory exists
os.makedirs(os.path.join(SCRIPT_DIR, "logs"), exist_ok=True)

# Known dataset configurations
DATASET_TYPES = {
    "ACDC": "ACDC",
    "CAMUS": "CAMUS",
    "MMs": "ACDC",  # Uses ACDC config
    "LGG": "LGG",
    "CXR": "CXR",
    "BreastUS": "BreastUS",
    "BUS": "BUS",
    "kits23": "kits23",
    "abdomenct": "abdomenct",
    "lungx": "lungx",
    "perfusion": "perfusion",
    "amos22_CT": "AMOS22_CT",
    "amos22_MRI": "AMOS22_MRI",
    # MSD Tasks
    "Task01_BrainTumour": "Task01_BrainTumour",
    "Task02_Heart": "Task02_Heart",
    "Task03_Liver": "Task03_Liver",
    "Task04_Hippocampus": "Task04_Hippocampus",
    "Task05_Prostate": "Task05_Prostate",
    "Task06_Lung": "Task06_Lung",
    "Task07_Pancreas": "Task07_Pancreas",
    "Task08_HepaticVessel": "Task08_HepaticVessel",
    "Task09_Spleen": "Task09_Spleen",
    "Task10_Colon": "Task10_Colon",
    # Six new datasets
    "BraTS2020": "BraTS2020",
    "SPIDER": "SPIDER",
    "STS-2D-Tooth": "STS-2D-Tooth",
    "STS-3D-Tooth": "STS-3D-Tooth",
    "TotalSegmentator": "TotalSegmentator",
    "TotalSegmentatorMRI": "TotalSegmentatorMRI",
    "TotalSegmentator_reduced": "TotalSegmentator_reduced",
    "TotalSegmentatorMRI_reduced": "TotalSegmentatorMRI_reduced"
}


def get_output_folder_name(dataset_name):
    """Output folder name in RL_dataset: remove '_reduced' for reduced datasets."""
    if dataset_name == "TotalSegmentator_reduced":
        return "totalsegmentator"
    if dataset_name == "TotalSegmentatorMRI_reduced":
        return "totalsegmentatormri"
    return dataset_name.lower()


def run_step1_convert(dataset_name, dataset_type, raw_data_dir, output_base_dir, output_folder_name=None, split_mode="default"):
    """
    Step 1: Convert raw BiomedParse format to HF-compatible format
    
    For MSD tasks, output structure is:
    output_base_dir/MSD_taskXX/TaskXX_Name/
    """
    print("\n" + "="*80)
    print(f"STEP 1: Converting {dataset_name} to HF format")
    print("="*80)
    
    # Special handling for MSD tasks
    if dataset_name.startswith("Task") and dataset_name.endswith(("BrainTumour", "Heart", "Liver", "Hippocampus", 
                                                                    "Prostate", "Lung", "Pancreas", "HepaticVessel", 
                                                                    "Spleen", "Colon")):
        # Extract task number (e.g., Task01_BrainTumour -> MSD_task01)
        task_num = dataset_name.split("_")[0]  # Task01
        msd_task_dir = os.path.join(output_base_dir, f"MSD_{task_num.lower()}")
        os.makedirs(msd_task_dir, exist_ok=True)
        # Output to MSD_taskXX/TaskXX_Name/
        output_dir = os.path.join(msd_task_dir, dataset_name)
    else:
        # Standard output: output_base_dir/<folder_name> (reduced 数据集去掉 _reduced)
        folder_name = output_folder_name if output_folder_name else get_output_folder_name(dataset_name)
        output_dir = os.path.join(output_base_dir, folder_name)
    
    # Check if already converted (must have dataset_info.json AND actual data files)
    dataset_info_path = os.path.join(output_dir, "dataset_info.json")
    train_images_dir = os.path.join(output_dir, "train", "images")
    has_data = False
    if os.path.exists(train_images_dir):
        image_files = [f for f in os.listdir(train_images_dir) if f.endswith(('.jpg', '.png', '.jpeg'))]
        has_data = len(image_files) > 0
    
    if os.path.exists(output_dir) and os.path.isfile(dataset_info_path) and has_data:
        print(f"✓ Dataset already converted: {output_dir}")
        return output_dir
    if os.path.exists(output_dir):
        if not os.path.isfile(dataset_info_path):
            print(f"⚠ Output dir exists but dataset_info.json missing (incomplete run); removing and re-running conversion.")
        elif not has_data:
            print(f"⚠ Output dir exists but no image files found (incomplete run); removing and re-running conversion.")
        import shutil
        shutil.rmtree(output_dir)
        print(f"[DEBUG] Removed incomplete directory: {output_dir}")
    
    cmd = [
        sys.executable,
        os.path.join(SCRIPT_DIR, "convert_biomedparse_to_hf_format.py"),
        "--input_dir", raw_data_dir,
        "--output_dir", output_dir,
        "--dataset_type", dataset_type,
        "--num_vis", "30",
        "--split_mode", split_mode
    ]
    
    print(f"Running: {' '.join(cmd)}")
    print(f"[DEBUG] Step 1 started at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    
    try:
        result = subprocess.run(cmd, check=True)
        
        if result.returncode == 0:
            print(f"[DEBUG] Step 1 completed successfully at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
            print(f"✓ Step 1 complete: {output_dir}")
            return output_dir
        else:
            print(f"[DEBUG] Step 1 failed with return code: {result.returncode}")
            raise RuntimeError(f"Step 1 failed with code {result.returncode}")
    except subprocess.CalledProcessError as e:
        print(f"[DEBUG] Step 1 subprocess error: {e}")
        print(f"[DEBUG] Step 1 failed at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        raise
    except Exception as e:
        print(f"[DEBUG] Step 1 unexpected error: {e}")
        print(f"[DEBUG] Step 1 failed at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        raise


def run_step2_packing(dataset_dir, output_base_dir, dataset_name, template_percent, processed_prefix=None,
                      template_extraction_device="cuda:0"):
    """
    Step 2: Pack dataset and extract template features
    
    For MSD tasks, output structure is:
    output_base_dir/MSD_taskXX/TaskXX_Name_10percent_processed/
    """
    print("\n" + "="*80)
    print(f"STEP 2: Packing {dataset_name} with {template_percent}% templates")
    print("="*80)
    
    # Special handling for MSD tasks
    if dataset_name.startswith("Task") and dataset_name.endswith(("BrainTumour", "Heart", "Liver", "Hippocampus", 
                                                                    "Prostate", "Lung", "Pancreas", "HepaticVessel", 
                                                                    "Spleen", "Colon")):
        # Extract task number (e.g., Task01_BrainTumour -> MSD_task01)
        task_num = dataset_name.split("_")[0]  # Task01
        msd_task_dir = os.path.join(output_base_dir, f"MSD_{task_num.lower()}")
        os.makedirs(msd_task_dir, exist_ok=True)
        # Output to MSD_taskXX/TaskXX_Name_10percent_processed/
        output_dir = os.path.join(msd_task_dir, f"{dataset_name}_{template_percent}percent_processed")
    else:
        # Standard output: output_base_dir/<folder_name>_Xpercent_processed
        folder_prefix = processed_prefix if processed_prefix else get_output_folder_name(dataset_name)
        output_dir = os.path.join(output_base_dir, f"{folder_prefix}_{template_percent}percent_processed")
    
    # Check if already packed
    if os.path.exists(output_dir):
        print(f"✓ Dataset already packed: {output_dir}")
        return output_dir
    
    print(f"[DEBUG] Step 2 started at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"[DEBUG] Input dataset_dir: {dataset_dir}")
    print(f"[DEBUG] Output dir: {output_dir}")
    
    # First, run packing without templates
    cmd_pack = [
        sys.executable,
        os.path.join(SCRIPT_DIR, "packing_HF_dataset.py"),
        "--dataset_root_dir", dataset_dir,
        "--output_dir", output_dir,
        "--subset_size", "0",
        "--device", "cpu"
    ]
    
    print(f"Running packing: {' '.join(cmd_pack)}")
    print(f"[DEBUG] Starting packing subprocess...")
    
    try:
        result = subprocess.run(cmd_pack, check=True)
        print(f"[DEBUG] Packing subprocess completed successfully (returncode: {result.returncode})")
    except subprocess.CalledProcessError as e:
        print(f"[DEBUG] Packing subprocess failed with return code: {e.returncode}")
        print(f"[DEBUG] Step 2 failed at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        raise
    except Exception as e:
        print(f"[DEBUG] Packing unexpected error: {e}")
        print(f"[DEBUG] Step 2 failed at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        raise
    
    # Then extract templates
    print(f"\nExtracting templates ({template_percent}% of training data)...")
    print(f"[DEBUG] Starting template extraction...")
    
    try:
        extract_templates_python(dataset_dir, output_dir, template_percent, device=template_extraction_device)
        print(f"[DEBUG] Template extraction completed successfully")
    except Exception as e:
        print(f"[DEBUG] Template extraction failed: {e}")
        print(f"[DEBUG] Step 2 failed at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        raise
    
    print(f"[DEBUG] Step 2 completed successfully at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"✓ Step 2 complete: {output_dir}")
    return output_dir


def extract_templates_python(dataset_dir, output_dir, template_percent, device="cuda:0"):
    """
    Extract template features using Python directly (DINOv2 on ``device``).
    """
    from packing_HF_dataset import extract_template_features
    from PIL import Image as PILImage
    
    image_dir = os.path.join(dataset_dir, "train", "images")
    mask_dir = os.path.join(dataset_dir, "train", "masks")
    template_dir = os.path.join(output_dir, "template")
    
    # Get all images
    all_images = sorted(glob.glob(os.path.join(image_dir, "*.jpg")))
    num_templates = int(len(all_images) * (template_percent / 100.0))
    
    # Randomly sample
    random.seed(42)  # Fixed seed for reproducibility
    selected_images = random.sample(all_images, num_templates)
    
    print(f"Total training images: {len(all_images)}")
    print(f"{template_percent}% = {num_templates} templates")
    
    # Load templates
    template_data_dict = {}
    for img_path in selected_images:
        template_id = os.path.splitext(os.path.basename(img_path))[0]
        mask_path = os.path.join(mask_dir, f"{template_id}.png")
        if os.path.exists(mask_path):
            img = PILImage.open(img_path).convert('RGB')
            mask = PILImage.open(mask_path).convert('L')
            # Get modality from dataset_info.json
            dataset_info_path = os.path.join(dataset_dir, "dataset_info.json")
            modality = "MRI"  # Default
            if os.path.exists(dataset_info_path):
                with open(dataset_info_path) as f:
                    info = json.load(f)
                    modality = info.get("modality", "MRI")
            template_data_dict[template_id] = {'image': img, 'mask': mask, 'modality': modality}
    
    # Extract features
    if template_data_dict:
        results = extract_template_features(template_data_dict=template_data_dict, output_dir=template_dir, device=device)
        print(f"✓ Extracted {len(results)} templates")
    else:
        print("✗ No templates loaded")


def get_present_label_ids_from_mask_path(mask_path: str) -> List[int]:
    """Read a label mask and return non-background class ids actually present."""
    try:
        if mask_path and os.path.exists(mask_path):
            arr = np.array(Image.open(mask_path))
            uniq = np.unique(arr)
            present = [int(v) for v in uniq.tolist() if int(v) != 0]
            return present
    except Exception:
        pass
    return []


def get_non_bg_label_ids(labels_json_str: str) -> List[int]:
    """Parse labels JSON and return non-background class IDs."""
    labels = json.loads(labels_json_str)
    return [int(k) for k, name in labels.items() if k != "0" and name != "background"]


def precompute_for_entries(
    entries: List[Dict[str, Any]],
    templates: Dict[str, Dict[str, Any]],
    index,
    template_ids_list: List[str],
    tm,
    topk: int = 1,
    match_strategy: str = "complete",
    max_search: int = 100,
    overwrite: bool = True,
    allow_fallback: bool = False
) -> int:
    """
    Precompute template matches for a list of entries.
    
    Args:
        entries: List of sample entries to process
        templates: Dictionary of template data
        index: FAISS index for similarity search
        template_ids_list: List of template IDs corresponding to index
        tm: TemplateManager instance
        topk: Number of top templates to store per sample
        match_strategy: "complete" (strict - all target classes must be in template) 
                       or "partial" (at least one class overlap)
        max_search: How many nearest neighbors to search before filtering
        overwrite: Whether to overwrite existing template entries
        allow_fallback: If True and not enough strict matches found, fall back to 
                       nearest templates regardless of class overlap (NOT RECOMMENDED)
    
    Returns:
        Number of entries updated
    """
    updated = 0
    for entry in entries:
        image_path = entry.get("image")
        if not image_path:
            continue
        
        # Skip if already has templates and not overwriting
        if (not overwrite) and ("template" in entry) and isinstance(entry["template"], list) and len(entry["template"]) >= topk:
            continue
        
        # Get present class IDs from mask
        labels_json = entry.get("labels", "{}")
        label_mask_path = entry.get("label_mask", None)
        present_ids = set(get_present_label_ids_from_mask_path(label_mask_path))
        if len(present_ids) == 0:
            present_ids = set(get_non_bg_label_ids(labels_json))
        target_label_ids = present_ids
        
        # Extract DINO features
        try:
            feat = tm._extract_feature(image_path)
        except Exception as e:
            print(f"Warning: Failed to extract feature from {image_path}: {e}")
            continue
        
        # FAISS search
        query = feat.reshape(1, -1).astype("float32")
        D, I = index.search(query, min(max_search, len(template_ids_list)))
        
        # Rank templates by similarity with STRICT class matching
        ranked = []
        for d, idx in zip(D[0], I[0]):
            tid = template_ids_list[idx]
            tentry = templates[tid]
            tclasses = set(int(c) for c in tentry.get("classes", []))
            overlap = target_label_ids.intersection(tclasses)
            
            # STRICT: Template must contain ALL target classes
            if len(target_label_ids) == 0:
                continue
            if not target_label_ids.issubset(tclasses):
                continue
            
            ranked.append({
                "template_id": tid,
                "similarity": float(d),
                "overlap": sorted(list(map(int, overlap))),
            })
            if len(ranked) >= topk:
                break
        
        # Optional fallback (disabled by default)
        # Only use if explicitly enabled via allow_fallback=True
        if len(ranked) < topk and allow_fallback:
            print(f"Warning: Only found {len(ranked)}/{topk} strict matches for {image_path}, using fallback")
            for d, idx in zip(D[0], I[0]):
                tid = template_ids_list[idx]
                if any(r["template_id"] == tid for r in ranked):
                    continue
                ranked.append({
                    "template_id": tid,
                    "similarity": float(d),
                    "overlap": [],
                })
                if len(ranked) >= topk:
                    break
        
        # Only update if we found at least one valid match
        if len(ranked) > 0:
            entry["template"] = ranked
            updated += 1
        else:
            # No matching template found - leave entry without template
            if "template" in entry:
                del entry["template"]
    
    return updated


def run_step3_precompute(processed_dir, device="cuda:0", topk=1):
    """
    Step 3: Precompute template matches using STRICT class matching
    
    Template matching strategy:
    - STRICT ONLY: Template must contain ALL classes present in target image
    - No fallback to partial matches
    - Samples without matching templates are left without template field
    """
    print("\n" + "="*80)
    print(f"STEP 3: Precomputing template matches (STRICT)")
    print("="*80)
    print(f"Strategy: Template must contain ALL target classes")
    print(f"Top-k matches per sample: {topk}")
    print(f"[DEBUG] Step 3 started at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"[DEBUG] Processed dir: {processed_dir}")
    
    hf_dataset_path = os.path.join(processed_dir, "hf_dataset")
    template_dir = os.path.join(processed_dir, "template")
    backup_dir = os.path.join(processed_dir, "backup_entries")
    
    print(f"[DEBUG] HF dataset path: {hf_dataset_path}")
    print(f"[DEBUG] Template dir: {template_dir}")
    print(f"[DEBUG] Backup dir: {backup_dir}")
    
    # Check if template matches already computed
    backup_train = os.path.join(backup_dir, "train_entries.json")
    if os.path.exists(backup_train):
        with open(backup_train) as f:
            entries = json.load(f)
            if len(entries) > 0 and "template" in entries[0] and isinstance(entries[0]["template"], list):
                print(f"[DEBUG] Templates already precomputed, skipping Step 3")
                print(f"✓ Templates already precomputed in {backup_train}")
                return
    
    # Verify required directories exist
    print(f"[DEBUG] Verifying required directories...")
    if not os.path.isdir(template_dir):
        print(f"[DEBUG] ERROR: Template directory not found: {template_dir}")
        print(f"✗ Template directory not found: {template_dir}")
        return
    if not os.path.isdir(hf_dataset_path):
        print(f"[DEBUG] ERROR: HF dataset not found: {hf_dataset_path}")
        print(f"✗ HF dataset not found: {hf_dataset_path}")
        return
    if not os.path.isdir(backup_dir):
        print(f"[DEBUG] ERROR: Backup entries not found: {backup_dir}")
        print(f"✗ Backup entries not found: {backup_dir}")
        return
    print(f"[DEBUG] All required directories verified")
    
    # Import dependencies
    print(f"[DEBUG] Importing dependencies...")
    try:
        from datasets import load_from_disk
        from template_manager import TemplateManager
        print(f"[DEBUG] Dependencies imported successfully")
    except ImportError as e:
        print(f"[DEBUG] ERROR: Failed to import modules: {e}")
        print(f"✗ Failed to import required modules: {e}")
        print(f"  Make sure 'datasets' and 'template_manager' are available")
        return
    
    print(f"\nLoading HF dataset from: {hf_dataset_path}")
    print(f"[DEBUG] Loading HF dataset...")
    try:
        ds = load_from_disk(hf_dataset_path)
        train_split = ds.get("train") or ds.get("training")
        val_split = ds.get("val") or ds.get("validation")
        print(f"[DEBUG] HF dataset loaded: train={train_split is not None}, val={val_split is not None}")
    except Exception as e:
        print(f"[DEBUG] ERROR: Failed to load HF dataset: {e}")
        raise
    
    # Load backup entries
    print(f"Loading backup entries from: {backup_dir}")
    print(f"[DEBUG] Loading backup entries...")
    try:
        with open(os.path.join(backup_dir, "train_entries.json"), "r") as f:
            train_entries = json.load(f)
        with open(os.path.join(backup_dir, "val_entries.json"), "r") as f:
            val_entries = json.load(f)
        print(f"[DEBUG] Backup entries loaded: train={len(train_entries)}, val={len(val_entries)}")
    except Exception as e:
        print(f"[DEBUG] ERROR: Failed to load backup entries: {e}")
        raise
    
    # Initialize TemplateManager and load templates
    print(f"Loading templates from: {template_dir}")
    print(f"Initializing TemplateManager on {device}...")
    print(f"[DEBUG] Initializing TemplateManager...")
    try:
        tm = TemplateManager(device=device)
        templates, index, template_ids_list = tm._load_templates(template_dir)
        print(f"[DEBUG] TemplateManager initialized, loaded {len(templates)} templates")
        print(f"✓ Loaded {len(templates)} templates for search")
    except Exception as e:
        print(f"[DEBUG] ERROR: Failed to initialize TemplateManager or load templates: {e}")
        raise
    
    # Precompute matches for train and val with STRICT matching
    total_updates = 0
    
    if train_entries:
        print(f"\nProcessing {len(train_entries)} train entries (strict class matching)...")
        print(f"[DEBUG] Starting train entries precomputation...")
        try:
            updated = precompute_for_entries(
                train_entries, templates, index, template_ids_list, tm,
                topk=topk, 
                match_strategy="complete",  # STRICT: all target classes must be in template
                max_search=100,  # Search more candidates to find strict matches
                overwrite=True,
                allow_fallback=False  # NO fallback to partial matches
            )
            total_updates += updated
            matched_pct = (updated / len(train_entries) * 100) if len(train_entries) > 0 else 0
            print(f"[DEBUG] Train precomputation completed: {updated}/{len(train_entries)} matched")
            print(f"✓ Matched {updated}/{len(train_entries)} train entries ({matched_pct:.1f}%)")
        except Exception as e:
            print(f"[DEBUG] ERROR: Train precomputation failed: {e}")
            raise
    
    if val_entries:
        print(f"\nProcessing {len(val_entries)} val entries (strict class matching)...")
        print(f"[DEBUG] Starting val entries precomputation...")
        try:
            updated = precompute_for_entries(
                val_entries, templates, index, template_ids_list, tm,
                topk=topk,
                match_strategy="complete",  # STRICT: all target classes must be in template
                max_search=100,  # Search more candidates to find strict matches
                overwrite=True,
                allow_fallback=False  # NO fallback to partial matches
            )
            total_updates += updated
            matched_pct = (updated / len(val_entries) * 100) if len(val_entries) > 0 else 0
            print(f"[DEBUG] Val precomputation completed: {updated}/{len(val_entries)} matched")
            print(f"✓ Matched {updated}/{len(val_entries)} val entries ({matched_pct:.1f}%)")
        except Exception as e:
            print(f"[DEBUG] ERROR: Val precomputation failed: {e}")
            raise
    
    # Write back updated JSONs
    print(f"[DEBUG] Writing back updated JSONs...")
    train_json = os.path.join(backup_dir, "train_entries.json")
    val_json = os.path.join(backup_dir, "val_entries.json")
    
    print(f"[DEBUG] Writing train_entries.json...")
    try:
        with open(train_json, "w") as f:
            json.dump(train_entries, f, indent=2)
        print(f"[DEBUG] train_entries.json written successfully")
    except Exception as e:
        print(f"[DEBUG] ERROR: Failed to write train_entries.json: {e}")
        raise
    
    print(f"[DEBUG] Writing val_entries.json...")
    try:
        with open(val_json, "w") as f:
            json.dump(val_entries, f, indent=2)
        print(f"[DEBUG] val_entries.json written successfully")
    except Exception as e:
        print(f"[DEBUG] ERROR: Failed to write val_entries.json: {e}")
        raise
    
    print(f"[DEBUG] Step 3 completed successfully at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"\n✓ Step 3 complete: {total_updates} samples with strict template matches")
    print(f"  Updated: {train_json}")
    print(f"  Updated: {val_json}")
    print(f"\nNote: Samples without matching templates have no 'template' field")


def verify_output(processed_dir, expected_train_count, expected_val_count, expected_template_count):
    """
    Verify the processed dataset has correct structure and counts
    """
    print("\n" + "="*80)
    print("VERIFICATION")
    print("="*80)
    
    checks = []
    
    # Check backup_entries
    train_json = os.path.join(processed_dir, "backup_entries", "train_entries.json")
    val_json = os.path.join(processed_dir, "backup_entries", "val_entries.json")
    
    if os.path.exists(train_json):
        with open(train_json) as f:
            train_entries = json.load(f)
            checks.append(("backup_entries/train", len(train_entries) == expected_train_count, 
                          f"{len(train_entries)} entries (expected {expected_train_count})"))
    else:
        checks.append(("backup_entries/train", False, "File not found"))
    
    if os.path.exists(val_json):
        with open(val_json) as f:
            val_entries = json.load(f)
            checks.append(("backup_entries/val", len(val_entries) == expected_val_count,
                          f"{len(val_entries)} entries (expected {expected_val_count})"))
    else:
        checks.append(("backup_entries/val", False, "File not found"))
    
    # Check hf_dataset
    hf_dir = os.path.join(processed_dir, "hf_dataset")
    if os.path.exists(hf_dir):
        checks.append(("hf_dataset", True, "Directory exists"))
    else:
        checks.append(("hf_dataset", False, "Directory not found"))
    
    # Check templates
    template_dir = os.path.join(processed_dir, "template")
    if os.path.exists(template_dir):
        npz_files = glob.glob(os.path.join(template_dir, "*.npz"))
        checks.append(("template", len(npz_files) == expected_template_count,
                      f"{len(npz_files)} templates (expected {expected_template_count})"))
    else:
        checks.append(("template", False, "Directory not found"))
    
    # Print results
    all_passed = True
    for name, passed, detail in checks:
        status = "✓" if passed else "✗"
        print(f"{status} {name}: {detail}")
        if not passed:
            all_passed = False
    
    return all_passed


def main():
    parser = argparse.ArgumentParser(description="End-to-end preprocessing pipeline")
    parser.add_argument("--dataset_name", type=str, required=True,
                       help="Dataset name (e.g., ACDC, CAMUS, MMs)")
    parser.add_argument("--template_percent", type=int, required=True,
                       help="Percentage of training data to use as templates (e.g., 5, 10)")
    parser.add_argument("--raw_data_dir", type=str, required=True,
                       help="Path to raw dataset (e.g., /projects/p32685/MSD/Task01_BrainTumour)")
    parser.add_argument("--output_base_dir", type=str, default="/projects/p32685/RL_dataset",
                       help="Base output directory")
    parser.add_argument("--device", type=str, default="cuda:0",
                       help="Device for template precomputation")
    parser.add_argument("--skip_convert", action="store_true",
                       help="Skip step 1 (conversion) if already done")
    parser.add_argument("--skip_packing", action="store_true",
                       help="Skip step 2 (packing) if already done")
    parser.add_argument("--skip_precompute", action="store_true",
                       help="Skip step 3 (precompute) if already done")
    parser.add_argument("--output_folder_name", type=str, default=None,
                       help="Optional explicit folder name for converted dataset (Step 1 output folder).")
    parser.add_argument("--processed_prefix", type=str, default=None,
                       help="Optional explicit prefix for processed output folders (Step 2 output prefix).")
    parser.add_argument("--split_mode", type=str, default="default",
                       choices=["default", "acdc_patient_100_50"],
                       help="Split mode passed to converter.")
    
    args = parser.parse_args()
    
    # Determine dataset type
    dataset_type = DATASET_TYPES.get(args.dataset_name, args.dataset_name)
    
    print("╔" + "="*78 + "╗")
    print("║" + " "*20 + "END-TO-END PREPROCESSING PIPELINE" + " "*25 + "║")
    print("╔" + "="*78 + "╗")
    print(f"\nDataset: {args.dataset_name}")
    print(f"Type: {dataset_type}")
    print(f"Template percentage: {args.template_percent}%")
    print(f"Raw data: {args.raw_data_dir}")
    print(f"Output base: {args.output_base_dir}")
    print(f"Device: {args.device}")
    print(f"[DEBUG] Pipeline started at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    
    try:
        # Step 1: Convert
        if not args.skip_convert:
            converted_dir = run_step1_convert(args.dataset_name, dataset_type, 
                                             args.raw_data_dir, args.output_base_dir,
                                             output_folder_name=args.output_folder_name,
                                             split_mode=args.split_mode)
        else:
            # Determine converted_dir based on dataset type
            if args.dataset_name.startswith("Task") and args.dataset_name.endswith(("BrainTumour", "Heart", "Liver", "Hippocampus", 
                                                                                      "Prostate", "Lung", "Pancreas", "HepaticVessel", 
                                                                                      "Spleen", "Colon")):
                task_num = args.dataset_name.split("_")[0]
                converted_dir = os.path.join(args.output_base_dir, f"MSD_{task_num.lower()}", args.dataset_name)
            else:
                folder_name = args.output_folder_name if args.output_folder_name else get_output_folder_name(args.dataset_name)
                converted_dir = os.path.join(args.output_base_dir, folder_name)
            print(f"Skipping Step 1 (conversion), using: {converted_dir}")
        
        # Step 2: Pack and extract templates
        if not args.skip_packing:
            processed_dir = run_step2_packing(converted_dir, args.output_base_dir,
                                             args.dataset_name, args.template_percent,
                                             processed_prefix=args.processed_prefix,
                                             template_extraction_device=args.device)
        else:
            # Determine processed_dir based on dataset type
            if args.dataset_name.startswith("Task") and args.dataset_name.endswith(("BrainTumour", "Heart", "Liver", "Hippocampus", 
                                                                                      "Prostate", "Lung", "Pancreas", "HepaticVessel", 
                                                                                      "Spleen", "Colon")):
                task_num = args.dataset_name.split("_")[0]
                processed_dir = os.path.join(args.output_base_dir, f"MSD_{task_num.lower()}", 
                                           f"{args.dataset_name}_{args.template_percent}percent_processed")
            else:
                folder_prefix = args.processed_prefix if args.processed_prefix else get_output_folder_name(args.dataset_name)
                processed_dir = os.path.join(args.output_base_dir,
                                        f"{folder_prefix}_{args.template_percent}percent_processed")
            print(f"Skipping Step 2 (packing), using: {processed_dir}")
        
        # Step 3: Precompute template matches
        if not args.skip_precompute:
            run_step3_precompute(processed_dir, args.device)
        else:
            print("Skipping Step 3 (precompute)")
        
        # Verification
        print("\n" + "="*80)
        print("PIPELINE COMPLETE")
        print("="*80)
        print(f"[DEBUG] Pipeline completed successfully at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print(f"\nOutput directories created:")
        print(f"  1. Converted: {converted_dir}")
        print(f"  2. Processed: {processed_dir}")
        
        # Calculate expected counts
        train_images = glob.glob(os.path.join(converted_dir, "train", "images", "*"))
        val_images = glob.glob(os.path.join(converted_dir, "val", "images", "*"))
        expected_templates = int(len(train_images) * (args.template_percent / 100.0))
        
        # Verify
        all_good = verify_output(processed_dir, len(train_images), len(val_images), expected_templates)
        
        if all_good:
            print("\n✓ All checks passed! Dataset ready for training.")
        else:
            print("\n✗ Some checks failed. Please review above.")
        
    except Exception as e:
        print(f"\n[DEBUG] Pipeline failed at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print(f"[DEBUG] Error type: {type(e).__name__}")
        print(f"[DEBUG] Error message: {e}")
        print(f"\n✗ Pipeline failed: {e}")
        import traceback
        print(f"[DEBUG] Full traceback:")
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()