diff --git a/src/common.odin b/src/common.odin
index 72a9345..3381a6d 100644
--- a/src/common.odin
+++ b/src/common.odin
@@ -17,7 +17,7 @@ Problem :: struct {
 	name:            string,
 	chromosome_size: int,
 	fitness_proc:    proc(_: Chromosome) -> f64,
-	maximize:        bool, // true = higher better, false = lower better
+	maximize:        bool,
 }
 
 Stats :: struct {
diff --git a/src/feature_selection.odin b/src/feature_selection.odin
index 888caf4..1664cd4 100644
--- a/src/feature_selection.odin
+++ b/src/feature_selection.odin
@@ -1,5 +1,6 @@
 package main
 
+import "core:math/rand"
 import "core:container/bit_array"
 import "core:encoding/csv"
 import "core:fmt"
@@ -63,16 +64,29 @@ fitness_features :: proc(chrom: Chromosome) -> f64 {
 		delete(y)
 	}
 
-	// Train/test split
-	X_train, X_test, y_train, y_test := train_test_split(X, y, 0.2, RANDOM_SEED)
-	defer {
-		for row in X_train {delete(row)}
-		for row in X_test {delete(row)}
-		delete(X_train)
-		delete(X_test)
-		delete(y_train)
-		delete(y_test)
-	}
+	// NOW split based on actual data size
+    n := len(X)
+    test_count := int(f64(n) * 0.2)
+    train_count := n - test_count
+
+    // Create indices and shuffle ONCE per fitness eval
+    indices := make([]int, n)
+    defer delete(indices)
+    for i in 0 ..< n {
+        indices[i] = i
+    }
+    rand.shuffle(indices[:])  // Uses your seeded generator
+
+    // Split using shuffled indices
+    X_train, X_test, y_train, y_test := split_by_indices(X, y, indices, train_count)
+    defer {
+        for row in X_train {delete(row)}
+        for row in X_test {delete(row)}
+        delete(X_train)
+        delete(X_test)
+        delete(y_train)
+        delete(y_test)
+    }
 
 	// Train and evaluate
 	beta := train_linear_regression(X_train, y_train)
@@ -84,6 +98,37 @@ fitness_features :: proc(chrom: Chromosome) -> f64 {
 	return rmse(predictions, y_test)
 }
 
+split_by_indices :: proc(
+    X: [][]f64,
+    y: []f64,
+    indices: []int,
+    train_count: int,
+) -> ([][]f64, [][]f64, []f64, []f64) {
+    n_features := len(X[0])
+    test_count := len(indices) - train_count
+
+    X_train := make([][]f64, train_count)
+    X_test := make([][]f64, test_count)
+    y_train := make([]f64, train_count)
+    y_test := make([]f64, test_count)
+
+    for i in 0 ..< train_count {
+        idx := indices[i]
+        X_train[i] = make([]f64, n_features)
+        copy(X_train[i], X[idx])
+        y_train[i] = y[idx]
+    }
+
+    for i in 0 ..< test_count {
+        idx := indices[train_count + i]
+        X_test[i] = make([]f64, n_features)
+        copy(X_test[i], X[idx])
+        y_test[i] = y[idx]
+    }
+
+    return X_train, X_test, y_train, y_test
+}
+
 get_selected_features :: proc(dataset: Dataset, chrom: Chromosome) -> ([][]f64, []f64) {
 	n_rows := len(dataset)
 
diff --git a/src/linreg.odin b/src/linreg.odin
index 8e29115..6697227 100644
--- a/src/linreg.odin
+++ b/src/linreg.odin
@@ -1,5 +1,6 @@
 package main
 
+import "core:hash"
 import "core:math"
 import "core:math/rand"
 
@@ -114,60 +115,3 @@ rmse :: proc(predictions: []f64, actual: []f64) -> f64 {
 	}
 	return math.sqrt(sum / f64(len(predictions)))
 }
-
-// Train/test split
-train_test_split :: proc(
-	X: [][]f64,
-	y: []f64,
-	test_size: f64 = 0.2,
-	random_seed: u64 = 0,
-) -> (
-	[][]f64,
-	[][]f64,
-	[]f64,
-	[]f64,
-) {
-	n := len(X)
-	if n == 0 || len(X[0]) == 0 {
-		return nil, nil, nil, nil
-	}
-
-	n_features := len(X[0])
-	test_count := int(f64(n) * test_size)
-	train_count := n - test_count
-
-	// Shuffle indices
-	indices := make([]int, n)
-	defer delete(indices)
-	for i in 0 ..< n {
-		indices[i] = i
-	}
-
-	rng := rand.create(random_seed)
-	context.random_generator = rand.default_random_generator(&rng)
-	rand.shuffle(indices[:])
-
-	// Allocate
-	X_train := make([][]f64, train_count)
-	X_test := make([][]f64, test_count)
-	y_train := make([]f64, train_count)
-	y_test := make([]f64, test_count)
-
-	// Copy training data
-	for i in 0 ..< train_count {
-		idx := indices[i]
-		X_train[i] = make([]f64, n_features)
-		copy(X_train[i], X[idx])
-		y_train[i] = y[idx]
-	}
-
-	// Copy test data
-	for i in 0 ..< test_count {
-		idx := indices[train_count + i]
-		X_test[i] = make([]f64, n_features)
-		copy(X_test[i], X[idx])
-		y_test[i] = y[idx]
-	}
-
-	return X_train, X_test, y_train, y_test
-}
diff --git a/src/main.odin b/src/main.odin
index 54c6d5d..42362ca 100644
--- a/src/main.odin
+++ b/src/main.odin
@@ -1,11 +1,16 @@
 package main
 
+import "core:math/rand"
 import "core:fmt"
+import "base:runtime"
 
 main :: proc() {
 	// Choose problem
 	problem_type := "feature_selection" // or "knapsack"
 
+	state := rand.create(RANDOM_SEED)
+	context.random_generator = runtime.default_random_generator(&state)
+
 	problem: Problem
 	switch problem_type {
 	case "knapsack":