protected function calculateMeans(array $data, array $classes) : array
    {
        $means = [];
        $counts= [];
        $overallMean = array_fill(0, count($data[0]), 0.0);

        foreach ($data as $index => $row) {
            $label = array_search($classes[$index], $this->labels);

            foreach ($row as $col => $val) {
                if (!isset($means[$label][$col])) {
                    $means[$label][$col] = 0.0;
                }
                $means[$label][$col] += $val;
                $overallMean[$col] += $val;
            }

            if (!isset($counts[$label])) {
                $counts[$label] = 0;
            }

            ++$counts[$label];
        }

        foreach ($means as $index => $row) {
            foreach ($row as $col => $sum) {
                $means[$index][$col] = $sum / $counts[$index];
            }
        }

        // Calculate overall mean of the dataset for each column
        $numElements = array_sum($counts);
        $map = function ($el) use ($numElements) {
            return $el / $numElements;
        };
        $this->overallMean = array_map($map, $overallMean);
        $this->counts = $counts;

        return $means;
    }