import os
import pandas as pd
import numpy as np
from tqdm import tqdm
from suno_utils.utils.text import write_jsonl
from joblib import Parallel, delayed
import multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
import time

version = "t13"
output_dir = "/home/christian/code/christian/metadata/reward_model"

dataframes = [
    "/home/tony/Data/Preference/up_v2_d3/interesting_clips_ahi_d3_20250608.pkl",
    "/home/tony/Data/Preference/up_v2_d4/interesting_clips_ahi_d4_20250714.pkl",
    # "/home/tony/Data/Preference/up_v2_d5/interesting_clips_ahi_d5_20250824.pkl"
    "/home/tony/Data/Preference/up_v2_d5/fully_merged_up_v2_d5.pkl",
]

root_dirs = [
    "/app2/suno/data/dpo/diff2_v2_d3",
    "/app2/suno/data/dpo/diff2_v2_d4",
    "/app2/suno/data/dpo/diff2_v2_d5",
]


def process_index(index, df, root_dir):
    pos_item_id = df.iloc[index]["id"]
    neg_item_id = df.iloc[index - 1]["id"]

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

    if os.path.exists(pos_vae_latents_filepath) and os.path.exists(
        neg_vae_latents_filepath
    ):
        return {
            "id": pos_item_id,
            "pos_vae_latents_filepath": pos_vae_latents_filepath,
            "neg_vae_latents_filepath": neg_vae_latents_filepath,
        }
    else:
        return None


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

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

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


def process_with_threading(indices, df, root_dir, n_jobs=None):
    """Use threading for I/O bound tasks"""
    if n_jobs is None:
        n_jobs = min(32, (os.cpu_count() or 1) + 4)  # More threads for I/O bound

    # Create batches to reduce overhead
    batch_size = max(1, len(indices) // (n_jobs * 4))
    index_batches = [
        indices[i : i + batch_size] for i in range(0, len(indices), batch_size)
    ]

    print(
        f"Using {n_jobs} threads with {len(index_batches)} batches of size {batch_size}"
    )

    with ThreadPoolExecutor(max_workers=n_jobs) as executor:
        futures = []
        for batch in index_batches:
            # Pass the full dataframe - threading shares memory so this is efficient
            future = executor.submit(process_index_batch, batch, df, root_dir)
            futures.append(future)

        results = []
        for future in tqdm(futures, desc="Processing batches"):
            batch_results = future.result()
            results.extend(batch_results)

    return results


def process_with_multiprocessing(indices, df, root_dir, n_jobs=None):
    """Use multiprocessing with better chunking"""
    if n_jobs is None:
        n_jobs = mp.cpu_count()

    # Create larger batches to reduce serialization overhead
    batch_size = max(10, len(indices) // (n_jobs * 2))
    index_batches = [
        indices[i : i + batch_size] for i in range(0, len(indices), batch_size)
    ]

    print(
        f"Using {n_jobs} processes with {len(index_batches)} batches of size {batch_size}"
    )

    with ProcessPoolExecutor(max_workers=n_jobs) as executor:
        futures = []
        for batch in index_batches:
            # Pass the full dataframe to each process
            future = executor.submit(process_index_batch, batch, df, root_dir)
            futures.append(future)

        results = []
        for future in tqdm(futures, desc="Processing batches"):
            batch_results = future.result()
            results.extend(batch_results)

    return results


def process_with_joblib_improved(
    indices, df, root_dir, n_jobs=None, backend="threading"
):
    """Improved joblib implementation with better chunking"""
    if n_jobs is None:
        n_jobs = -1

    # Create batches to reduce overhead
    batch_size = max(5, len(indices) // (abs(n_jobs) * 3))
    index_batches = [
        indices[i : i + batch_size] for i in range(0, len(indices), batch_size)
    ]

    print(
        f"Using joblib with {abs(n_jobs)} workers, {len(index_batches)} batches of size {batch_size}, backend={backend}"
    )

    def process_batch_wrapper(batch):
        return process_index_batch(batch, df, root_dir)

    results = Parallel(n_jobs=n_jobs, backend=backend, batch_size=1)(
        delayed(process_batch_wrapper)(batch)
        for batch in tqdm(index_batches, desc="Processing batches")
    )

    # Flatten results
    flattened_results = []
    for batch_results in results:
        flattened_results.extend(batch_results)

    return flattened_results


if __name__ == "__main__":
    # Configuration - you can change these to test different approaches
    PARALLEL_METHOD = "threading"  # Options: "threading", "multiprocessing", "joblib_threading", "joblib_loky"
    N_JOBS = None  # None for auto-detection, or specify a number

    print(f"CPU count: {mp.cpu_count()}")
    print(f"Using parallel method: {PARALLEL_METHOD}")

    train_metas = []
    val_metas = []

    for root_dir, df_path in zip(root_dirs, dataframes):
        print(f"\nProcessing {df_path}")
        df = pd.read_pickle(df_path)
        print(f"DataFrame size: {len(df)}")

        indices = np.arange(len(df))
        indices = indices[indices % 2 == 1]
        print(f"Processing {len(indices)} indices")

        start_time = time.time()

        # Choose parallel processing method
        if PARALLEL_METHOD == "threading":
            subset_metas = process_with_threading(indices, df, root_dir, N_JOBS)
        elif PARALLEL_METHOD == "multiprocessing":
            subset_metas = process_with_multiprocessing(indices, df, root_dir, N_JOBS)
        elif PARALLEL_METHOD == "joblib_threading":
            subset_metas = process_with_joblib_improved(
                indices, df, root_dir, N_JOBS, "threading"
            )
        elif PARALLEL_METHOD == "joblib_loky":
            subset_metas = process_with_joblib_improved(
                indices, df, root_dir, N_JOBS, "loky"
            )
        else:
            # Fallback to original method
            print("Using original joblib method")
            results = Parallel(n_jobs=-1, backend="loky")(
                delayed(process_index)(index, df, root_dir) for index in tqdm(indices)
            )
            subset_metas = [meta for meta in results if meta is not None]

        end_time = time.time()
        print(
            f"Processed {len(subset_metas)} valid items in {end_time - start_time:.2f} seconds"
        )
        print(f"Rate: {len(subset_metas) / (end_time - start_time):.2f} items/second")

        train_metas.extend(subset_metas[: int(len(subset_metas) * 0.95)])
        val_metas.extend(subset_metas[int(len(subset_metas) * 0.95) :])

    print("\nFinal results:")
    print(f"Train metas: {len(train_metas)}")
    print(f"Val metas: {len(val_metas)}")

    write_jsonl(train_metas, output_dir + f"/metas_tr_{version}.jsonl")
    write_jsonl(val_metas, output_dir + f"/metas_val_{version}.jsonl")