from multiprocessing import Process
from multiprocessing import Lock
import os
import time
import random
import numpy as np
import SharedArray as sa
import fcntl
import errno



class Cluster:

    def __init__(self, numThread = None):
        if numThread == None:
           numThread = os.cpu_count()
        self.numThread = numThread
        
        self._initbarr(numThread)
        self.fd = os.open('FILE_LOCK', os.O_CREAT)
        self.mlock = Lock()

        self.sharedList = []  # a list of shared variables

        self.resources = {}
        self.resources["lock"] = self.mlock
        self.resources["id"] = -1
        
        

    def __enter__(self):
        return self

    
    def __exit__(self, exc_type, exc_value, traceback):
        self.terminate()



    # paras need to be a tuple
    def runProcesses(self, func, paras = ()):

        p_list = []
        for id in range(self.numThread):
            # id starts at zero, and ends at numThread - 1
            self.resources = self.resources.copy() # a shallow copy of itself
            self.resources['id'] = id
            p_list.append(Process(target=func, args=(self.resources, )+paras))
        
        for p in p_list:
            p.start()
        
        for p in p_list:
            p.join()
        
        

    def _initbarr(self, numThread):
        nt = sa.create("shm://numThread2048", 1, int)
        nt[0] = numThread
        sense = sa.create("shm://sense2048", 1, int)
        sense[0] = 1
        counter = sa.create("shm://counter2048", 1, int)
        counter[0] = 0


    # Sense Reversal barrier
    @classmethod
    def barrier(cls):
        fd = cls.filelock()
        sense = sa.attach("shm://sense2048")
        counter = sa.attach("shm://counter2048")
        loc_sense = sense[0]
        nt = sa.attach("numThread2048")
        numThread = nt[0]

        if counter[0] == numThread - 1:  # last thread finished
            # Now every thread has reached the barrier
            # Reset everything
            counter[0] = 0
            sense[0] = 1 - sense[0]
            cls.fileunlock(fd)
        else:
            counter[0] += 1
            cls.fileunlock(fd)
            while loc_sense == sense[0]:
                pass



    # Create a shared array initialized to zeros
    # this shared array will be automatically deleted once processes finish
    def createShared(self, name, shape, dataType = int):
        # Check name != 'id' and name != 'lock' and no repeated names
        try:
            if name in self.resources:
                self.deleteShared()
                raise KeyError("Name '" + str(name) + "' repetition:" +
                    "Please use a different name.")
        except KeyError as e:
            exit(str(e))
            
        sharedLoc = "shm://" + name
        sharedAry = sa.create(sharedLoc, shape, dataType)
        self.sharedList.append(sharedAry)
        self.resources[name] = sharedAry


        return sharedAry

    
    @classmethod
    def getShared(cls, name):
        return sa.attach(name)

    def deleteShared(self):
        for key, value in self.resources.items():
            if key != 'id' and key != 'lock':
                sa.delete(key)
                del value
        sa.delete("shm://counter2048")
        sa.delete("shm://sense2048")
        sa.delete("shm://numThread2048")




    def terminate(self):
        os.close(self.fd)
        os.unlink("./FILE_LOCK")
        self.deleteShared()
        # In the case of abrupt termination of program, "shm://spInd2048" may
        # not be deleted. So need to delete it here
        try:
            sa.delete("spInd2048")
        except (IOError, OSError): # FileNotFoundError
            pass



    # If length is not divisible by n, those extra terms
    # will be placed in the list of the last process
    # For example, given a length of 10, and numThread of 3
    # and after shuffling, ary = [1, 3, 4, 0, 2, 7, 9, 8, 6, 5]
    # Then, 
    # process with id 0 will get [1, 3, 4]
    # process with id 1 will get [0, 2, 7]
    # process with id 2 will get [9, 8, 6, 5]
    @classmethod
    def splitIndices(cls, length, id, random = True):
        n = sa.attach("numThread2048")[0]

        if id == 0:
            ary = sa.create("shm://spInd2048", length, int)
            ary[:] = np.arange(length)
            if random:
                np.random.shuffle(ary)

        cls.barrier()

        if id != 0:
            ary = sa.attach("spInd2048")


        # Copy the computed numpy array to a new iterable
        # then delete the shared array
        # return a slice of iterable to each process

        chunkSize = int(length / n)
        result = list(ary)

        cls.barrier()

        if id == 0:
            sa.delete("spInd2048")
            del ary

        if id == n - 1:
            return result[id*chunkSize:]
        else:
            return result[id*chunkSize:(id+1)*chunkSize]


    @classmethod
    def filelock(cls):
        fd = os.open('FILE_LOCK', os.O_RDONLY)
        while True:
            try:
                # Acquire the lock
                fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
                break
            except (IOError, OSError) as e:
                # Raise if the error is not EAGAIN 
                # (i.e. Resource temporarily unavailable)
                if e.errno != errno.EAGAIN:
                    raise
            time.sleep(0.01) # Sleep to avoid busy waiting
        return fd
        

    @classmethod
    def fileunlock(cls, fd):
        # Release the lock
        fcntl.flock(fd, fcntl.LOCK_UN)
        os.close(fd)
     