# usage of multiprocessing for calculating the power of jordan matrices
# by Stefan Koch

# German: http://eliteinformatiker.de/2012/06/07/mehrere-cpu-kerne-mit-python-nutzen-jordan-matrizenpotenzierung/


import multiprocessing
import random

class Matrixmultiplier(multiprocessing.Process):
    def __init__(self, task_queue, result_queue):
        multiprocessing.Process.__init__(self)
        self.task_queue = task_queue
        self.result_queue = result_queue

    def run(self):
        proc_name = self.name
        while True:
            (counter, next_task) = self.task_queue.get()
            if next_task is None:
                # Poison pill means we should exit
                print('%s: Exiting' % proc_name)
                break
            answer = next_task.calculate()
            self.result_queue.put((counter, answer))
        return


class Task(object):
    def __init__(self, matrix):
        self.matrix = matrix
    
    def calculate(self):
        return matrix_multiply(self.matrix, self.matrix)


def matrix_multiply(matrix1, matrix2):
    new_matrix = zero_matrix(len(matrix1), len(matrix2[0]))
    
    for i in range(len(matrix1)):
        for j in range(len(matrix2[0])):
            for k in range(len(matrix2)):
                new_matrix[i][j] += matrix1[i][k]*matrix2[k][j]
    
    return new_matrix

def zero_matrix(width, height):
    return [[0 for j in range(width)] for i in range(height)]

def zerofill_matrix(matrix, width, height):
    matrix_width = 0
    if len(matrix) != 0:
        matrix_width = len(matrix[0])
    
    for i in range(height):
        matrix.append([0 for i in range(matrix_width)])
    
    for i in range(len(matrix)):
        for j in range(width):
            matrix[i].append(0)
    return matrix

def extract_matrix(matrix, fro, to):
    new_matrix = zero_matrix(to-fro, to-fro)
    for i in range(len(new_matrix)):
        for j in range(len(new_matrix[0])):
            new_matrix[i][j] = matrix[fro+i][fro+j]
    return new_matrix
    
def split_matrix(matrix):
    sub_matrices = []
    
    last_split = 0
    for i in range(len(matrix)):
        if i == len(matrix)-1:
            sub_matrices.append(extract_matrix(matrix, last_split, len(matrix)))
        elif matrix[i][i+1] == 0:
            sub_matrices.append(extract_matrix(matrix, last_split, i+1))
            last_split = i+1
    return sub_matrices

def union_matrix(matrices):
    full_matrix = []
    offset = 0
    for matrix in matrices:
        zerofill_matrix(full_matrix, len(matrix), len(matrix))
        for i in range(len(matrix)):
            for j in range(len(matrix)):
                full_matrix[offset+i][offset+j] = matrix[i][j]
            
        # Set new offset
        offset = len(matrix)
    
    return full_matrix


def create_matrix(size = 1000, max_jordan_size = 100, max_eigenvalue = 1000):
    cur = 0
    matrix = zero_matrix(size, size)
    while cur < size:
        # length of the first jordan part
        length = random.randint(1, max_jordan_size)
        eigenvalue = random.randint(-max_eigenvalue, max_eigenvalue)
        
        for i in range(length):
            if cur + i < size:
                matrix[cur + i][cur + i] = eigenvalue
        for i in range(length-1):
            if cur + i + 1 < size:
                matrix[cur + i][cur + i + 1] = 1
        
        cur += length
    
    return matrix


if __name__ == '__main__':
    matrix = create_matrix(10)
    submatrices = split_matrix(matrix)
    
    # Establish communication queues
    tasks = multiprocessing.Queue()
    results = multiprocessing.Queue()
    
    # Start processes
    num_processes = multiprocessing.cpu_count()
    processes = [Matrixmultiplier(tasks, results) for i in range(num_processes)]
    for w in processes:
        w.start()
    
    # Enqueue jobs
    counter = 0
    for submatrix in submatrices:
        tasks.put((counter, Task(submatrix)))
        counter += 1
    
    # Add a poison pill for each consumer
    for i in range(num_processes):
        tasks.put((None, None))
    
    # Start printing results
    num_jobs = len(submatrices)
    resultlist = []
    while num_jobs:
        result = results.get() # all partial results
        resultlist.append(result)
        num_jobs -= 1
    
    # Each result has its position at [0] of the tuple and the result at [1]
    sorted(resultlist)
    
    print("A = ")
    print(matrix)
    print("A^2 = ")
    print(union_matrix([result[1] for result in resultlist]))