preprocessing_controller.py

"""
Created on Fri Feb 25 2022
@author:Tristan 
Fichier permettant l'extraction des angles du corpus complet (ou du  s8 )
prepare les données d'entrainements du réseau émulateur puis réalise 
un premier entrainement.
"""
#%% Imports
# - Third-party
import numpy as np 
import pandas as pd
import os 
import json

# path = "C:/Users/trist/Documents/Stage/perso/git_code/cobiopro_tristan"  # office path
# path = "C:/Users/trist/Documents/ENSAM/Stage INRIA/cobiopro_tristan"  # home path
# os.chdir(path) 

#%% 1. Lecture des données du csv

def read_df(n):   
    """ Fonction qui ouvre le corpus s8  si n = 0 
    et le corpus complet sinon (~15sec) """
    
    if n == 0:
        df = pd.read_csv('datasets/s8_corpus_ProxyDistal_RightHanded_Train.csv') #df contenant le corpus de s8  
        print("Loaded corpus s8 of 58844 ech")
    else:
        df = pd.read_csv('datasets/corpus_ProxyDistal_RightHanded_Train.csv') #df contenant le corpus complet
        print("-> Loaded full corpus of 933925 ech")
    df0 = df[["subjectId","timestamp","tgtNumber","tgtRed"]]
    df1 = df.loc[:,"tgtPos_x":"endEffCustQuat_w"] 
    df = pd.concat([df0, df1], axis=1) # df contenant data nécessaires pour le controleur et l'émulateur
    
    return df

# on choisit quel corpus on veut lire, prend 15s pour corpus complet
# df = read_df(n = 1)

#%% 2. Obtention de l'input du controleur

## pratique pour lancement par cellules #%% (permet de commencer ici)
## le preprocessing en amont pour ces données est commun a l'emulateur 

# if "index_tronc.csv" not in locals():
#     index_tronc = np.array(pd.read_csv('datasets/index_tronc.csv', index_col = 0))
        
# if 'mat_states' not in locals():
#     mat_states = np.array(pd.read_csv('datasets/state_remapped_refInit.csv' , index_col = 0)) # matrice contenant le corpus de tgt,
                       
# if 'tgt_states' not in locals():
#     tgt_states = np.array(pd.read_csv('datasets/tgt_states_remapped_refInit.csv' , index_col = 0)) # matrice contenant le corpus de tgt, PB: taille (n_sample , 8)
                

def split_truncate_traj_c(df, index_tronc, mat_states, tgt_states, n):
    """ Fonction qui renvoie les trajectoires tronquées et séparées 
    par cibles du contrôleur. Proche de split_truncate_traj_e mais différentes
    input et output
    (compilation time: 90s) """
    
    ## on réorganise le dataset pour dissocier les trajectoires par cible
    vect_cibles = np.array(df[["tgtNumber"]]).squeeze() # vecteur comportant le nb de cible atteintes a chaque configuration
    
    # initialisation de matrices/vecteurs contenant les trajectoires selon le corpus
    # (propre au preprocess controleur)
    tgt_ang = tgt_states[: , 3:5] # on conserve uniquement tgtP et tgtR 
    tgt_pos = tgt_states[: , :3] # on a besoin de la Pos pour changement repère dans refTgt plus tard  
    mat_inputs = np.concatenate((mat_states , tgt_ang), axis = 1) # on ajoute tgtP et tgtR
    ref_quat_in_W = np.array(df.loc[:,"refInitQuat_x":"refInitQuat_w"])
    
    if n == 1:
        vect_subj = np.array(df[["subjectId"]]) # liste contenant les subj id associé a chaque ech
        # useful pour reconstruction pos main a partir angle predit
        mat_inputs = np.concatenate((mat_inputs, tgt_pos), axis = 1) # on ajoute les pos de tgt
        mat_inputs = np.concatenate((mat_inputs, ref_quat_in_W), axis = 1) # on ajoute les quat de refInit pour changement repère dans refTgt plus tard
       
        mat_inputs = np.concatenate((mat_inputs, vect_subj), axis = 1) # on ajoute les index des sujets 
        mat_tgts = np.zeros((4305,), dtype = object) # matrice contenant les cibles pour le corpus complet
    
    elif n == 0: # Unfinished for corpus s8 
        mat_tgts = np.zeros((290,), dtype = object) # matrice contenant les cibles pour s8
    
    # on construit la matrice
    mat_cible = mat_inputs[0] # on initialise la matrice contenant une trajectoires de cible
    n_ech = 1 # incrément parcourant les échantillons
    nb_cible = 0 # nb de cible réellement parcourues
    i_c = 0 # incrément pour voir les variations dans vect_cible
    
    for i in vect_cibles[1:]: # boucle parcourant tgtNumber
        if i != i_c:
            mat_tgts[nb_cible] = mat_cible # on ajoute la trajectoire de cible 
            i_c = i
            nb_cible += 1
            mat_cible = mat_inputs[n_ech]
        else: 
            mat_cible = np.vstack(( mat_cible, mat_inputs[n_ech] )) # on ajoute les ech d'une trajectoire
            
        n_ech += 1
    mat_tgts[-1] = mat_cible # on ajoute la dernière cible pour i == i_c
    
    ## on tronque le corpus avec le choix du corpus de cible qu'on a fait précédémment ( en 1 ligne: + compact que pour emu)
    mat_tgts = np.array([ mat_tgts[k][: len(mat_tgts[k]) - int(index_tronc[k])  ] for k in range( len(mat_tgts) )] , dtype = object)
    
    # on tronque les inputs (propre au preprocess controleur) (variable pas utilisee)
    # mat_inputs = mat_tgts[0][0 , :14]
    # for c in mat_tgts:
    #     mat_inputs = np.vstack((mat_inputs , c[: , :14]))

    # mat_inputs = mat_inputs[ 1 :] # on enleve le premier ech qui se repète 
    
    len_tgts = [len(c) for c in mat_tgts] # matrice contenant la longueur des tgt 
    nb_ech_tronc = sum(len_tgts)    
    print("Corpus has been truncated and is now of ", nb_ech_tronc, " ech")
    
    if n == 1  : # corpus complet 
    
        # on tronque les données qu'on récup du corpus  
        vect_subj = mat_tgts[0][: , -1]
        for c in mat_tgts[1 : ]: 
            vect_subj = np.concatenate((vect_subj , c[: , -1]))
        vect_subj = list(vect_subj) # vecteur contenant le numéro de sujet par échantillon
        
        # vect_tgt = mat_tgts[0][: , 12:17]
        # for c in mat_tgts[1 : ]: 
        #     vect_tgt = np.concatenate((vect_tgt , c[: ,  12:17])) # Tgt Pitch, Roll, PosX, Y, Z
        
        # on remet mat_tgts avec les données qui nous interessent
        mat_tgts = np.array([ mat_tgts[k][: , :21] for k in range(len(mat_tgts))], dtype = object)
        
        
        ## UNCOMMENT IF you want to extract corpus tgt remapped_refInit to csv 
        # if not(os.path.exists("/datasets/mat_tgts_c.npy")): # sauvegarde de mat_tgts
        #     np.save('./datasets/mat_tgts_c.npy', mat_tgts)
        
        return mat_tgts, len_tgts, mat_inputs, vect_subj

# mat_tgts, len_tgts, mat_inputs, vect_subj = split_truncate_traj_c(df, index_tronc, mat_states, tgt_states, n = 1) 

#%% 3. Preparation des données pour le réseau controleur

## pratique pour lancement par cellules #%% (permet de commencer ici)
if 'mat_tgts' not in locals():
    mat_tgts_c = np.load("./datasets/mat_tgts_c.npy", allow_pickle = True)
  
if "Index_subj" not in locals():
    f = open("datasets/Index des sujets.json") 
    Index_subj = json.load(f)
    f.close()
    
def prepare_sets_c(mat_tgts, vect_subj, ID, n, nb_pas): # NOT USED ANYMORE   
    """Fonction qui renvoie le dataset d'entrainement et de test du controleur 
    avec k=17 fold et un dico avec les index de départ 
    du sujet test ID extrait (échantillon et trajectoire) 
    Très proche de prepare_sets_e1, elle nécessite le fichier vect_subj.json 
    donnant les index des sujets par cible. Elle diffère sur la construction 
    (surement pas au point pour n = 0)
    (THIS FUNCTION WAS NOT RUNNED RECENTLY ; previous compilation time: ~ 30 sec)
    """
    
    if n == 1 : 
      ## on créé les ensembles d'entrainements et de test 
        len_tgts = [len(c) for c in mat_tgts] 
        # on recupere les numéros des sujets
        subj_id = [10] # liste qui contiendra les numéros des 17 sujets
        k = subj_id[0] 
        
        for j in vect_subj: # boucle parcourant les ID de sujets
            if k != j:
                subj_id.append(j) # on ajoute les numéros de sujets
                k = j
        
        print("Le sujet extrait du training data est le sujet n°", subj_id[ID])
        
        # on détermine le 1er et dernier échantillon du set de test 
        ech_start = vect_subj.index(subj_id[ID]) # indice signalant a quel ech commence le sujet choisit 
        
        if ID == 16: # cas pour le dernier sujet
            ech_stop = vect_subj.index(subj_id[-1])
            
        else:
            ech_stop = vect_subj.index(subj_id[ ID + 1 ]) # indice signalant a quel ech finis le sujet 
            
        # on détermine l'indice de départ et de fin sur les cibles car on a séparé les tgt par cibles independamment des sujets     
        nb_ech = 0
        cible_start = 0 # indice qui sera celui de la 1ere cible DU SUJET extrait
        
        # on détermine les indices pour séparer le train et test set
        while nb_ech < ech_start:
            nb_ech += len_tgts[cible_start]
            cible_start += 1
            
        cible_stop = cible_start # indice qui sera celui de la derniere cible du sujet 
        while nb_ech < ech_stop:
            nb_ech += len_tgts[cible_stop]
            cible_stop += 1
            
        if ID == 16: # cas pour le dernier sujet    
            cible_stop = len(mat_tgts)
      
            
    elif n == 0: # corpus s8 (PAS ACTUALISE)
        cible_start = ID
        cible_stop = ID + 1
        
        # on commence à l'ech 0 
        if ID == 0:
            ech_start = 0
            ech_stop = len(mat_tgts[0]) - 1 # la trajectoire commence a 0 
            
        # on commence à l'ech où commence la cible 
        else: 
            ech_start = sum(len_tgts[ : ID]) 
            ech_stop = sum(len_tgts[ : ID + 1]) 
    
        print("La cible extraite du training data est la cible n°", ID + 1 ) # car on commence a cible 0 
        
    ## on construit les matrices train, test de fin de trajectoire
    ## if data shouldn't have a link with emulator
    # if ID == 0:
    #     train_inputs = np.array([ a for a in mat_inputs[ ech_stop : ]])
        
    # else:
    #     train_inputs = np.array([ a for a in mat_inputs[ : ech_start]])
    #     train_inputs = np.concatenate((train_inputs, 
    #                                 np.array([a for a in mat_inputs[ ech_stop : ]])))
        
    # test_inputs = np.array([ a for a in mat_inputs[ ech_start : ech_stop]])
    
    k = nb_pas
    ## construction de fin de trajectoires ( les indices peuvent être fonction de K pour choisir les portions de traj) 
    train_inputs = mat_tgts[0][ -k ] # initialisation des vecteurs d'état z(t) et z(t+1) du train set
    train_outputs = mat_tgts[0][-k + 1 , :10]
    if ID == 0 or cible_start == 0 : 
        train_inputs = mat_tgts[cible_stop][ -k ] # initialisation des vecteurs d'état z(t) et z(t+1) du train set
        train_outputs = mat_tgts[cible_stop][-k + 1 , :10]
        
    for c in mat_tgts[ : cible_start ]: # boucle parcourant les trajectoires par cible
        if np.array_equal(c, mat_tgts[0]):
            train_inputs = np.vstack((train_inputs, c[ -k + 1 : -1]))
            train_outputs = np.vstack((train_outputs, c[ -k + 2 : , :10]))
        else:
            train_inputs = np.vstack((train_inputs,  c[ -k : -1]))
            train_outputs = np.vstack((train_outputs, c[ -k + 1: , :10]))

    # on conserve un sujet (n°10) ou une cible pour le test
    test_inputs = mat_tgts[cible_start][ -k ]
    test_outputs = mat_tgts[cible_start][-k + 1 , :10]
    
    for c in mat_tgts[ cible_start : cible_stop ]:
        if np.array_equal(c, mat_tgts[cible_start]):
            test_inputs = np.vstack((test_inputs, c[ -k + 1 : -1]))
            test_outputs = np.vstack((test_outputs, c[ -k + 2 : , :10]))
            
        else:
            test_inputs = np.vstack((test_inputs, c[ -k : -1]))
            test_outputs = np.vstack((test_outputs, c[ -k + 1: , :10]))
            
            
    for c in mat_tgts[ cible_stop :  ]: 
        if np.array_equal(c, mat_tgts[cible_stop]): 
            train_inputs = np.vstack((train_inputs, c[ -k + 1 : -1]))
            train_outputs = np.vstack((train_outputs, c[ -k + 2 : , :10]))
        else:
            train_inputs = np.vstack((train_inputs,  c[ -k : -1]))
            train_outputs = np.vstack((train_outputs, c[ -k + 1: , :10]))

    # on reforme le dataset pour plus de simplicite
    x_train_c, y_train_c  = train_inputs, train_outputs
    x_test_c, y_test_c = test_inputs, test_outputs
    
    ## on stocke les index de départs 
    dic = dict(ech_start = ech_start, ech_stop = ech_stop  , cible_start = cible_start , cible_stop = cible_stop) 
    
    ## UNCOMMENT IF you want to extract previous controller training and test data
    # if not os.path.exists("datasets/previous_training_data/corpus_xtrain_c.npy"): 
        
    #     column_names1 = ["ShouPitch","ShouRoll","ArmYaw","ElbPitch", "ForearmYaw",
    #                         "WriPitch","WriRoll", "HandRoll", "HandPitch", "PosX",
    #                         "PosY", "PosZ", "TgtPitch","TgtRoll"]
    #     column_names2 =["ShouPitch","ShouRoll","ArmYaw","ElbPitch", "ForearmYaw",
    #                       "WriPitch","WriRoll", "HandRoll", "HandPitch", "PosX",
    #                       "PosY", "PosZ", "TgtPitch","TgtRoll"]
        
    #     df_xtrain_c = pd.DataFrame(x_train_c, columns = column_names1)
    #     df_ytrain_c = pd.DataFrame(y_train_c, columns = column_names2)
        
    #     df_xtest_c = pd.DataFrame(x_test_c, columns = column_names1)
    #     df_ytest_c = pd.DataFrame(y_test_c, columns = column_names2)
       
    #     df_xtrain_c.to_csv(path + "/datasets/previous_training_data/corpus_xtrain_c.csv", index = True)
    #     df_ytrain_c.to_csv(path + "/datasets/previous_training_data/corpus_ytrain_c.csv", index = True)
    #     df_xtest_c.to_csv(path + "/datasets/previous_training_data/corpus_xtest_c.csv", index = True)
    #     df_ytest_c.to_csv(path + "/datasets/previous_training_data/corpus_ytest_c.csv", index = True)
        
    return ((x_train_c, y_train_c), (x_test_c, y_test_c)), dic


def prepare_sets_c2(mat_tgts, Index_subj, kfold): # NOT USED ANYMORE
    """Fonction qui construit les train, val et test set avec 5 Fold 
    elle est très proche de prepare_sets_e3 et donc de prepare_sets_c1 mais 
    elle nécessite le fichier Index des sujets.json qui contient les indices 
    d'échantillons et trajectoires de chaque sujet (ou: d'avoir lancé 
    prepare_sets_e1 pour les 17 sujets et récupéré tout les vect_subj.json)
    (THIS FUNCTION WAS NOT RUNNED RECENTLY ; previous compilation time: ~ 30 sec)
    """
    
    k = kfold
    
    ## creation des mat_tgts train, test et val
    test_index_start, test_index_stop = Index_subj[-2]["cible_start"],  Index_subj[-1]["cible_stop"]
    mat_tgts_test = mat_tgts[test_index_start : test_index_stop]
    mat_tgts = mat_tgts[ : test_index_start]
    
    Index = [i["cible_start"] for i in Index_subj]
    Index = Index[: -1] # on enleve le dernier sujet test 
    
    if k == 1 : 
        k -= 1
    elif k == 2:
        k += 1
    elif k == 3 :
        k += 3
    elif k == 4 : 
        k += 5
    elif k == 5:
        k += 7
        
    index_start, index_stop =Index[k] , Index[k + 3] 
    mat_tgts_train = np.concatenate((mat_tgts[ : index_start] , mat_tgts[ index_stop : ]))
    mat_tgts_val = mat_tgts[ index_start : index_stop ]
    
    ## construction des train, val et test de fin de trajectoires 
    train_states_init = mat_tgts_train[0][ -40 ] # initialisation des vecteurs d'état z(t) et z(t+1) du train set
    train_states_final = mat_tgts_train[0][ -39 , :10 ]
    for c in mat_tgts_train: # boucle parcourant les trajectoires par cible
        if np.array_equal(c, mat_tgts_train[0]):
            train_states_init = np.vstack((train_states_init, c[ -39: -1]))
            train_states_final = np.vstack((train_states_final, c[ -38: , :10]))
        else:
            train_states_init = np.vstack((train_states_init, c[ -40: -1]))
            train_states_final = np.vstack((train_states_final, c[ -39: , :10]))
    train_control = train_states_final - train_states_init
    
    
    x_train = np.concatenate((train_states_init, train_control), axis = 1)
    y_train = train_states_final
    
    val_states_init = mat_tgts_val[0][ -40 ]
    val_states_final = mat_tgts_val[0][ -39 , :10 ]
    for c in mat_tgts_val:
        if np.array_equal(c, mat_tgts_val[0]):
            val_states_init = np.vstack((val_states_init, c[ -39: -1]))
            val_states_final = np.vstack((val_states_final, c[ -38: , :10]))
        else:
            val_states_init = np.vstack((val_states_init, c[ -40: -1]))
            val_states_final = np.vstack((val_states_final, c[ -40: -1]))
    val_control = val_states_final - val_states_init
    
    x_val = np.concatenate((val_states_init, val_control), axis = 1)
    y_val = val_states_final
    
    test_states_init = mat_tgts_test[0][0]
    test_states_final = mat_tgts_test[0][1]
    for c in mat_tgts_test:
        if np.array_equal(c, mat_tgts_test[0]):
            test_states_init = np.vstack((test_states_init, c[ -39: -1]))
            test_states_final = np.vstack((test_states_final, c[ -38: , :10]))
        else:
            test_states_init = np.vstack((test_states_init, c[ -40: -1]))
            test_states_final = np.vstack((test_states_final, c[ -40: -1]))
    test_control = test_states_final - test_states_init
    
    x_test = np.concatenate((test_states_init, test_control), axis = 1)
    y_test = test_states_final
    
    ## UNCOMMENT IF you want to extract previous training, val and test data
    # if not os.path.exists("datasets/training_data/corpus_xtest.csv"): # substituer plus tard par not in repertory 
        
    #     column_names1 = ["ShouPitch","ShouRoll","ArmYaw","ElbPitch", "ForearmYaw","WriPitch",
    #                                           "WriRoll", "PosX", "PosY", "PosZ", "u0","u1",
    #                                           "u2","u3","u4","u5","u6","u7","u8","u9"]
    #     column_names2 = column_names1[:10]
        
    #     df_xtrain = pd.DataFrame(x_train, columns = column_names1)
    #     df_ytrain = pd.DataFrame(y_train, columns = column_names2)
        
    #     df_xval = pd.DataFrame(x_val, columns = column_names1)
    #     df_yval = pd.DataFrame(y_val, columns = column_names2)
        
    #     df_xtest = pd.DataFrame(x_test, columns = column_names1)
    #     df_ytest = pd.DataFrame(y_test, columns = column_names2)
        
    #     df_xtrain.to_csv(path + "/datasets/training_data/corpus_xtrain_k="+str(k+1)+".csv", index = True)
    #     df_ytrain.to_csv(path + "/datasets/training_data/corpus_ytrain_k="+str(k+1)+".csv", index = True)
    #     df_xval.to_csv(path + "/datasets/training_data/corpus_xval_k="+str(k+1)+".csv", index = True)
    #     df_yval.to_csv(path + "/datasets/training_data/corpus_yval_k="+str(k+1)+".csv", index = True)
    #     df_xtest.to_csv(path + "/datasets/training_data/corpus_xtest.csv", index = True)
    #     df_ytest.to_csv(path + "/datasets/training_data/corpus_ytest.csv", index = True)

    print("-> Le Fold ",str(k + 1), " du train, val set et test set a été créé (cross val 5Fold)")               
    return ((x_train, y_train), (x_val, y_val), (x_test, y_test))


def prepare_sets_c3(mat_tgts, Index_subj, kfold, nb_pas):
    """ Fonction identique a prepare_sets_c2 mais qui prepare le dataset au 
    format notamment accepté par les RNN c'est à dire avec 3 dimensions : 
    (batch_size, time_steps, variable). Ce format est possible pour les 
    trajectoires de taille maximale nb_pas = 36 (plus petite trajectoire du corpus). 
    - On peut néanmoins construire des datasets au format correct jusque nb_pas = 50. 
    Si nb_pas >50 on perd les 3 dimensions. 
    - Pour les trajectoires plus grandes que 36 échantillons on propose uniquement 
    du padding pour l'instant (on comble les petites trajectoires par des 0 , INCOMPLET ).
    - Nb_pas part de la fin de la trajectoire
    Temps de compilation ~ 5sec
    """
    
    k = kfold
    
    ## creation des mat_tgts train, test et val
    test_index_start, test_index_stop = Index_subj[-2]["cible_start"],  Index_subj[-1]["cible_stop"]
    mat_tgts_test = mat_tgts[test_index_start : test_index_stop]
    mat_tgts = mat_tgts[ : test_index_start]
    
    Index = [i["cible_start"] for i in Index_subj]
    Index = Index[: -1] # on enleve le sujet du test set qui reste dans l'index
    
    if k == 1 : 
        k -= 1
    elif k == 2:
        k += 1
    elif k == 3 :
        k += 3
    elif k == 4 : 
        k += 5
    elif k == 5:
        k += 7
        
    index_start, index_stop = Index[k] , Index[k + 3] 
    mat_tgts_train = np.concatenate((mat_tgts[ : index_start] , mat_tgts[ index_stop : ]))
    mat_tgts_val = mat_tgts[ index_start : index_stop ]
    
    if True:
        ## training data with increment of one between output et input 

        x_train = np.array([c[- nb_pas : -1, :14] for c in mat_tgts_train])
        y_train = np.array([c[- nb_pas + 1 : , :12] for c in mat_tgts_train])
        
        x_val = np.array([c[- nb_pas : -1, :14] for c in mat_tgts_val])
        y_val = np.array([c[- nb_pas + 1 : , :12] for c in mat_tgts_val])
        
        x_test = np.array([c[- nb_pas : -1, :14] for c in mat_tgts_test])
        y_test = np.array([c[- nb_pas + 1 : , :12] for c in mat_tgts_test])
        
        
        ## UNCOMMENT IF changing of ref in refTgt necessary
        # tgt_train = mat_tgts_train[0][1:,12:17]
        # for t in mat_tgts_train[1:]:
        #     tgt_train = np.vstack( ( tgt_train, t[1: , 12:17]))
            
        # tgt_val = mat_tgts_val[0][1:,12:17]
        # for t in mat_tgts_val[1:]:
        #     tgt_val = np.vstack( ( tgt_val, t[1: , 12:17]))
        
        # tgt_test = mat_tgts_test[0][1:,12:17]
        # for t in mat_tgts_test[1:]:
        #     tgt_test = np.vstack( ( tgt_test, t[1: , 12:17]))
        
        # refInit_train = mat_tgts_train[0][1:,17:21]
        # for t in mat_tgts_train[1:]:
        #     refInit_train = np.vstack( ( refInit_train, t[1:,17:21]))
            
        # refInit_val = mat_tgts_val[0][1:,17:21]
        # for t in mat_tgts_val[1:]:
        #     refInit_val = np.vstack( ( refInit_val, t[1:,17:21]))
            
        # refInit_test = mat_tgts_test[0][1:,17:21]
        # for t in mat_tgts_test[1:]:
        #     refInit_test = np.vstack( ( refInit_test, t[1:,17:21]))
        
        
        ## UNCOMMENT IF you want to extract training, val and test data
        if not os.path.exists("datasets/training_data/x_train_c_k="+str(k)+"_"+str(nb_pas-1)+"steps.npy"):
        
            if k == 0 : 
                k = 1
            elif k == 3:
                k = 2
            elif k == 6 :
                k = 3
            elif k == 9 : 
                k = 4
            elif k == 12:
                k = 5
                
            np.save("datasets/training_data/x_train_c_k="+str(k)+"_"+str(nb_pas-1)+"steps.npy", x_train)
            np.save("datasets/training_data/y_train_c_k="+str(k)+"_"+str(nb_pas-1)+"steps.npy", y_train)
            np.save("datasets/training_data/x_val_c_k="+str(k)+"_"+str(nb_pas-1)+"steps.npy", x_val)
            np.save("datasets/training_data/y_val_c_k="+str(k)+"_"+str(nb_pas-1)+"steps.npy", y_val)
            np.save("datasets/training_data/x_test_c_k="+str(k)+"_"+str(nb_pas-1)+"steps.npy", x_test)
            np.save("datasets/training_data/y_test_c_k="+str(k)+"_"+str(nb_pas-1)+"steps.npy", y_test)
            print("-> Le Fold ",str(k), " du train, val set et test set a été créé (cross val 5Fold)")    
            
    ## UNCOMMENT IF you want to pad the smaller trajectories 
    # if nb_pas > 36: # Début de piste pour les trajectoires plus longues que 36 ech  (INCOMPLETE)
    #     ## trajectoires comblée par des zéros 
    #     # initialisation des matrices de train val et test
    #     x_train = np.zeros((mat_tgts_train.shape[0], nb_pas-1 , 14), dtype = object)
    #     y_train = np.zeros((mat_tgts_train.shape[0], nb_pas-1 , 12), dtype = object)
        
    #     x_val = np.zeros((mat_tgts_val.shape[0], nb_pas-1 , 14), dtype = object)
    #     y_val = np.zeros((mat_tgts_val.shape[0], nb_pas-1 , 12), dtype = object)
        
    #     x_test = np.zeros((mat_tgts_test.shape[0], nb_pas-1 , 14), dtype = object)
    #     y_test = np.zeros((mat_tgts_test.shape[0], nb_pas-1 , 12), dtype = object)
        
    #     for k in range(len(mat_tgts_train)-1):
    #         x_train[k] = mat_tgts_train[k][ - nb_pas : -1]
    #         y_train[k] = mat_tgts_train[k][ - nb_pas + 1 :, :12]
            
    #     for k in range(len(mat_tgts_val)-1):
    #         x_val[k] =mat_tgts_val[k][ - nb_pas :-1]
    #         y_val[k] =mat_tgts_val[k][ - nb_pas + 1 :, :12]
            
    #     for k in range(len(mat_tgts_test)-1):
    #         x_test[k] = mat_tgts_test[k][ - nb_pas :-1]
    #         y_test[k] = mat_tgts_test[k][ - nb_pas + 1 :, :12]
    
    ## UNCOMMENT IF changing of ref in refTgt necessary  
    # return ((x_train, y_train), (x_val,y_val), (x_test, y_test)), tgt_train, tgt_val, tgt_test, refInit_train, refInit_val, refInit_test
    return ((x_train, y_train), (x_val,y_val), (x_test, y_test))

## utilisation de prepareset_1 : 
## choix du sujet pr le test set dans [10, 11, 13, 14, 3, 5, 6, 7, 8, 9, 27, 28, 29, 30, 31, 32, 33]
## sujet ID = 0 est le sujet n°10
# dataset_c , dic = prepare_sets_c(mat_tgts, vect_subj, ID = 0, n = 1, nb_pas = 15) 

## utilisation de prepareset_3
# dataset_c  = prepare_sets_c3(mat_tgts_c, Index_subj, kfold = 1, nb_pas = 15)  

## UNCOMMENT IF changing of ref in refTgt necessary
# ((x_train, y_train), (x_val,y_val), (x_test, y_test)), tgt_train, tgt_val, tgt_test, refInit_train, refInit_val, refInit_test = prepare_sets_c3(mat_tgts, Index_subj, kfold = 1, nb_pas = 15) 
 
#%% Main function

def prepare_data_c():
    """ Fonction qui fait le preprocessing du dataset du controleur avec me 
    découpage pour cross validation fixé à 5fold (12 train, 3 val et 2 sujets test).
    Elle necessite au minimum le csv du corpus complet, Index des sujets.json
    et les 4 csv issus de preprocessing_emulator suivant : index_tronc.csv, 
    state_remapped_refInit.csv, tgt_states_remapped_refInit.csv. 
    Temps de compilation rapide
    """
    
    print("Pour extraire le dataset vérifier que la partie réservée dans prepare_sets_c3 est décommentée")
    
    if True: #  # Mettre a True si on part du corpus brut
        ## choix du corpus
        df = read_df(n = 1) 
        
        ## on load le dataset commun avec l'émulateur
        index_tronc = np.array(pd.read_csv('datasets/index_tronc.csv', index_col = 0))
        mat_states = np.array(pd.read_csv('datasets/state_remapped_refInit.csv' , index_col = 0)) # matrice contenant le corpus de tgt,
        tgt_states = np.array(pd.read_csv('datasets/tgt_states_remapped_refInit.csv' , index_col = 0)) # matrice contenant le corpus de tgt, PB: taille (n_sample , 8)
        
        ## construction des input propres au controleur
        mat_tgts, _, _, _ = split_truncate_traj_c(df, index_tronc, mat_states, tgt_states, n = 1) 
        
        ## construction des données d'entrainements en 5fold
        f = open("datasets/Index des sujets.json")
        Index_subj = json.load(f)
        f.close()
        
        dataset_c = prepare_sets_c3(mat_tgts, Index_subj, kfold = 1, nb_pas = 15)  
        return dataset_c
    
    if False: # Mettre a True si mat_tgts_c déja enregistrée
        mat_tgts = np.load("datasets/mat_tgts_c.npy", allow_pickle = True)
        f = open("datasets/Index des sujets.json")
        Index_subj = json.load(f)
        f.close()
        
        ## choix du kfold
        dataset_c = prepare_sets_c3(mat_tgts, Index_subj, kfold = 1, nb_pas = 15)  
        return dataset_c

def main():
    prepare_data_c()

if __name__=="__main__":
    main()