package nn

import (
	"math"
	"math/rand"

	"github.com/go-gota/gota/dataframe"
)

func (nn *NN) trainTestSplit() (dataframe.DataFrame, dataframe.DataFrame, dataframe.DataFrame, dataframe.DataFrame) {
	// now we split the data into training
	// and testing based on user specified
	// nn.TestSize.
	nRows := nn.Df.Nrow()
	testRows := int(math.Floor(float64(nRows) * nn.TestSize))

	// subset the testing data
	// randomly select trainRows number of rows
	randStrt := rand.Intn(int(math.Floor(float64(nRows) * nn.TestSize)))
	test := nn.Df.Subset([]int{randStrt, randStrt + testRows})

	// use what is left for training
	allIndices := make([]int, nRows)
	for i := range allIndices {
		allIndices[i] = i
	}

	// Remove the test indices using slice append and variadic parameter
	trainIndices := append(allIndices[:randStrt], allIndices[randStrt+testRows:]...)

	// Create the train DataFrame using the trainIndices
	train := nn.Df.Subset(trainIndices)

	XTrain = train.Select(nn.Features)
	YTrain = train.Select(nn.Target)
	XTest = test.Select(nn.Features)
	YTest = test.Select(nn.Target)

	return XTrain, XTest, YTrain, YTest
}