import h5py
import json
import numpy as np
from pathlib import Path
from sklearn.model_selection import StratifiedKFold, train_test_split

def generate_multitask_metadata(h5_input_path, json_output_path):
    input_file = Path(h5_input_path)
    output_file = Path(json_output_path)

    if not input_file.exists():
        raise FileNotFoundError(f"Source data not found at: {input_file.resolve()}")

    with h5py.File(input_file, 'r') as f:
        emotion_labels = np.array(f['label_emotion'])
        sign_labels = np.array(f['label_sign'])
        # Video IDs correspond to the index in the label datasets
        video_ids = np.arange(len(sign_labels))

    composite_labels = [f"{s}_{e}" for s, e in zip(sign_labels, emotion_labels)]

    ids_train_val, ids_test, labels_train_val, labels_test = train_test_split(
        video_ids, 
        composite_labels, 
        test_size=0.20, 
        random_state=42, 
        stratify=composite_labels
    )

    skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
    folds_data = []

    for train_idx, val_idx in skf.split(ids_train_val, labels_train_val):
        folds_data.append({
            "train": [
                {"id": int(i), "sign": int(sign_labels[i]), "emotion": int(emotion_labels[i])} 
                for i in ids_train_val[train_idx]
            ],
            "val": [
                {"id": int(i), "sign": int(sign_labels[i]), "emotion": int(emotion_labels[i])} 
                for i in ids_train_val[val_idx]
            ]
        })

    metadata = {
        "test_set": [
            {"id": int(i), "sign": int(sign_labels[i]), "emotion": int(emotion_labels[i])} 
            for i in ids_test
        ],
        "folds": folds_data
    }

    output_file.parent.mkdir(parents=True, exist_ok=True)

    with open(output_file, 'w') as jf:
        json.dump(metadata, jf, indent=4)
    
    print(f"Split metadata successfully written to: {output_file.resolve()}")

if __name__ == "__main__":
    generate_multitask_metadata(
        h5_input_path="fsl-data/multitask_mediapipe.h5", 
        json_output_path="metadata/multitask_splits.json"
    )