# TRANSFORME LES  FICHIERS .TXT EN MATRICE 8x512, CE SCRIPT EST UTILISE DANS TRAINDATA.PY

import torch
from torch.utils.data import DataLoader, TensorDataset
import numpy as np
import os

train_batch_size = 16
data_dir = "C:/DATA/M1/Stages/Fablab/Dataset"

LST2 = [] 
label = []

for ID, sub_dir in enumerate(os.listdir(data_dir)):
    for file in os.listdir(os.path.join(data_dir, sub_dir)):
        label.append(ID)
        
        f = open(os.path.join(data_dir, sub_dir, file), 'r+')
        lst_f = []

        for i in f:
            i = i.replace('\n', '')
            lst_f.append(int(i))
            
        # matrice = np.array(lst_f).reshape(8, 512)
        # LST_2D.append(matrice)

        matrice = np.array(lst_f[:4096]).reshape(8, 512)
        LST2.append(matrice)

# X devient (nombre de fichiers, channel, listes, données), avec 1 pour le canal d'entrée requis par EEGNet
X = torch.Tensor(np.array(LST2)).unsqueeze(1) 
Y = torch.LongTensor(label)

# créer le DataLoader
dataset = TensorDataset(X, Y)
data_loader = DataLoader(dataset, batch_size=train_batch_size, shuffle=True)

print(f' Longueur de la liste: {len(LST2)}')
print(f' Nombre de fichiers: {len(label)}')
print(f' Forme de la liste des tensor: {X.shape}')
print("train batch size:",data_loader.batch_size,
     ", num of batch:", len(data_loader))