from pathlib import Path
from concurrent.futures import ProcessPoolExecutor, as_completed
import os

import pandas as pd
import soundfile as sf
from tqdm import tqdm


# ============================================================
# Paths
# ============================================================

AUDIO_ROOT = Path("/home/debarpanb1/TREA_2.0/GISE/isolated_events")
CSV_PATH = Path("/home/debarpanb1/TREA_2.0/GISE/meta/gise_metadata.csv")

OUTPUT_PATH = Path("/home/debarpanb1/TREA_2.0/GISE/meta/gise_audio_format_check.csv")

NUM_WORKERS = min(16, os.cpu_count() or 4)


# ============================================================
# Worker
# ============================================================

def check_one(row_dict):
    relative_path = Path(row_dict["relative_path"])
    audio_path = AUDIO_ROOT / relative_path

    row_dict["audio_path"] = str(audio_path)
    row_dict["exists"] = audio_path.exists()
    row_dict["read_success"] = False
    row_dict["read_error"] = ""

    if not audio_path.exists():
        row_dict["read_error"] = "missing_audio_file"
        return row_dict

    try:
        info = sf.info(str(audio_path))

        row_dict["sample_rate"] = info.samplerate
        row_dict["channels"] = info.channels
        row_dict["frames"] = info.frames
        row_dict["duration"] = info.frames / info.samplerate
        row_dict["format"] = info.format
        row_dict["subtype"] = info.subtype
        row_dict["endian"] = info.endian
        row_dict["sections"] = info.sections
        row_dict["read_success"] = True

    except Exception as e:
        row_dict["read_error"] = repr(e)

    return row_dict


# ============================================================
# Main
# ============================================================

def main():
    df = pd.read_csv(CSV_PATH)
    records = df.to_dict("records")

    print(f"Checking {len(records)} GISE files with {NUM_WORKERS} workers...")

    results = []

    with ProcessPoolExecutor(max_workers=NUM_WORKERS) as executor:
        futures = [executor.submit(check_one, row) for row in records]

        for future in tqdm(as_completed(futures), total=len(futures)):
            results.append(future.result())

    out_df = pd.DataFrame(results)
    out_df.to_csv(OUTPUT_PATH, index=False)

    print("\nDone.")
    print(f"Saved detailed check to: {OUTPUT_PATH}")

    print("\nExists:")
    print(out_df["exists"].value_counts(dropna=False))

    print("\nRead success:")
    print(out_df["read_success"].value_counts(dropna=False))

    print("\nSample rates:")
    print(out_df["sample_rate"].value_counts(dropna=False).sort_index())

    print("\nChannels:")
    print(out_df["channels"].value_counts(dropna=False).sort_index())

    print("\nFormats:")
    print(out_df["format"].value_counts(dropna=False))

    print("\nSubtypes:")
    print(out_df["subtype"].value_counts(dropna=False))

    print("\nDuration summary:")
    print(out_df["duration"].describe())

    print("\nDuration filter counts:")
    print("duration < 0.5s:", (out_df["duration"] < 0.5).sum())
    print("0.5s <= duration <= 10s:", ((out_df["duration"] >= 0.5) & (out_df["duration"] <= 10.0)).sum())
    print("duration > 10s:", (out_df["duration"] > 10.0).sum())

    print("\nNon-22050 files:")
    non_22050 = out_df[out_df["sample_rate"] != 22050]
    print(len(non_22050))
    if len(non_22050) > 0:
        print(non_22050[["relative_path", "sample_rate", "channels", "format", "subtype", "duration"]].head(20))

    print("\nNon-PCM_16 files:")
    non_pcm16 = out_df[out_df["subtype"] != "PCM_16"]
    print(len(non_pcm16))
    if len(non_pcm16) > 0:
        print(non_pcm16[["relative_path", "sample_rate", "channels", "format", "subtype", "duration"]].head(20))

    print("\nNon-mono files:")
    non_mono = out_df[out_df["channels"] != 1]
    print(len(non_mono))
    if len(non_mono) > 0:
        print(non_mono[["relative_path", "sample_rate", "channels", "format", "subtype", "duration"]].head(20))


if __name__ == "__main__":
    main()