import getopt, sys, math

def encode_graphlet(nodeSequence):
    encoding = ''
    nodeCodes = dict()
    nodeWith1 = None
    nodeWith2 = None
    cMax = 1
    i=0
    while i<len(nodeSequence):
        u = nodeSequence[i]
        v = nodeSequence[i+1]
        updated = False
        if not nodeCodes.has_key(u):
            nodeCodes[u] = cMax
            if cMax == 1:
                nodeWith1 = u
            if cMax == 2:
                nodeWith2 = u
            cMax += 1
            updated = True
        if not nodeCodes.has_key(v):
            nodeCodes[v] = cMax
            if cMax == 1:
                nodeWith1 = v
            if cMax == 2:
                nodeWith2 = v
            cMax += 1
            updated = True
        if cMax == 4 and updated:
            if (nodeCodes[u]==1 and nodeCodes[v]==3) or (nodeCodes[u]==3 and nodeCodes[v]==3):
                nodeCodes[nodeWith1] = 2
                nodeCodes[nodeWith2] = 1

        encoding += '%s%s' % (min(nodeCodes[u],nodeCodes[v]),max(nodeCodes[u],nodeCodes[v]))
        i+=2
    return encoding








def get_node_orbit_dependency_count(nodeOrbit):
    [g,v] = nodeOrbit.split('_')
    encounteredGraphlets = [g]
    o = 1 #count itself

    i=0
    while i<len(g):
        j=len(g)
        while j>i and v in g[i:j]:
            encoding = encode_graphlet(list(g[i:j]))
            if not encoding in encounteredGraphlets:
                encounteredGraphlets.append(encoding)
                o+=1
            j-=2
        i+=2


    return o


def generate_temporal_graphlets(kMax, nMax = None):
    print 'Generating all temporal graphlets with up to %s nodes and %s events' % (nMax if nMax is not None else 'MAX', kMax)

    temporalGraphlets = {
        1: [[1,2]]
    }
    i=1

    for k in range(2,kMax+1):

        temporalGraphlets[k] = []
        for g in temporalGraphlets[k-1]:
            s = g[-2]
            t = g[-1]
            u = max(g)+1

            nodes = set(g)
            if len(nodes)==2:
                temporalGraphlets[k].append(g + [s,t])
                if nMax is None or nMax > 2:
                    temporalGraphlets[k].append(g + [t,u])
            else:
                for v in nodes:
                    if v==s:
                       temporalGraphlets[k].append(g + [s,t])
                    elif v!=t:
                       temporalGraphlets[k].append(g + sorted([s,v]))
                       temporalGraphlets[k].append(g + sorted([t,v]))
                if nMax is None or nMax > max(g):
                    temporalGraphlets[k].append(g + [s,u])
                    temporalGraphlets[k].append(g + [t,u])
        print '\t%s temporal graphlet(s) with %s events' % (len(temporalGraphlets[k]),k)
    return temporalGraphlets


def get_node_orbits(temporalGraphlets):
     orbits = []
     for k in sorted(temporalGraphlets.keys()):
        tList = temporalGraphlets[k]
        tList = [(''.join([str(v) for v in  t]), max(t)) for t in tList]
        for tPair in sorted(tList): #sort encoding in a lexicographic order
            t = tPair[0]
            maxV = tPair[1]
            if maxV == 2:
                orbits.append(t + '_1')
            else:
                for v in range(1,maxV+1):
                    orbits.append(t + '_%s' % v)
     return orbits




if __name__ == "__main__":
    outputDir = '../vector_files/'
    nMax = None
    saveOrbits= False

    opts, args = getopt.getopt(sys.argv[1:], "k:o:n:vh", ["help"])
    for o, a in opts:
        if o == '-n':
            nMax = int(a)
        elif  o == '-v':
            saveOrbits = True
        elif o=='-o':
            outputDir = a + '/'
        elif o in ('-h','--help'):
            print """Usage:
-n nMax    Maximum number of nodes in a temporal graphlet (at least 2)
-v saveOrbits   Save node orbits with their dependency counts and weights
-o outputDir   Output directory ('../vector_files' by default)
-h,--help        Print this"""
            sys.exit(0)

    kMax = int(args[0])

    if kMax<1 or kMax>9:
        raise Exception('Number of events should be between 1 and 8')

    if nMax is None or nMax > kMax +1:
        nMax = kMax + 1
        print 'Setting nMax to %s' % nMax

    if nMax<2:
        raise Exception('Number of nodes should be at least 2')


    temporalGraphlets = generate_temporal_graphlets(kMax, nMax)

    if saveOrbits:
        orbits = get_node_orbits(temporalGraphlets)
        print '%s orbits' % len(orbits)
        with open(outputDir + 'orbits_%s_%s.txt' % (kMax,nMax), 'w') as f:
            for nodeOrbit in orbits:
                orbitDependencyCount = get_node_orbit_dependency_count(nodeOrbit)
                orbitWeight = 1 - math.log(orbitDependencyCount)/math.log(len(orbits))
                #f.write('%s\n' % nodeOrbit)
                f.write('%s\t%s\t%s\n' % (nodeOrbit, orbitDependencyCount, orbitWeight))

    #save graphlets
    with open(outputDir + 'graphlets_%s_%s.txt' % (kMax,nMax), 'w') as f:
        for k in sorted(temporalGraphlets.keys()):
            tList = temporalGraphlets[k]
            tList = [(''.join([str(v) for v in  t]), max(t)) for t in tList]
            for tPair in sorted(tList): #sort encoding in a lexicographic order
                t = tPair[0]
                f.write(t + '\n')





