# initialization
from bitarray import bitarray
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import random
from time import time

def init_sim(p_len):
    # processors data-structure - [ time, delay, memory, income-buffer, id ]
    p = map(lambda x: {'mem':{},'rx':[],'id':x}, list(range(p_len)))
    return p

def run_gossip(p_list, exec_gossip, turn_eval, params):
    p_len = len(p_list)
    msg_queue = [[] for _ in range(params['L'] + 1)]
    finish = False

    inactive_bitmap = bitarray('0' * p_len)
    online_error_bitmap = bitarray('0' * p_len)

    inactive_list = [0]
    while 0 in inactive_list:
        inactive_list = np.random.permutation(p_len)
        inactive_list = inactive_list[:int(round(params['inactive_per'] * p_len)//1)]

    inactive_list = inactive_list.tolist()

    for a in inactive_list:
        inactive_bitmap[a] = True

    of_ind_time = []
    for i in range(params['online_failures']):
        tmp = 2
        t = np.random.randint(1, params['total_time'])
        while tmp in inactive_list or tmp in map(lambda x: x[0], of_ind_time):
            tmp = np.random.randint(1, p_len)
        online_error_bitmap[tmp] = True
        of_ind_time.append((tmp,t))
    # print 'online failures - ', of_ind_time

    for t in range(params['total_time']):
        # online failures
        for of_i, of_t in of_ind_time:
            if of_t < t:
                inactive_list.append(of_i)
        of_ind_time = filter(lambda x: x[1] >= t, of_ind_time)

        #handles the messeging mechanism
        for msg in msg_queue.pop(0):
            msg['buffer-size'] = len(p_list[msg['dest'] % params['p']]['rx'])
            p_list[msg['dest'] % p_len]['rx'].append(msg)
        msg_queue.append([])

        turn_eval(p_list, msg_queue, t)

        #handles the execution mechanism
        for i in range(len(p_list)):
            if i not in inactive_list:
                _, msg = exec_gossip(p_list[i], params)
                if msg == 'finish':
                    finish = True
                    break
                if msg:
                    msg_queue[-1].append(msg)

        if finish:
            break

        #if t % (params['total_time'] // 10) == 0:
            # print 'i = ', t, '/', params['total_time']

    return (p_list, inactive_bitmap, online_error_bitmap, t, finish)

def run_sim(exec_gossip, params, sim_init = lambda x: x, turn_eval = lambda x,y,z: 1):
    # print '-- init sim --'
    p_list = init_sim(params['p'])
    p_list = sim_init(p_list, params)
    # print '-- run gossip sim --'
    p_list, inactive_bitmap, online_error_bitmap, t, finish = run_gossip(p_list, exec_gossip, turn_eval, params)
    # print '-- done --'
    return (p_list, inactive_bitmap, online_error_bitmap, t, finish)

def fcg(p, params):
    if not 'timeout' in p['mem'] and not 'type' in p['mem']: # if no message was received until now
        if len(p['rx']) > 0:
            msg = p['rx'].pop(0)
            if msg['type'] == 'ccg-gossip':
                p['mem']['timeout'] = msg['timeout']
                p['mem']['time']    = msg['time'] + params['L'] # need to fix time with latency?
                p['mem']['type']    = 'ccg-gossip'
                # p['mem']['data']    = msg['data']
            elif msg['type'] == 'ccg-correct':
                p['mem']['type']    = 'ccg-finish'
                p['mem']['timeout'] = 0
                p['mem']['time'] = 0
        return (p, None)

    p['mem']['time'] += 1

    if p['mem']['type'] == 'ccg-gossip':
        if p['mem']['time'] > p['mem']['timeout']:
            p['mem']['type'] = 'ccg-correct'
            p['mem']['send'] = [True, True]
            p['mem']['dist'] = [0, 0]
            p['mem']['received'] = [[],[]]


    if p['mem']['type'] == 'ccg-gossip':
        # handle received
        if len(p['rx']) > 0:
            msg = p['rx'].pop(0)
            if msg['type'] == 'ccg-gossip':
                pass
#                 p['mem']['time'] = min(p['mem']['time'], msg['time'] + latency)
            elif msg['type'] == 'ccg-correct':
                p['mem']['type']    = 'ccg-correct'
                p['mem']['send'] = [True, True]
                p['mem']['dist'] = [0, 0]
                p['mem']['received'] = [[],[]]
                p['mem']['received'][msg['dir']] = [msg['dist']]
#                 p['mem']['log'].append(msg)
            return (p, None)
        # handle send
        return (p, {'dest': random.randint(1 if params['fcg_without_first'] else 0, params['p']) , 'timeout': p['mem']['timeout'], 'time': p['mem']['time'], 'type': p['mem']['type']})

    if p['mem']['type'] == 'ccg-correct':
        # handle received
        if len(p['rx']) > 0:
            msg = p['rx'].pop(0)
            if msg['type'] == 'ccg-gossip':
                pass
            elif msg['type'] == 'ccg-correct':
                p['mem']['received'][msg['dir']].append(msg['dist'])
#                 p['mem']['log'].append(msg)
            return (p, None)

        p['mem']['send'][0] = len(p['mem']['received'][0]) < params['f'] or p['mem']['dist'][0] < max(p['mem']['received'][0])
        p['mem']['send'][1] = len(p['mem']['received'][1]) < params['f'] or p['mem']['dist'][1] < max(p['mem']['received'][1])
        if any(p['mem']['send']):
            # handle send
            if (all(p['mem']['send']) and p['mem']['dist'][0] < p['mem']['dist'][1]) or not p['mem']['send'][1]:
                p['mem']['dist'][0] += 1
                dest = (p['id'] - p['mem']['dist'][0]) % params['p']
                dest -= 1 if dest == 0 and params['fcg_without_first'] else 0
                msg = {'dest': dest, 'dist': p['mem']['dist'][0], 'dir': 1, 'type': p['mem']['type']}
#                 p['mem']['send-log'].append(msg)
                return (p, msg)
            else:
                p['mem']['dist'][1] += 1
                dest = (p['id'] +  p['mem']['dist'][1]) % params['p']
                dest += 1 if dest == 0 and params['fcg_without_first'] else 0
                msg = {'dest': dest, 'dist': p['mem']['dist'][1], 'dir': 0, 'type': p['mem']['type']}
#                 p['mem']['send-log'].append(msg)
                return (p, msg)
        else:
            p['mem']['type'] = 'ccg-finish'
    return (p, None)

def get_closest_alive(bitmap, index, p_len):
    li = index - 1
    ri = index + 1
    while not bitmap[li % p_len]:
        li = (li - 1) % p_len
    while not bitmap[ri % p_len]:
        ri = (ri + 1) % p_len
    return (li, ri)

def fast_barrier(p, params):
    if not 'master' in p['mem']:
        # initialize slave
        p['mem']['time'] = 0
        p['mem']['ll'] = dict()
        p['mem']['li'] = (p['id'] - 1) % params['p']
        p['mem']['ri'] = (p['id'] + 1) % params['p']
        p['mem']['time_to_start'] = random.randint(0 ,params['spread'])
        p['mem']['timeout'] = params['Tf'] + p['mem']['time'] + p['mem']['time_to_start']
        p['mem']['bitmap_timeout'] = params['Ti'] + p['mem']['time'] + p['mem']['time_to_start']
        p['mem']['master'] = 0
        p['mem']['type'] = 'ring-valid'
        p['mem']['step'] = 0
        p['mem']['last_bitmap'] = bitarray('1' * params['p'])
        p['mem']['bitmap'] = bitarray('0' * params['p'])
        p['mem']['bitmap'][p['id']] = True

    if p['mem']['master'] == 1 and len(p['rx']) > params['master_buffer_size']:
        p['rx'] = p['rx'][:params['master_buffer_size']]


    if p['mem']['type'] != 'ring-valid':
        return fcg(p, params)

    p['mem']['time'] += 1
    if p['mem']['master']:
        return fast_barrier_master(p, params)
#     if p['mem']['time'] > p['mem']['time_to_start']:
    return fast_barrier_slave(p, params)
    return (p, None)


def get_bitmap_from_ll(ll, p_len):
    if len(ll) < 1:
        return bitarray('0' * p_len)
    bitmap = ll[0]['bitmap']
    for b in map(lambda x: x['bitmap'], ll):
        bitmap |= b
    for i in range(len(ll)):
        for j in range(len(ll[i]['flags'])):
            if ll[i]['flags'][j]:
                bitmap[ll[i]['ind'][j] % p_len] =  False
    return bitmap

def is_valid_bitmap(bitmap):
    return True

def fast_barrier_master(p, params):
    if 'gossip' in p['mem']:
        return (p, {'dest': random.randint(1, params['p']) , 'timeout': p['mem']['timeout'], 'time': p['mem']['time'], 'type': 'ccg-gossip'})

    if 'last_update' not in p['mem']:
        p['mem']['last_update'] = 0

    if len(p['rx']) > 0:
        msg = p['rx'].pop(0)
        if msg['type'] in ['error']:
            p['mem']['last_received'].append((p['mem']['time'], msg))
        return (p, [])

    p['mem']['last_received'] = filter(lambda x: x[0] > p['mem']['time'] - params['Tf'], p['mem']['last_received'])
    p['mem']['bitmap2'] = p['mem']['bitmap1']
    p['mem']['bitmap1'] = get_bitmap_from_ll(map(lambda x: x[1], p['mem']['last_received']), params['p'])

    if p['mem']['bitmap1'] != p['mem']['bitmap2']:
        p['mem']['last_update'] = p['mem']['time']

    #if p['mem']['time'] - p['mem']['last_update']  > Tf and \
    if  len(p['mem']['last_received']) < 2 * (params['f'] + 1) and \
        len(p['mem']['last_received']) >= 2: # p['mem']['bitmap1'].count(True) > p_len // 2:
        p['mem']['timeout'] = p['mem']['time'] + params['Tg']
        p['mem']['gossip'] = True
        # print 'valid!', p['mem']['time']
        if not params['barrier_fcg']:
            return (p, 'finish')

    return (p, None)

def fast_barrier_slave(p, params):
    if p['mem']['time'] < p['mem']['bitmap_timeout']:
        p['mem']['li'], p['mem']['ri'] = get_closest_alive(p['mem']['last_bitmap'], p['id'], params['p'])
    else:
        p['mem']['li'], p['mem']['ri'] = get_closest_alive(p['mem']['bitmap'], p['id'], params['p'])


    if p['mem']['step'] == 0:
        if p['mem']['time'] < p['mem']['time_to_start']:
            if len(p['rx']) > 0:
                msg = p['rx'].pop(0)
                if msg['type'] == 'ring-valid':
                    p['mem']['ll'][msg['id']] = p['mem']['time']
                    p['mem']['bitmap'] |= msg['bitmap']
            return (p, [])
        else:
            p['mem']['step'] += 1

    # send alive messages
    if p['mem']['step'] == 1:
        p['mem']['step'] += 1
        return (p, {'dest': p['mem']['li'], 'id' : p['id'], 'type': 'ring-valid', 'bitmap': p['mem']['bitmap']})
    if p['mem']['step'] == 2:
        p['mem']['step'] += 1
        p['mem']['timeout'] = params['Tf'] + p['mem']['time']
        return (p, {'dest': p['mem']['ri'], 'id' : p['id'], 'type': 'ring-valid', 'bitmap': p['mem']['bitmap']})

    # wait for alive messages
    if p['mem']['step'] == 3:
        if p['mem']['time'] < p['mem']['timeout']:
            if len(p['rx']) > 0:
                msg = p['rx'].pop(0)
                if msg['type'] == 'ring-valid':
                    p['mem']['ll'][msg['id']] = p['mem']['time']
                    p['mem']['bitmap'] |= msg['bitmap']
                elif msg['type'] == 'ccg-gossip':
                    p['mem']['timeout'] = msg['timeout']
                    p['mem']['time']    = msg['time'] + params['L'] + 1 # need to fix time with latency?
                    p['mem']['type']    = 'ccg-gossip'
                    # p['mem']['data']    = msg['data']
                elif msg['type'] == 'ccg-correct':
                    p['mem']['type']    = 'ccg-finish'
                return (p, None)
            else: # len(p['rx']) == 0:
                return (p, {'dest': random.randint(1, params['p']), 'id' : p['id'], 'type': 'ring-valid', 'bitmap': p['mem']['bitmap']})
        else:
            p['mem']['step'] += 1

    # check for non-responsive nodes
    if p['mem']['step'] == 4:
        p['mem']['timeout'] = params['Tf'] + p['mem']['time']
        for k in p['mem']['ll'].keys():
            if p['mem']['time'] - p['mem']['ll'][k] > params['Tf'] * 2:
                del p['mem']['ll'][k]
        lf = True if not p['mem']['li'] in p['mem']['ll'] else False
        rf = True if not p['mem']['ri'] in p['mem']['ll'] else False
        flags = [lf, rf]
        p['mem']['step'] = 1
        if any(flags):
            # send error message to the master
            cur_bitmap = p['mem']['last_bitmap'] if p['mem']['time'] < p['mem']['bitmap_timeout'] else p['mem']['bitmap']
            return (p, {'dest':0, 'bitmap': cur_bitmap, 'ind' : [p['mem']['li'], p['mem']['ri']], 'flags': flags, 'id': p['id'], 'type': 'error'})
    return (p, None)

def sim_init(p_list, params):
    p_list[0]['mem']['master'] = 1
    p_list[0]['mem']['type'] = 'ring-valid'
    p_list[0]['mem']['time'] = 0
    p_list[0]['mem']['timeout'] = params['Tf']
    p_list[0]['mem']['last_received'] = []
    p_list[0]['mem']['bitmap1'] = bitarray('0' * params['p'])
    p_list[0]['mem']['bitmap2'] = bitarray('0' * params['p'])
    return p_list

def run_one(params):
    # t0 = time()
    p_list, inactive_bitmap, online_error_bitmap, t, finish = run_sim(fast_barrier, params, sim_init=sim_init)
    # print 'time - ', (time()-t0), 'seconds'

    bitmap = p_list[0]['mem']['bitmap1']
    bitmap[0] = True
    correct = (bitmap ^ (~inactive_bitmap)) & (~online_error_bitmap)
    return (finish, correct.count(), t)