import random
import time
from sortedcontainers import SortedDict

# These functions are used to populate the database with
# values that are kinda-sorta like the ones in my app.

def xdi(B,N,M):
    '''
            eXponential Distribution of Integers, XDI
    Return a list of length N of integers with max value M, from an exponential
    distribution with beta B.
    '''
    M = int(M) # just in case
    # pre-allocate the list rather than building it with .append()
    r = [0]*N
    for j in range(N) :
        r[j] = int( random.expovariate(B) ) % M
    return r

RWTALPH = 'aeiouåáïbcdfghjklmnpqrstvwxyz_BCDFGHJKLMNPQRSTVWZ'
def rwt(N):
    '''
            Random Word-like Token, RWT
    Return a word-like token comprising N Latin-1 letters with an
    emphasis on vowels.
    '''
    ii = xdi( 0.1, N, len(RWTALPH)-1 )
    return ''.join( [ RWTALPH[i] for i in ii ] )

def rft():
    '''
            Random Flag Token
    Return an 8-char token composed mostly of dashes with a sprinkle of Xs,
    e.g. "---X--X-" or "XX------".

    '''
    ii = xdi( 2, 8, 2 )
    return ''.join( ['-X'[i] for i in ii ] )

#
# Define the database. All this code is out here at the module level --
# instead of inside the Table Model class -- because in the real app, the
# database is managed in a separate module entirely.
#
# The database comprises a single SortedDict, DATABASE, in which the keys are
# rwt(8) "word" tokens, and the values are tuples comprising ( word, int in
# 0-99, rft() flag string). (Yes this is redundant in having the word as both
# key and value but it simplifies coding the data() method of the model.)
#
DATABASE = SortedDict()
NUMBER_OF_ROWS = 10000 # how many rows to create
#
# The contents are accessed via a valuesView object, VALUES_BY_ROW. This can
# be indexed, so VALUES_BY_ROW[J] is the values for the Jth entry in the
# ascending order of the dict keys. VALUES_BY_ROW[J][2] is the column-2 data
# from row J (the flag string column).
#
VALUES_BY_ROW = None # valuesview here
#
# In order to provide the indirection required for sorting, access to VALUES_BY_ROW is
# by way of a sort vector of indices. As many as 6 such vectors, one for each column
# for each sort order, are stored in a two lists, SORT_UP_VECTORS and SORT_DOWN_VECTORS.
# When the database is initialized, only one of these is created, the sort-up for
# column 0. Others are made as needed.
#
SORT_UP_VECTORS = [ None, None, None ]
SORT_DOWN_VECTORS = [ None, None, None ]
#
# When the database is "refreshed" (filled with new data), all vectors are
# discarded and a new SORT_UP_VECTORS[0] created as simply
# range(len(DATABASE)). Thus VALUES_BY_ROW[ SORT_UP_VECTORS[0][J] ] is
# exactly equivalent to VALUES_BY_ROW[J], except it takes a bit longer.
#

# Clear the database, called when the model initializes and before Refresh

def clear_the_db() :
    global DATABASE, VALUES_BY_ROW, SORT_UP_VECTORS, SORT_DOWN_VECTORS
    DATABASE.clear()
    VALUES_BY_ROW = None
    SORT_UP_VECTORS = [ None, None, None ]
    SORT_DOWN_VECTORS = [ None, None, None ]

# Rebuild the database by repopulating it with phony data. This emulates what
# worddata.py does when it scans the book and tabulates its vocabulary.

def rebuild_the_db():
    global DATABASE, VALUES_BY_ROW, NUMBER_OF_ROWS, SORT_UP_VECTORS
    # Throw everything away. In the real app, worddata doesn't throw everything away,
    # it only clears the counts in column 1.
    clear_the_db()
    # Create values for column 1, small random integers. Use the length
    # of column 1 to control the loop.
    column_1 = xdi( 0.05, NUMBER_OF_ROWS, 100 )
    # Populate the database with rows of stuff
    for k in column_1 :
        w = rwt(8) # a random "word"
        f = rft()  # a random "flag" string
        DATABASE[ w ] = ( w, k, f )
    # Create the indexable view of the data
    VALUES_BY_ROW = DATABASE.values()
    # Initialize the first sort vector to "idempotent" values
    SORT_UP_VECTORS[0] = range( len( DATABASE ) )

#
# On request from the table model, return a sort vector for a given column
# (0, 1 or 2) and given sort order. If we have a vector for that column and
# sort order, just return it.
#
# Otherwise, ask if we have a vector for ascending order on this column. If
# we do not, build a vector for accessing the values of DATABASE in ascending
# sequence over that column. If the request is for ascending order, return
# that.
#
# Otherwise, make a vector for descending order by applying reversed() to the
# ascending order vector and place it in SORT_DOWN_VECTORS[col]
#
# We build the ascending vector as follows. Make a SortedDict in which each
# key is the concatenation of the column value for a row (which may have
# dups) and the primary key for the same row. The combination is unique and
# sorts properly.
#
# The value for each key is the index of that row in DATABASE. Because this
# is a SortedDict, the values are a sort vector for the DATABASE rows in
# column-value order.
#
# Place in SORT_UP_VECTORS[col] a valuesView on this sorted dict. This is
# effectively a vector of DATABASE row numbers in the desired sequence.

# What we want is a valuesView on a dict, but we have to stow the dict
# itself somewhere:

SORT_UP_DICTS = [ None, None, None ] # only 1 & 2 used

def get_column_vector( col, order ):
    global DATABASE, VALUES_BY_ROW, SORT_UP_VECTORS, SORT_DOWN_VECTORS
    vector = SORT_UP_VECTORS[ col ] if order == Qt.AscendingOrder else SORT_DOWN_VECTORS[ col ]
    if vector is None : # we need to make one
        if SORT_UP_VECTORS[ col ] is None : # we have no up-vector
            # build a format string for col 1 int:str or col 2 str:str
            key_format = '{:05}{}' if col == 1 else '{}{}'
            sort_dict = SortedDict()
            for j in range( len(DATABASE) ) :
                k = key_format.format( VALUES_BY_ROW[j][col], VALUES_BY_ROW[j][0] )
                sort_dict[k] = j
            SORT_UP_DICTS[ col ] = sort_dict # save the dict itself
            SORT_UP_VECTORS[ col ] = sort_dict.values()
        vector = SORT_UP_VECTORS[ col ] # be optimistic
        # We have a sort-up vector, or was it sort-down we need?
        if order == Qt.DescendingOrder : # ok, make the reverse
            SORT_DOWN_VECTORS[ col ] = [j for j in reversed( SORT_UP_VECTORS[ col ] ) ]
            vector = SORT_DOWN_VECTORS [ col ]
    return vector

# Define the Table Model, instrumented to count certain calls

from PyQt5.QtCore import Qt, QAbstractTableModel, pyqtSignal

class Model( QAbstractTableModel ):
    tableSorted = pyqtSignal()

    def __init__ ( self, parent=None ) :
        super().__init__( parent )
        clear_the_db()
        self.access_counts = [0, 0, 0]
        self.sort_order = Qt.AscendingOrder
        self.sort_column = 0
        self.sort_vector = []
        self.sort_time = 0.0 # TIMING

    def rowCount( self, index ) :
        global DATABASE
        if index.isValid() : return 0
        return len(DATABASE) # 0 until Refresh is called

    def columnCount( self, index ) :
        if index.isValid() : return 0
        return 3

    def data(self, index, role ) :
        global VALUES_BY_ROW
        if role != Qt.DisplayRole :
            return None
        row = index.row()
        col = index.column()
        self.access_counts[col] += 1
        sort_row = self.sort_vector[ row ]
        return VALUES_BY_ROW[ sort_row ][ col ]

    def clear_counts( self ) :
        self.access_counts = [0, 0, 0]

    def counts( self ) :
        return list( self.access_counts )

    # reimplement QAbstractItemModel.sort() which receives two arguments,
    # the column number to sort, and the sort order, Qt.AscendingOrder
    # or Qt.DescendingOrder. Then emit layoutChanged to repaint the view.

    def sort( self, col, order ) :
        t0 = time.process_time() # TIMING
        self.sort_column = col
        self.sort_order = order
        self.sort_vector = get_column_vector( col, order )
        self.layoutChanged.emit()
        self.sort_time = time.process_time() - t0    # TIMING
        self.tableSorted.emit()    # TIMING

    # Override endResetModel() to make sure to refresh the sort
    # when there is new data.

    def endResetModel( self ) :
        super().endResetModel()
        self.sort( self.sort_column, self.sort_order )

# Define the Table View, which at this point is quite minimal.
# The View is instantiated from MainWindow, which also connects
# the model to it.

from PyQt5.QtWidgets import QTableView

class View( QTableView ):
    def __init__ ( self, parent=None ) :
        super().__init__( parent )
        self.setSortingEnabled( True )
        self.sortByColumn( 0, Qt.AscendingOrder )

# Define the main window which is the visual face of this app.

from PyQt5.QtWidgets import (
    QLabel,
    QMainWindow,
    QPushButton,
    QVBoxLayout,
    QHBoxLayout,
    QWidget
    )

class Main( QMainWindow ) :
    def __init__ ( self ) :
        super().__init__( )
        self.times = [0, 0, 0]
        clear_the_db() # make sure to start with 0 rows
        self._uic() # all the layout stuff out of line
        self.refresh_button.clicked.connect( self.do_refresh )
        self.table_model.tableSorted.connect( self.update_sort_time )

    # Slot called when Refresh is clicked:
    # * clear the counts
    # * start model reset
    # * rebuild the DATABASE, timing it
    # * end the model reset, timing that
    # * update the labels displaying call counts and times

    def do_refresh( self ) :
        self.table_model.clear_counts()
        self.table_model.beginResetModel()
        self.times[0] = time.process_time()
        rebuild_the_db()
        self.times[1] = time.process_time()
        self.table_model.endResetModel()
        QApplication.processEvents()
        self.times[2] = time.process_time()
        self.update_labels()

    # Slot called when a sort happens. Get the process time from
    # the table model and put it in the sort_time_label.

    def update_sort_time( self ) :
        self.sort_time_label.setText ('{:02.5f}'.format( self.table_model.sort_time  ) )

    def update_labels(self) :
        [c0, c1, c2] = self.table_model.counts()
        [t0, t1, t2] = self.times
        self.c0_label.setText( str(c0) )
        self.c1_label.setText( str(c1) )
        self.c2_label.setText( str(c2) )
        self.td_label.setText( '{:02.5f}'.format(t1-t0) )
        self.tr_label.setText( '{:02.5f}'.format(t2-t1) )

    def _make_label( self, text='0' ) :
        # just make a right-aligned label out of line
        L = QLabel(text)
        L.setAlignment( Qt.AlignRight | Qt.AlignVCenter )
        return L

    def _uic( self ) :
        # create the table
        self.table_view = View( parent=self )
        self.table_model = Model( parent=self )
        self.table_view.setModel( self.table_model )
        # create the refresh button, put it in an hbox by with the sort time
        self.refresh_button = QPushButton( "Refresh" )
        hb0 = QHBoxLayout()
        hb0.addWidget(self.refresh_button, 0)
        hb0.addStretch(1)
        hb0.addWidget( self._make_label( 'sort time:'), 0  )
        self.sort_time_label = self._make_label()
        hb0.addWidget( self.sort_time_label )
        # create a set of labels to display counts of
        # entry to the model.data() method and times
        self.c0_label = self._make_label() # display role calls to column 0
        self.c1_label = self._make_label() # 1
        self.c2_label = self._make_label() # 2
        self.tr_label = self._make_label() # time to reset the model
        self.td_label = self._make_label() # time to rebuild the db
        # build the row of call numbers
        hb1 = QHBoxLayout()
        hb1.addStretch(1) # push this row to the right
        hb1.addWidget( self._make_label( 'Display role calls col 0:' ) )
        hb1.addWidget( self.c0_label )
        hb1.addStretch(0)
        hb1.addWidget( self._make_label( 'col 1:' ) )
        hb1.addWidget( self.c1_label )
        hb1.addStretch(0)
        hb1.addWidget( self._make_label( 'col 2:' ) )
        hb1.addWidget( self.c2_label )
        # build the row of times
        hb2 = QHBoxLayout()
        hb2.addStretch(1)
        hb2.addWidget( self._make_label( 'Seconds to build DB:' ) )
        hb2.addWidget( self.td_label )
        hb2.addStretch(0)
        hb2.addWidget( self._make_label( 'to reset model:' ) )
        hb2.addWidget( self.tr_label )
        # stack up the central layout
        vb = QVBoxLayout()
        vb.addLayout( hb0, 0 )
        vb.addWidget( self.table_view, 1 )
        vb.addLayout( hb1, 0 )
        vb.addLayout( hb2, 0 )
        # put all that in a widget and make the widget our central layout
        wij = QWidget()
        wij.setLayout( vb )
        wij.setMinimumSize( 500, 500 )
        self.setCentralWidget( wij )


if __name__ == '__main__' :
    from PyQt5.QtWidgets import QApplication
    the_app = QApplication([])
    main_window = Main()
    main_window.show()
    the_app.exec_()