''' --- TEMPORAL MEMORY FUNCTIONS --- '''
def cellsFromCol(col):
    global TM_cellsPerCol
    minCol = (col*TM_cellsPerCol)+1
    cells = [i+minCol for i in range(TM_cellsPerCol)]
    return cells

def colFromCell(cell):
    global TM_cellsPerCol
    col = int((float(cell) / TM_cellsPerCol) - 0.01)
    return col

def segmentsFromCell(cell,cells_segments):
    segs = cells_segments[cell]
    return segs

def segmentsFromCol(col,cells_segments):
    cells_Segs = {}
    cells = cellsFromCol(col)
    for c in cells:
        cells_Segs[c] = segmentsFromCell(c,cells_segments)
    return cells_Segs

def activatePredColCells(col,prevActiveCells,prevWinnerCells,colActiveCells,colMatchingCells,cells_segments):
    global TM_segActiveThresh,TM_permINC,TM_permDEC,TM_segNewSynCount,TM_learn
    activeCs = list()
    winnerCs = list()
    for cell in colMatchingCells:
        activeCs.append(cell)
        winnerCs.append(cell)
        if TM_learn:
            segList = cells_segments[cell]
            for seg in segList:
                segIndex = segList.index(seg)
                if len(activeSynapsesOnSegment(seg,prevActiveCells)) > TM_segActiveThresh:
                    for preSynCell,perm in seg.items():
                        if preSynCell in prevActiveCells:
                            if (perm + TM_permINC) > 1.0:
                                cells_segments[cell][segIndex][preSynCell] = 1.0
                            else:
                                cells_segments[cell][segIndex][preSynCell] += TM_permINC
                        else:
                            if (perm - TM_permDEC) < 0:
                                cells_segments[cell][segIndex][preSynCell] = 0
                            else:
                                cells_segments[cell][segIndex][preSynCell] -= TM_permDEC
                    key = str(cell) + '_' + str(segIndex)
                    newSynCount = TM_segNewSynCount - cellsSegInds_numActivePotentialSyns[key]
                    growSynapses(seg,cell,newSynCount,prevWinnerCells)
    return activeCs,winnerCs,cells_segments

def activeSynapsesOnSegment(segment,prevActiveCells):
    preSynCells = list(segment.keys())
    actives = [c for c in preSynCells if c in prevActiveCells]
    return actives

def growSynapses(segment,cell,newSynapseCount,prevWinnerCells):
    global TM_segNewSynCount,TM_permINIT
    candidates = copy.copy(prevWinnerCells)
    while len(candidates) > 0 and newSynapseCount > 0:
        presynapticCell = rd.choice(candidates)
        candidates.remove(presynapticCell)
        alreadyConnected = False
        for preSynCell,perm in segment.items():
            if preSynCell == presynapticCell:
                alreadyConnected = True
        if not alreadyConnected:
            createNewSynapse(segment,cell,presynapticCell,TM_permINIT)
            newSynapseCount -= 1

def createNewSynapse(segment,cell,presynapticCell,TM_permINIT):
    global cells_segments
    segDicts = cells_segments[cell]
    seg_index = segDicts.index(segment)
    cells_segments[cell][seg_index][presynapticCell] = TM_permINIT

def burstCol(col,prevActiveCells,prevWinnerCells,activeSegCells,matchingSegCells,cells_segments,cellsSegInds_numActivePotentialSyns):
    global TM_segNewSynCount
    activeCs = list()
    winnerCs = list()
    for cell in cellsFromCol(col):
        activeCs.append(cell)
    if len(matchingSegCells) > 0:
        learningSegment,winnerCell = bestMatchingSegmentAndCell(col,matchingSegCells,prevWinnerCells,cells_segments,cellsSegInds_numActivePotentialSyns)
    else:
        winnerCell = leastUsedCell(col,cells_segments)
        if TM_learn and len(prevWinnerCells) > 0:
            learningSegment,cells_segments = growNewSegment(winnerCell,prevWinnerCells,prevActiveCells,cells_segments)
    winnerCs.append(winnerCell)
    if TM_learn and len(prevWinnerCells) > 0:
        segIndex = cells_segments[winnerCell].index(learningSegment)
        for preSynCell,perm in learningSegment.items():
            if preSynCell in prevActiveCells:
                if (perm + TM_permINC) > 1.0:
                    cells_segments[winnerCell][segIndex][preSynCell] = 1.0
                else:
                    cells_segments[winnerCell][segIndex][preSynCell] += TM_permINC
            else:
                if (perm - TM_permDEC) < 0:
                    cells_segments[winnerCell][segIndex][preSynCell] = 0
                else:
                    cells_segments[winnerCell][segIndex][preSynCell] -= TM_permDEC
        key = str(cell) + '_' + str(segIndex)
        numActiveSynsOnSeg = 0
        if key in cellsSegInds_numActivePotentialSyns:
            numActiveSynsOnSeg = cellsSegInds_numActivePotentialSyns[key]
        newSynapseCount = TM_segNewSynCount - numActiveSynsOnSeg
        growSynapses(learningSegment,winnerCell,newSynapseCount,prevWinnerCells)
    return activeCs,winnerCs,cells_segments

def getActiveMatchingCellsFromCol(col,prevActiveCells,cells_segments):
    global TM_permConnectThresh
    cells = cellsFromCol(col)
    activeSegCells, matchingSegCells = [[],[]]
    for c in cells:
        if len(cells_segments[c]) > 0:
            for seg in cells_segments[c]:
                connectedSyns = []
                for preSynCell, perm in seg.items():
                    if perm >= TM_permConnectThresh:
                        connectedSyns.append(preSynCell)
                connected_prevActive_syns = [cell for cell in connectedSyns if cell in prevActiveCells]
                if len(connected_prevActive_syns) >= TM_segActiveThresh:
                    activeSegCells.append(c)
                if len(connected_prevActive_syns) >= TM_segMatchingThres:
                    matchingSegCells.append(c)
    return activeSegCells,matchingSegCells

def bestMatchingSegmentAndCell(col,matchingCells,prevActiveCells,cells_segments,cellsSegInds_numActivePotentialSyns):
    bestMatchSeg = None
    bestMatchCell = None
    bestScore = -1
    for cell in matchingCells:
        segDicts = cells_segments[cell]
        for seg in segDicts:
            segIndex = segDicts.index(seg)
            key = str(cell) + '_' + str(segIndex)
            numActiveSyns = 0
            if key in cellsSegInds_numActivePotentialSyns:
                numActiveSyns = cellsSegInds_numActivePotentialSyns[key]
            if numActiveSyns > bestScore:
                bestMatchSeg = seg
                bestMatchCell = cell
                bestScore = numActiveSyns
    return bestMatchSeg,bestMatchCell

def leastUsedCell(col,cells_segments):
    global TM_cellMaxSegs
    cells = cellsFromCol(col)
    leastSegCells = list()
    minSegs = TM_cellMaxSegs
    for c in cells:
        seg_count = len(cells_segments[c])
        minSegs = min(minSegs,seg_count)
    for c in cells:
        if len(cells_segments[c]) == minSegs:
            leastSegCells.append(c)
    return rd.choice(leastSegCells)

def growNewSegment(cell,prevCells,prevActives,cells_segments):
    global TM_segNewSynCount,TM_permINIT
    if len(prevCells) < len(prevActives):
        prevCells = prevActives
    preSynCells = rd.sample(prevCells,TM_segNewSynCount)
    seg = {}
    for c in preSynCells:
        seg[c] = TM_permINIT
    cells_segments[cell].append(seg)
    return seg,cells_segments

def decSeg(seg,cell,segIndex,prevActiveCells,cells_segments):
    for preSynCell,perm in seg.items():
        if preSynCell in prevActiveCells:
            currentPerm = cells_segments[cell][segIndex][preSynCell]
            if (currentPerm - TM_permDEC) < 0:
                cells_segments[cell][segIndex][preSynCell] = 0
            else:
                cells_segments[cell][segIndex][preSynCell] -= TM_permDEC
    return cells_segments

def decColSegments(col,prevActiveCells,matchingSegCells,cells_segments):
    global TM_segMatchingThres,TM_learn
    if TM_learn:
        for cell in matchingSegCells:
            segList = cells_segments[cell]
            for seg in segList:
                segIndex = segList.index(seg)
                segPreSynCells = list(seg.keys())
                numActivePotentialSyns = len([c for c in segPreSynCells if c in prevActiveCells])
                if numActivePotentialSyns > TM_segMatchingThres:
                    cells_segments = decSeg(seg,cell,segIndex,prevActiveCells,cells_segments)
    return cells_segments