"""Flexible enumeration of C types."""
from __future__ import division, print_function

from Enumeration import *

# TODO:

#  - struct improvements (flexible arrays, packed &
#    unpacked, alignment)
#  - objective-c qualified id
#  - anonymous / transparent unions
#  - VLAs
#  - block types
#  - K&R functions
#  - pass arguments of different types (test extension, transparent union)
#  - varargs

###
# Actual type types

class Type(object):
    def isBitField(self):
        return False

    def isPaddingBitField(self):
        return False

    def getTypeName(self, printer):
        name = 'T%d' % len(printer.types)
        typedef = self.getTypedefDef(name, printer)
        printer.addDeclaration(typedef)
        return name

class BuiltinType(Type):
    def __init__(self, name, size, bitFieldSize=None):
        self.name = name
        self.size = size
        self.bitFieldSize = bitFieldSize

    def isBitField(self):
        return self.bitFieldSize is not None

    def isPaddingBitField(self):
        return self.bitFieldSize is 0

    def getBitFieldSize(self):
        assert self.isBitField()
        return self.bitFieldSize

    def getTypeName(self, printer):
        return self.name

    def sizeof(self):
        return self.size

    def __str__(self):
        return self.name

class EnumType(Type):
    unique_id = 0

    def __init__(self, index, enumerators):
        self.index = index
        self.enumerators = enumerators
        self.unique_id = self.__class__.unique_id
        self.__class__.unique_id += 1

    def getEnumerators(self):
        result = ''
        for i, init in enumerate(self.enumerators):
            if i > 0:
                result = result + ', '
            result = result + 'enum%dval%d_%d' % (self.index, i, self.unique_id)
            if init:
                result = result + ' = %s' % (init)

        return result

    def __str__(self):
        return 'enum { %s }' % (self.getEnumerators())

    def getTypedefDef(self, name, printer):
        return 'typedef enum %s { %s } %s;'%(name, self.getEnumerators(), name)

class RecordType(Type):
    def __init__(self, index, isUnion, fields):
        self.index = index
        self.isUnion = isUnion
        self.fields = fields
        self.name = None

    def __str__(self):
        def getField(t):
            if t.isBitField():
                return "%s : %d;" % (t, t.getBitFieldSize())
            else:
                return "%s;" % t

        return '%s { %s }'%(('struct','union')[self.isUnion],
                            ' '.join(map(getField, self.fields)))

    def getTypedefDef(self, name, printer):
        def getField(it):
            i, t = it
            if t.isBitField():
                if t.isPaddingBitField():
                    return '%s : 0;'%(printer.getTypeName(t),)
                else:
                    return '%s field%d : %d;'%(printer.getTypeName(t),i,
                                               t.getBitFieldSize())
            else:
                return '%s field%d;'%(printer.getTypeName(t),i)
        fields = [getField(f) for f in enumerate(self.fields)]
        # Name the struct for more readable LLVM IR.
        return 'typedef %s %s { %s } %s;'%(('struct','union')[self.isUnion],
                                           name, ' '.join(fields), name)
                                           
class ArrayType(Type):
    def __init__(self, index, isVector, elementType, size):
        if isVector:
            # Note that for vectors, this is the size in bytes.
            assert size > 0
        else:
            assert size is None or size >= 0
        self.index = index
        self.isVector = isVector
        self.elementType = elementType
        self.size = size
        if isVector:
            eltSize = self.elementType.sizeof()
            assert not (self.size % eltSize)
            self.numElements = self.size // eltSize
        else:
            self.numElements = self.size

    def __str__(self):
        if self.isVector:
            return 'vector (%s)[%d]'%(self.elementType,self.size)
        elif self.size is not None:
            return '(%s)[%d]'%(self.elementType,self.size)
        else:
            return '(%s)[]'%(self.elementType,)

    def getTypedefDef(self, name, printer):
        elementName = printer.getTypeName(self.elementType)
        if self.isVector:
            return 'typedef %s %s __attribute__ ((vector_size (%d)));'%(elementName,
                                                                        name,
                                                                        self.size)
        else:
            if self.size is None:
                sizeStr = ''
            else:
                sizeStr = str(self.size)
            return 'typedef %s %s[%s];'%(elementName, name, sizeStr)

class ComplexType(Type):
    def __init__(self, index, elementType):
        self.index = index
        self.elementType = elementType

    def __str__(self):
        return '_Complex (%s)'%(self.elementType)

    def getTypedefDef(self, name, printer):
        return 'typedef _Complex %s %s;'%(printer.getTypeName(self.elementType), name)

class FunctionType(Type):
    def __init__(self, index, returnType, argTypes):
        self.index = index
        self.returnType = returnType
        self.argTypes = argTypes

    def __str__(self):
        if self.returnType is None:
            rt = 'void'
        else:
            rt = str(self.returnType)
        if not self.argTypes:
            at = 'void'
        else:
            at = ', '.join(map(str, self.argTypes))
        return '%s (*)(%s)'%(rt, at)

    def getTypedefDef(self, name, printer):
        if self.returnType is None:
            rt = 'void'
        else:
            rt = str(self.returnType)
        if not self.argTypes:
            at = 'void'
        else:
            at = ', '.join(map(str, self.argTypes))
        return 'typedef %s (*%s)(%s);'%(rt, name, at)

###
# Type enumerators

class TypeGenerator(object):
    def __init__(self):
        self.cache = {}

    def setCardinality(self):
        abstract

    def get(self, N):
        T = self.cache.get(N)
        if T is None:
            assert 0 <= N < self.cardinality
            T = self.cache[N] = self.generateType(N)
        return T

    def generateType(self, N):
        abstract

class FixedTypeGenerator(TypeGenerator):
    def __init__(self, types):
        TypeGenerator.__init__(self)
        self.types = types
        self.setCardinality()

    def setCardinality(self):
        self.cardinality = len(self.types)

    def generateType(self, N):
        return self.types[N]

# Factorial
def fact(n):
    result = 1
    while n > 0:
        result = result * n
        n = n - 1
    return result

# Compute the number of combinations (n choose k)
def num_combinations(n, k): 
    return fact(n) // (fact(k) * fact(n - k))

# Enumerate the combinations choosing k elements from the list of values
def combinations(values, k):
    # From ActiveState Recipe 190465: Generator for permutations,
    # combinations, selections of a sequence
    if k==0: yield []
    else:
        for i in range(len(values)-k+1):
            for cc in combinations(values[i+1:],k-1):
                yield [values[i]]+cc

class EnumTypeGenerator(TypeGenerator):
    def __init__(self, values, minEnumerators, maxEnumerators):
        TypeGenerator.__init__(self)
        self.values = values
        self.minEnumerators = minEnumerators
        self.maxEnumerators = maxEnumerators
        self.setCardinality()

    def setCardinality(self):
        self.cardinality = 0
        for num in range(self.minEnumerators, self.maxEnumerators + 1):
            self.cardinality += num_combinations(len(self.values), num)

    def generateType(self, n):
        # Figure out the number of enumerators in this type
        numEnumerators = self.minEnumerators
        valuesCovered = 0
        while numEnumerators < self.maxEnumerators:
            comb = num_combinations(len(self.values), numEnumerators)
            if valuesCovered + comb > n:
                break
            numEnumerators = numEnumerators + 1
            valuesCovered += comb

        # Find the requested combination of enumerators and build a
        # type from it.
        i = 0
        for enumerators in combinations(self.values, numEnumerators):
            if i == n - valuesCovered:
                return EnumType(n, enumerators)
                
            i = i + 1

        assert False

class ComplexTypeGenerator(TypeGenerator):
    def __init__(self, typeGen):
        TypeGenerator.__init__(self)
        self.typeGen = typeGen
        self.setCardinality()
    
    def setCardinality(self):
        self.cardinality = self.typeGen.cardinality

    def generateType(self, N):
        return ComplexType(N, self.typeGen.get(N))

class VectorTypeGenerator(TypeGenerator):
    def __init__(self, typeGen, sizes):
        TypeGenerator.__init__(self)
        self.typeGen = typeGen
        self.sizes = tuple(map(int,sizes))
        self.setCardinality()

    def setCardinality(self):
        self.cardinality = len(self.sizes)*self.typeGen.cardinality

    def generateType(self, N):
        S,T = getNthPairBounded(N, len(self.sizes), self.typeGen.cardinality)
        return ArrayType(N, True, self.typeGen.get(T), self.sizes[S])

class FixedArrayTypeGenerator(TypeGenerator):
    def __init__(self, typeGen, sizes):
        TypeGenerator.__init__(self)
        self.typeGen = typeGen
        self.sizes = tuple(size)
        self.setCardinality()

    def setCardinality(self):
        self.cardinality = len(self.sizes)*self.typeGen.cardinality

    def generateType(self, N):
        S,T = getNthPairBounded(N, len(self.sizes), self.typeGen.cardinality)
        return ArrayType(N, false, self.typeGen.get(T), self.sizes[S])

class ArrayTypeGenerator(TypeGenerator):
    def __init__(self, typeGen, maxSize, useIncomplete=False, useZero=False):
        TypeGenerator.__init__(self)
        self.typeGen = typeGen
        self.useIncomplete = useIncomplete
        self.useZero = useZero
        self.maxSize = int(maxSize)
        self.W = useIncomplete + useZero + self.maxSize
        self.setCardinality()

    def setCardinality(self):
        self.cardinality = self.W * self.typeGen.cardinality

    def generateType(self, N):
        S,T = getNthPairBounded(N, self.W, self.typeGen.cardinality)
        if self.useIncomplete:
            if S==0:
                size = None
                S = None
            else:
                S = S - 1
        if S is not None:
            if self.useZero:
                size = S
            else:
                size = S + 1        
        return ArrayType(N, False, self.typeGen.get(T), size)

class RecordTypeGenerator(TypeGenerator):
    def __init__(self, typeGen, useUnion, maxSize):
        TypeGenerator.__init__(self)
        self.typeGen = typeGen
        self.useUnion = bool(useUnion)
        self.maxSize = int(maxSize)
        self.setCardinality()

    def setCardinality(self):
        M = 1 + self.useUnion
        if self.maxSize is aleph0:
            S =  aleph0 * self.typeGen.cardinality
        else:
            S = 0
            for i in range(self.maxSize+1):
                S += M * (self.typeGen.cardinality ** i)
        self.cardinality = S

    def generateType(self, N):
        isUnion,I = False,N
        if self.useUnion:
            isUnion,I = (I&1),I>>1
        fields = [self.typeGen.get(f) for f in getNthTuple(I,self.maxSize,self.typeGen.cardinality)]
        return RecordType(N, isUnion, fields)

class FunctionTypeGenerator(TypeGenerator):
    def __init__(self, typeGen, useReturn, maxSize):
        TypeGenerator.__init__(self)
        self.typeGen = typeGen
        self.useReturn = useReturn
        self.maxSize = maxSize
        self.setCardinality()
    
    def setCardinality(self):
        if self.maxSize is aleph0:
            S = aleph0 * self.typeGen.cardinality()
        elif self.useReturn:
            S = 0
            for i in range(1,self.maxSize+1+1):
                S += self.typeGen.cardinality ** i
        else:
            S = 0
            for i in range(self.maxSize+1):
                S += self.typeGen.cardinality ** i
        self.cardinality = S
    
    def generateType(self, N):
        if self.useReturn:
            # Skip the empty tuple
            argIndices = getNthTuple(N+1, self.maxSize+1, self.typeGen.cardinality)
            retIndex,argIndices = argIndices[0],argIndices[1:]
            retTy = self.typeGen.get(retIndex)
        else:
            retTy = None
            argIndices = getNthTuple(N, self.maxSize, self.typeGen.cardinality)
        args = [self.typeGen.get(i) for i in argIndices]
        return FunctionType(N, retTy, args)

class AnyTypeGenerator(TypeGenerator):
    def __init__(self):
        TypeGenerator.__init__(self)
        self.generators = []
        self.bounds = []
        self.setCardinality()
        self._cardinality = None
        
    def getCardinality(self):
        if self._cardinality is None:
            return aleph0
        else:
            return self._cardinality
    def setCardinality(self):
        self.bounds = [g.cardinality for g in self.generators]
        self._cardinality = sum(self.bounds)
    cardinality = property(getCardinality, None)

    def addGenerator(self, g):
        self.generators.append(g)
        for i in range(100):
            prev = self._cardinality
            self._cardinality = None
            for g in self.generators:
                g.setCardinality()
            self.setCardinality()
            if (self._cardinality is aleph0) or prev==self._cardinality:
                break
        else:
            raise RuntimeError("Infinite loop in setting cardinality")

    def generateType(self, N):
        index,M = getNthPairVariableBounds(N, self.bounds)
        return self.generators[index].get(M)

def test():
    fbtg = FixedTypeGenerator([BuiltinType('char', 4),
                               BuiltinType('char', 4, 0),
                               BuiltinType('int',  4, 5)])

    fields1 = AnyTypeGenerator()
    fields1.addGenerator( fbtg )

    fields0 = AnyTypeGenerator()
    fields0.addGenerator( fbtg )
#    fields0.addGenerator( RecordTypeGenerator(fields1, False, 4) )

    btg = FixedTypeGenerator([BuiltinType('char', 4),
                              BuiltinType('int',  4)])
    etg = EnumTypeGenerator([None, '-1', '1', '1u'], 0, 3)

    atg = AnyTypeGenerator()
    atg.addGenerator( btg )
    atg.addGenerator( RecordTypeGenerator(fields0, False, 4) )
    atg.addGenerator( etg )
    print('Cardinality:',atg.cardinality)
    for i in range(100):
        if i == atg.cardinality:
            try:
                atg.get(i)
                raise RuntimeError("Cardinality was wrong")
            except AssertionError:
                break
        print('%4d: %s'%(i, atg.get(i)))

if __name__ == '__main__':
    test()