#!/usr/bin/env python3
"""
Test script to compare different parallel processing methods
Run this to see which method works best for your use case
"""

import os
import time
import tempfile
import pandas as pd
import numpy as np
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
from joblib import Parallel, delayed
import multiprocessing as mp


def create_test_data(n_files=1000):
    """Create temporary test files and dataframe"""
    temp_dir = tempfile.mkdtemp()
    df_data = []

    for i in range(n_files):
        file_id = f"test_{i:06d}"
        # Create some files, skip others to simulate missing files
        if i % 3 != 0:  # Create 2/3 of files
            filepath = os.path.join(temp_dir, f"{file_id}_vae.npz")
            with open(filepath, "w") as f:
                f.write("dummy")
        df_data.append({"id": file_id})

    df = pd.DataFrame(df_data)
    return temp_dir, df


def process_index_simple(index, df, root_dir):
    """Simple file existence check"""
    pos_item_id = df.iloc[index]["id"]
    neg_item_id = df.iloc[index - 1]["id"]

    pos_file = os.path.join(root_dir, f"{pos_item_id}_vae.npz")
    neg_file = os.path.join(root_dir, f"{neg_item_id}_vae.npz")

    if os.path.exists(pos_file) and os.path.exists(neg_file):
        return {"id": pos_item_id, "pos_file": pos_file, "neg_file": neg_file}
    return None


def process_index_batch_simple(index_batch, df, root_dir):
    """Process a batch of indices"""
    results = []
    for index in index_batch:
        pos_item_id = df.iloc[index]["id"]
        neg_item_id = df.iloc[index - 1]["id"]

        pos_file = os.path.join(root_dir, f"{pos_item_id}_vae.npz")
        neg_file = os.path.join(root_dir, f"{neg_item_id}_vae.npz")

        if os.path.exists(pos_file) and os.path.exists(neg_file):
            results.append(
                {"id": pos_item_id, "pos_file": pos_file, "neg_file": neg_file}
            )
    return results


def test_original_method(indices, df, root_dir):
    """Original joblib method"""
    print("Testing original joblib method...")
    start = time.time()
    results = Parallel(n_jobs=-1, backend="loky")(
        delayed(process_index_simple)(index, df, root_dir) for index in indices
    )
    valid_results = [r for r in results if r is not None]
    end = time.time()
    print(f"Original method: {len(valid_results)} results in {end-start:.2f}s")
    return valid_results, end - start


def test_threading_method(indices, df, root_dir, n_jobs=16):
    """Threading method"""
    print(f"Testing threading method with {n_jobs} threads...")
    start = time.time()

    with ThreadPoolExecutor(max_workers=n_jobs) as executor:
        futures = [
            executor.submit(process_index_simple, index, df, root_dir)
            for index in indices
        ]
        results = [future.result() for future in futures]

    valid_results = [r for r in results if r is not None]
    end = time.time()
    print(f"Threading method: {len(valid_results)} results in {end-start:.2f}s")
    return valid_results, end - start


def test_multiprocessing_method(indices, df, root_dir, n_jobs=None):
    """Multiprocessing method"""
    if n_jobs is None:
        n_jobs = mp.cpu_count()
    print(f"Testing multiprocessing method with {n_jobs} processes...")
    start = time.time()

    with ProcessPoolExecutor(max_workers=n_jobs) as executor:
        futures = [
            executor.submit(process_index_simple, index, df, root_dir)
            for index in indices
        ]
        results = [future.result() for future in futures]

    valid_results = [r for r in results if r is not None]
    end = time.time()
    print(f"Multiprocessing method: {len(valid_results)} results in {end-start:.2f}s")
    return valid_results, end - start


def test_joblib_threading(indices, df, root_dir):
    """Joblib with threading backend"""
    print("Testing joblib with threading backend...")
    start = time.time()
    results = Parallel(n_jobs=-1, backend="threading")(
        delayed(process_index_simple)(index, df, root_dir) for index in indices
    )
    valid_results = [r for r in results if r is not None]
    end = time.time()
    print(f"Joblib threading: {len(valid_results)} results in {end-start:.2f}s")
    return valid_results, end - start


if __name__ == "__main__":
    print(f"CPU count: {mp.cpu_count()}")
    print("Creating test data...")

    # Create test data
    temp_dir, df = create_test_data(1000)
    indices = np.arange(1, len(df))  # Start from 1 since we need index-1

    print(f"Testing with {len(indices)} indices")
    print("=" * 50)

    # Test different methods
    methods = [
        ("Original (joblib loky)", lambda: test_original_method(indices, df, temp_dir)),
        (
            "Threading (16 threads)",
            lambda: test_threading_method(indices, df, temp_dir, 16),
        ),
        (
            "Threading (32 threads)",
            lambda: test_threading_method(indices, df, temp_dir, 32),
        ),
        ("Multiprocessing", lambda: test_multiprocessing_method(indices, df, temp_dir)),
        ("Joblib threading", lambda: test_joblib_threading(indices, df, temp_dir)),
    ]

    results = {}
    for name, test_func in methods:
        try:
            valid_results, duration = test_func()
            results[name] = duration
            print(f"✓ {name}: {duration:.2f}s")
        except Exception as e:
            print(f"✗ {name}: Failed - {e}")
        print("-" * 30)

    # Cleanup
    import shutil

    shutil.rmtree(temp_dir)

    print("\nPerformance Summary:")
    print("=" * 50)
    sorted_results = sorted(results.items(), key=lambda x: x[1])
    for i, (name, duration) in enumerate(sorted_results, 1):
        print(f"{i}. {name}: {duration:.2f}s")

    if sorted_results:
        fastest = sorted_results[0]
        print(f"\n🏆 Fastest method: {fastest[0]} ({fastest[1]:.2f}s)")
        print(f"\n💡 Recommendation: Use '{fastest[0]}' for your actual script")