"""
Split harmonized feature/mask pairs into train/val/test sets.
Copies matched pairs into output/{train,val,test}/{features,masks}/.

Usage:
    python split_dataset.py
"""

import os
import glob
import yaml
import shutil
import random

def load_config():
    config_path = os.path.join(os.path.dirname(__file__), 'config.yaml')
    with open(config_path, 'r') as f:
        return yaml.safe_load(f)

def main():
    config = load_config()
    
    base_dir = os.path.join(os.path.dirname(__file__), config['paths']['output_dir'])
    feature_dir = os.path.join(base_dir, 'features')
    mask_dir = os.path.join(base_dir, 'masks')
    dataset_dir = os.path.join(base_dir, 'dataset')
    
    train_r = config['splitting']['train_ratio']
    val_r = config['splitting']['val_ratio']
    seed = config['splitting']['random_seed']
    
    # Find all matched pairs
    feature_files = sorted(glob.glob(os.path.join(feature_dir, 'feature_*.tif')))
    
    # Extract IDs and verify masks exist
    paired_ids = []
    for fp in feature_files:
        idx_str = os.path.basename(fp).split('_')[1].split('.')[0]
        mp = os.path.join(mask_dir, f"mask_{idx_str}.tif")
        if os.path.exists(mp):
            paired_ids.append(idx_str)
        else:
            print(f"Warning: No mask for feature_{idx_str}.tif, skipping.")
    
    print(f"Found {len(paired_ids)} matched feature-mask pairs.")
    
    # Shuffle and split
    random.seed(seed)
    random.shuffle(paired_ids)
    
    n = len(paired_ids)
    n_train = int(n * train_r)
    n_val = int(n * (train_r + val_r)) - n_train
    
    train_ids = paired_ids[:n_train]
    val_ids = paired_ids[n_train:n_train + n_val]
    test_ids = paired_ids[n_train + n_val:]
    
    splits = {'train': train_ids, 'val': val_ids, 'test': test_ids}
    
    for split_name, ids in splits.items():
        feat_out = os.path.join(dataset_dir, split_name, 'features')
        mask_out = os.path.join(dataset_dir, split_name, 'masks')
        os.makedirs(feat_out, exist_ok=True)
        os.makedirs(mask_out, exist_ok=True)
        
        for idx_str in ids:
            src_feat = os.path.join(feature_dir, f"feature_{idx_str}.tif")
            src_mask = os.path.join(mask_dir, f"mask_{idx_str}.tif")
            shutil.copy2(src_feat, feat_out)
            shutil.copy2(src_mask, mask_out)
        
        print(f"  {split_name}: {len(ids)} pairs")
    
    print(f"\nSplit complete! Dataset saved to {dataset_dir}")
    print(f"  Train: {len(train_ids)} | Val: {len(val_ids)} | Test: {len(test_ids)}")
    
    # Leak check
    tr, va, te = set(train_ids), set(val_ids), set(test_ids)
    assert len(tr & va) == 0, "LEAK: train ∩ val is not empty!"
    assert len(tr & te) == 0, "LEAK: train ∩ test is not empty!"
    assert len(va & te) == 0, "LEAK: val ∩ test is not empty!"
    print("Leak check PASSED — no overlap between splits.")

if __name__ == '__main__':
    main()