Добавить чекбоксы в заголовки столбцов таблицы

from PyQt5.QtWidgets import (QApplication, QTableView, QTableWidget,
QHeaderView, QStyleOptionButton, QStyle, QWidget, QGridLayout, QPushButton)
from PyQt5.QtGui import QStandardItemModel, QFont
from PyQt5.QtCore import Qt, pyqtSignal, QSize, QRect
 
class HeaderView(QHeaderView):
    def __init__(self, orientation=Qt.Horizontal, parent=None):
        super().__init__(orientation, parent)
        self.setSectionsClickable(True)
        self.sectionClicked.connect(self.on_clicked)
        self.sectionCountChanged.connect(self.on_count_changed)
        
    def on_count_changed(self):
        self.state = [True for _ in range(self.count())]
        
    def get_state(self):
        return [i for i in range(len(self.state)) if self.state[i]]
        
    def on_clicked(self, idx):
        self.state[idx] ^= True
        
    def paintSection(self, painter, rect, logicalIndex):
        painter.save()
        QHeaderView.paintSection(self, painter, rect, logicalIndex)
        painter.restore()
        option = QStyleOptionButton()
        option.initFrom(self)
        option.rect = QRect(rect.left(), rect.top(), 20, 20)
        if self.state[logicalIndex]: option.state |= QStyle.State_On
        else: option.state |= QStyle.State_Off
        self.style().drawPrimitive(QStyle.PE_IndicatorCheckBox, option, painter)
 
class W(QWidget):
    def __init__(self):
        super().__init__()
        self.head = HeaderView()
        grid = QGridLayout(self)
        table = QTableWidget(3,4)
        table.setHorizontalHeader(self.head)
        table.setHorizontalHeaderLabels(['один','два','шестнадцать','четыре'])
        btn = QPushButton('State')
        btn.clicked.connect(self.on_btn)
        grid.addWidget(table)
        grid.addWidget(btn)
        
    def on_btn(self):
        t = self.head.get_state()
        print(t)
 
 
app = QApplication([])
w = W()
w.show()
app.exec()