import os
from tqdm import tqdm
import numpy as np
import shutil

BASE_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), 'datasets')

ISRUC_SG1_DIR = os.path.join(BASE_PATH, 'isruc_sg1')
ISRUC_SG2_DIR = os.path.join(BASE_PATH, 'isruc_sg2')
ISRUC_SG3_DIR = os.path.join(BASE_PATH, 'isruc_sg3')

ISRUC_GATHER_DIR = os.path.join(BASE_PATH, 'isruc_all')

os.makedirs(ISRUC_GATHER_DIR, exist_ok=True)
os.makedirs(os.path.join(ISRUC_GATHER_DIR, 'annotations'), exist_ok=True)


def gather_isruc(prefix='sg1'):
    source_dir = ISRUC_SG1_DIR if prefix == 'sg1' else ISRUC_SG2_DIR if prefix == 'sg2' else ISRUC_SG3_DIR
    for ann_file in os.listdir(os.path.join(source_dir, 'annotations')):
        new_name = f"{prefix}_{ann_file}"
        if not os.path.exists(os.path.join(ISRUC_GATHER_DIR, 'annotations', new_name)):
            os.symlink(os.path.join(source_dir, 'annotations', ann_file), 
                        os.path.join(ISRUC_GATHER_DIR, 'annotations', new_name))

    source_dir = os.path.join(source_dir, 'signals_MAE')

    for chn_grp in os.listdir(source_dir):
        for chn in os.listdir(os.path.join(source_dir, chn_grp)):
            os.makedirs(os.path.join(ISRUC_GATHER_DIR, 'signals_MAE', chn_grp), exist_ok=True)
            os.makedirs(os.path.join(ISRUC_GATHER_DIR, 'signals_MAE', chn_grp), exist_ok=True)
            os.makedirs(os.path.join(ISRUC_GATHER_DIR, 'signals_MAE', chn_grp, chn), exist_ok=True)
            for subj_file in os.listdir(os.path.join(source_dir, chn_grp, chn)):
                new_name = f"{prefix}_{subj_file}"
                if not os.path.exists(os.path.join(ISRUC_GATHER_DIR, 'signals_MAE', chn_grp, chn, new_name)):
                    os.symlink(os.path.join(source_dir, chn_grp, chn, subj_file), 
                                os.path.join(ISRUC_GATHER_DIR, 'signals_MAE', chn_grp, chn, new_name))
                
    
        


if __name__ == "__main__":
    print("Gathering ISRUC SG1 datasets into a single directory...")
    gather_isruc('sg1')
    print("Gathering ISRUC SG2 datasets into a single directory...")
    gather_isruc('sg2')
    print("Gathering ISRUC SG3 datasets into a single directory...")
    gather_isruc('sg3')