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

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_ROOT = Path("/home/debarpanb1/TREA_2.0/GISE_preprocessed_duration")
OUTPUT_AUDIO_ROOT = OUTPUT_ROOT / "isolated_events"

FULL_METADATA_OUT = OUTPUT_ROOT / "gise_metadata_with_duration_filter.csv"
KEPT_METADATA_OUT = OUTPUT_ROOT / "gise_metadata_kept_0p5_10s.csv"
DROPPED_METADATA_OUT = OUTPUT_ROOT / "gise_metadata_dropped_duration.csv"

OUTPUT_AUDIO_ROOT.mkdir(parents=True, exist_ok=True)


# ============================================================
# Config
# ============================================================

MIN_DURATION = 0.5
MAX_DURATION = 10.0

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


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

def process_one(row_dict):
    rel_path = Path(row_dict["relative_path"])

    src_path = AUDIO_ROOT / rel_path
    dst_path = OUTPUT_AUDIO_ROOT / rel_path

    row_dict["original_audio_path"] = str(src_path)
    row_dict["clean_audio_path"] = str(dst_path)

    row_dict["keep"] = False
    row_dict["drop_reason"] = ""
    row_dict["copy_success"] = False
    row_dict["error"] = ""

    if not src_path.exists():
        row_dict["drop_reason"] = "missing_audio_file"
        row_dict["error"] = "missing_audio_file"
        return row_dict

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

        duration = info.frames / info.samplerate

        row_dict["sample_rate"] = info.samplerate
        row_dict["channels"] = info.channels
        row_dict["frames"] = info.frames
        row_dict["duration"] = duration
        row_dict["format"] = info.format
        row_dict["subtype"] = info.subtype

        if duration < MIN_DURATION:
            row_dict["drop_reason"] = "duration_lt_0.5s"
            return row_dict

        if duration > MAX_DURATION:
            row_dict["drop_reason"] = "duration_gt_10s"
            return row_dict

        # Keep and copy
        row_dict["keep"] = True
        row_dict["drop_reason"] = ""

        dst_path.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(src_path, dst_path)

        row_dict["copy_success"] = True

    except Exception as e:
        row_dict["drop_reason"] = "processing_error"
        row_dict["error"] = repr(e)

    return row_dict


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

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

    print(f"Input metadata: {CSV_PATH}")
    print(f"Input audio root: {AUDIO_ROOT}")
    print(f"Output root: {OUTPUT_ROOT}")
    print(f"Filtering: {MIN_DURATION} <= duration <= {MAX_DURATION}")
    print(f"Workers: {NUM_WORKERS}")

    results = []

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

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

    out_df = pd.DataFrame(results)

    kept_df = out_df[out_df["keep"] == True].copy()
    dropped_df = out_df[out_df["keep"] == False].copy()

    OUTPUT_ROOT.mkdir(parents=True, exist_ok=True)

    out_df.to_csv(FULL_METADATA_OUT, index=False)
    kept_df.to_csv(KEPT_METADATA_OUT, index=False)
    dropped_df.to_csv(DROPPED_METADATA_OUT, index=False)

    print("\nDone.")
    print(f"Full metadata: {FULL_METADATA_OUT}")
    print(f"Kept metadata: {KEPT_METADATA_OUT}")
    print(f"Dropped metadata: {DROPPED_METADATA_OUT}")
    print(f"Copied kept audio: {OUTPUT_AUDIO_ROOT}")

    print("\nCounts:")
    print(out_df["keep"].value_counts(dropna=False))

    print("\nDrop reasons:")
    print(out_df["drop_reason"].value_counts(dropna=False))

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

    print("\nClass counts kept, top 20:")
    print(kept_df["class"].value_counts().head(20))

    print("\nSplit counts kept:")
    print(kept_df["split"].value_counts())


if __name__ == "__main__":
    main()