function behaviorCloningTrainInfo = behaviorCloningPrepareDataImFlyingRobot(dataStruct)
% Prepare training data for flying robot with behavior cloning
% 
% Copyright 2019 The MathWorks, Inc.

% Extract the dataStruct variables
numObservations = dataStruct.numObservations;
numActions = dataStruct.numActions;
actionIndex = dataStruct.actionIndex;
stateIndex = dataStruct.stateIndex;
data = dataStruct.data;

% calculate number of validation rows
totalRows = size(data,1);
validationSplitPercent = dataStruct.validationSplitPercent;
numValidationDataRows = floor(validationSplitPercent*totalRows);

% Test split
testSplitPercent = dataStruct.testSplitPercent;
numTestDataRows = floor(testSplitPercent*totalRows);

% create validation and training data
randomIdx = randperm(totalRows,numValidationDataRows + numTestDataRows);
randomData = data(randomIdx,:);    
trainDataIdx = setdiff(1:totalRows,randomIdx);

newValidationData = randomData(1:numValidationDataRows,:);
newTestData = randomData(numValidationDataRows + 1:end,:);
newTrainData = data(trainDataIdx,:);

numTrainDataRows = size(newTrainData,1);

% prepare new dataset for training
behaviorCloningTrainInfo.newTrainInput = newTrainData(:,stateIndex);
behaviorCloningTrainInfo.newTrainOutput = newTrainData(:,actionIndex);

newValidationInput = newValidationData(:,stateIndex);
newValidationOutput = newValidationData(:,actionIndex);
behaviorCloningTrainInfo.newValidationCellArray = {newValidationInput, newValidationOutput};

behaviorCloningTrainInfo.newTestDataInput = newTestData(:,stateIndex);
behaviorCloningTrainInfo.newTestDataOutput = newTestData(:,actionIndex);

end