#!/usr/bin/python3
# -*- coding: utf-8 -*-

from PyQt4.QtCore import *
from PyQt4.QtGui import *


class TreeItem:
    def __init__(self, data, parent = None):
        self.parentItem = parent
        if type(data) == list:
            self.itemData = data
        else:
            self.itemData = [data]
        self.childItems = []
        # -1 for root item
        if parent == None:
            self.row = -1

    def appendChild(self, item):
        item.row = len(self.childItems)
        self.childItems.append(item)

    def row(self):
        return self.row

    def columnCount(self):
        return(len(self.itemData))

    def parent(self):
        return self.parent

    def child(self, row):
        if row < len(self.childItems):
            return self.childItems[row]
        else:
            return None
    def childCount(self):
        return len(self.childItems)

    def data(self, column):
        if column < len(self.itemData):
            return self.itemData[column]
        else:
            return None

class TreeModel(QAbstractItemModel):
    def __init__(self, parent = None):
        QAbstractItemModel.__init__(self, parent)
        self.rootItem = TreeItem(["Title"], None)
        self.populateModel()

    def index(self, row, column, parent = QModelIndex()):
        if not parent.isValid():
            parentItem = self.rootItem
        else:
            parentItem = parent.internalPointer()
        childItem = parentItem.child(row)
        if childItem:
            return self.createIndex(row, column, childItem)
        else:
            return QModelIndex()

    def parent(self, index):
        if not index.isValid():
            return QModelIndex()
        childItem = index.internalPointer()
        parentItem = childItem.parent()
        if parentItem == rootItem:
            return QModelIndex()
        return self.createIndex(parentItem.row(), 0, parentItem)

    def columnCount(self, parent = QModelIndex()):
        if parent.isValid():
            return parent.internalPointer().columnCount()
        else:
            return self.rootItem.columnCount()

    def rowCount(self, parent = QModelIndex()):
        if parent.column():
            return 0
        if not parent.isValid():
            parentItem = self.rootItem
        else:
            parentItem = parent.internalPointer()
        return parentItem.childCount()

    def headerData(self, section, orientation, role = Qt.DisplayRole):
        if orientation == Qt.Horizontal and role == Qt.DisplayRole:
            return self.rootItem.data(section)
        return None

    def flags(self, index):
        if not index.isValid():
            return Qt.NoItemFlags
        return Qt.ItemIsEditable | Qt.ItemIsEnabled | Qt.ItemIsSelectable

    def data(self, index, role):
        if not index.isValid():
            return None
        if role != Qt.DisplayRole:
            return None
        item = index.internalPointer()
        return item.data(index.column())

    def appendChild(self, item, parent = None):
        if not parent or not parent.isValid():
            parentItem = self.rootItem
        else:
            parentItem = parent.internalPointer()
        parentItem.appendChild(item)

    def populateModel(self):
        items = ["Good", "Better", "Best"]
        for i, str in enumerate(items):
            treeItem = TreeItem([str])
            self.appendChild(treeItem)
        items = ["Never", "Rest"]
        index = self.index(0, 0, QModelIndex())
        for i, str in enumerate(items):
            self.appendChild(TreeItem([str], treeItem), index)


if __name__ == "__main__":
    app = QtGui.QApplication(sys.argv)
    treeView = QTreeView()
    treeModel = TreeModel()
    treeView.setModel(treeModel)
    treeView.show()
    index = treeModel.index(0,0)
    print(treeModel.data(index, Qt.DisplayRole))
    index = treeModel.index(1,0, index)
    print(treeModel.data(index, Qt.DisplayRole))
    sys.exit(app.exec_())

if parent.column():
    return 0