"""Data Splitting and Metadata Management with Data Leakage Protection.

Enforces:
1. Exact Duplicate Quarantine: Identical audio clips (all 22 duplicate clusters)
   are strictly confined to the TRAIN set to prevent any data leakage into Validation or Test.
2. Stratified 70/15/15 Splitting: Ensures balanced class distribution across Train, Validation, and Test.
3. Cryptographic Hash Lock: Generates splits_v1.json with a SHA-256 hash for strict reproducibility.
4. Generates dataset/metadata/metadata.csv with all established fields.
"""

import os
import json
import hashlib
from pathlib import Path
import pandas as pd
import numpy as np

from src.utils.config import (
    RAW_DATA_DIR,
    PROCESSED_DATA_DIR,
    METADATA_DIR,
    RANDOM_SEED,
    EMOTION_CLASSES,
    PROJECT_ROOT
)


def create_splits_and_metadata():
    print("=" * 70)
    print("STARTING DATA SPLITTING & LEAKAGE PROTECTION")
    print(f"Random Seed: {RANDOM_SEED}")
    print("=" * 70)

    # 1. Load the comprehensive dataset audit CSV
    audit_csv_path = METADATA_DIR / "dataset_audit.csv"
    if not audit_csv_path.exists():
        audit_csv_path = PROJECT_ROOT / "dataset_audit.csv"
    
    df = pd.read_csv(audit_csv_path)
    print(f"Loaded audit records for {len(df)} files.")

    # 2. Add processed file path and verify existence
    df["processed_file_path"] = df.apply(
        lambda row: str(PROCESSED_DATA_DIR / row["emotion"] / row["filename"]),
        axis=1
    )
    
    missing_processed = [p for p in df["processed_file_path"] if not os.path.exists(p)]
    if missing_processed:
        raise FileNotFoundError(f"Missing {len(missing_processed)} processed files! Run standardization first.")
    print("Verified all 5,000 processed files exist in dataset/processed/.")

    # 3. Identify duplicate audio files that MUST be quarantined to TRAIN
    # Any clip that belongs to a duplicate cluster goes to Train to prevent test leakage
    duplicate_mask = df["is_duplicate"] == True
    duplicate_df = df[duplicate_mask]
    non_duplicate_df = df[~duplicate_mask]
    
    print(f"Exact duplicate files to quarantine to TRAIN: {len(duplicate_df)} files ({df['duplicate_cluster_id'].nunique() - 1} clusters)")
    print(f"Unique files available for stratified partitioning: {len(non_duplicate_df)} files")

    # 4. Stratified Split on non-duplicate pool:
    # Target overall: Train ~70% (3500), Val ~15% (750), Test ~15% (750)
    # Since 44 duplicate files go into Train, remaining quota for Train is 3500 - 44 = 3456.
    # Total unique files = 4956.
    # Val = 750 (approx 150 per class)
    # Test = 750 (approx 150 per class)
    # Train unique = 4956 - 1500 = 3456
    
    np.random.seed(RANDOM_SEED)
    
    train_indices = list(duplicate_df.index)
    val_indices = []
    test_indices = []

    for emotion in EMOTION_CLASSES:
        emo_indices = non_duplicate_df[non_duplicate_df["emotion"] == emotion].index.tolist()
        np.random.shuffle(emo_indices)
        
        n_val = 150
        n_test = 150
        
        emo_val = emo_indices[:n_val]
        emo_test = emo_indices[n_val:n_val + n_test]
        emo_train = emo_indices[n_val + n_test:]
        
        val_indices.extend(emo_val)
        test_indices.extend(emo_test)
        train_indices.extend(emo_train)

    train_indices.sort()
    val_indices.sort()
    test_indices.sort()

    df.loc[train_indices, "split"] = "train"
    df.loc[val_indices, "split"] = "val"
    df.loc[test_indices, "split"] = "test"

    # Add source, group_id, speaker_id explicitly documenting limitation
    # As mandated by prompt: "Only include fields that can actually be established. Do not invent speaker IDs or source IDs."
    df["source"] = "BASE_multimedia_corpus"
    df["group_id"] = df["duplicate_cluster_id"].apply(lambda x: x if pd.notna(x) and x != "" else "independent")
    df["speaker_id"] = "not_annotated"

    # Verify duplicate leakage
    val_hashes = set(df.loc[val_indices, "md5_hash"])
    test_hashes = set(df.loc[test_indices, "md5_hash"])
    train_hashes = set(df.loc[train_indices, "md5_hash"])

    train_val_overlap = train_hashes.intersection(val_hashes)
    train_test_overlap = train_hashes.intersection(test_hashes)
    val_test_overlap = val_hashes.intersection(test_hashes)

    print("\n--- DATA LEAKAGE VERIFICATION ---")
    print(f"Identical Hash Overlap Train-Test: {len(train_test_overlap)} (Must be 0)")
    print(f"Identical Hash Overlap Train-Val:  {len(train_val_overlap)} (Must be 0)")
    print(f"Identical Hash Overlap Val-Test:    {len(val_test_overlap)} (Must be 0)")
    assert len(train_test_overlap) == 0, "CRITICAL ERROR: Data leakage detected between Train and Test!"
    assert len(train_val_overlap) == 0, "CRITICAL ERROR: Data leakage detected between Train and Val!"
    assert len(val_test_overlap) == 0, "CRITICAL ERROR: Data leakage detected between Val and Test!"
    print("PASSED: Zero audio duplicate leakage across splits!")

    # Print split summary
    print("\n--- SPLIT DISTRIBUTION ---")
    split_summary = df.groupby(["split", "emotion"]).size().unstack(fill_value=0)
    split_summary["Total"] = split_summary.sum(axis=1)
    print(split_summary)

    # 5. Save metadata.csv
    out_cols = [
        "audio_id",
        "filename",
        "file_path",
        "processed_file_path",
        "emotion",
        "duration_sec",
        "sample_rate",
        "channels",
        "bit_depth",
        "silence_percentage",
        "source",
        "group_id",
        "speaker_id",
        "quality_status",
        "is_duplicate",
        "split"
    ]
    metadata_csv_path = METADATA_DIR / "metadata.csv"
    metadata_csv_root = PROJECT_ROOT / "metadata.csv"
    df[out_cols].to_csv(metadata_csv_path, index=False)
    df[out_cols].to_csv(metadata_csv_root, index=False)
    print(f"\nSaved metadata.csv to {metadata_csv_path}")

    # 6. Save splits_v1.json
    splits_dict = {
        "split_version": "v1.0",
        "random_seed": RANDOM_SEED,
        "strategy": "Stratified 70/15/15 with Exact Duplicate Quarantine to Train",
        "split_counts": {
            "train": len(train_indices),
            "val": len(val_indices),
            "test": len(test_indices),
            "total": len(df)
        },
        "class_breakdown": {
            split: df[df["split"] == split]["emotion"].value_counts().to_dict()
            for split in ["train", "val", "test"]
        },
        "train_audio_ids": df.loc[train_indices, "audio_id"].tolist(),
        "val_audio_ids": df.loc[val_indices, "audio_id"].tolist(),
        "test_audio_ids": df.loc[test_indices, "audio_id"].tolist()
    }

    splits_json_str = json.dumps(splits_dict, indent=2)
    splits_sha256 = hashlib.sha256(splits_json_str.encode("utf-8")).hexdigest()
    splits_dict["sha256_checksum"] = splits_sha256

    splits_json_path = METADATA_DIR / "splits_v1.json"
    splits_json_root = PROJECT_ROOT / "splits_v1.json"
    with open(splits_json_path, "w", encoding="utf-8") as f:
        json.dump(splits_dict, f, indent=2)
    with open(splits_json_root, "w", encoding="utf-8") as f:
        json.dump(splits_dict, f, indent=2)

    print(f"Saved splits_v1.json to {splits_json_path} (SHA-256: {splits_sha256[:16]}...)")
    print("=" * 70)


if __name__ == "__main__":
    create_splits_and_metadata()