import copy

import framework.transf_visitor as tvisitor
import framework.visitor

import lan
import lan.ast_buildingblock as ast_bb


class MyError(Exception):
    def __init__(self, value):
        self.value = value

    def __str__(self):
        return repr(self.value)


class Transformation():
    """ Apply transformations to the original AST. Includes:
    1. Local Memory
    2. Stencil Local Memory
    3. Placing data in registers
    4. Transposing arrays
    5. Unrolling loops
    6. Adding defines
    7. Setting the number of dimensions to parallelize
    8. Setting the local work-group size
    9. Setting if we should read data back from the GPU
    10. Setting which kernel arguments changes
    """

    def __init__(self, rw):
        # The rewriter
        self.rw = rw

    def SetParDim(self, number):
        rw = self.rw
        rw.ParDim = number

    def SetNoReadBack(self):
        rw = self.rw
        rw.NoReadBack = True

    def Unroll(self, looplist):
        rw = self.rw
        # find loops and check that the loops given in the argument
        # exist
        loopIndices = framework.visitor.LoopIndices()
        loopIndices.visit(rw.Kernel)
        kernelLoops = loopIndices.Loops
        for l in looplist:
            if l not in kernelLoops:
                print "Valid loops are %r. Your input contained %r. Aborting..." % (kernelLoops.keys(), l)
                return

        rw.UnrollLoops.extend(looplist)

    def Unroll2(self, looplist):
        rw = self.rw

        for n in looplist:
            outerloop = rw.Loops[n]
            outeridx = n
            compound = outerloop.compound
            outerloop.compound = lan.Compound([])
            innerloop = copy.deepcopy(outerloop)
            innerloop.compound = compound
            outerstats = outerloop.compound.statements
            # Add the TypeId declarations to the outer loop
            # First find the typeids
            typeIds = tvisitor.TypeIds2()
            typeIds.visit(innerloop.compound)
            for m in typeIds.trans_ids:
                outerstats.append(m)
            outerstats.append(innerloop)
            upperbound = str(looplist[n])
            if upperbound == '0':
                upperbound = outerloop.cond.rval.name
            # change increment of outer loop
            outerloop.inc = lan.Increment(lan.Id(outeridx), '+=' + upperbound)
            inneridx = outeridx * 2
            # For adding to this index in other subscripts
            rw.Add[outeridx] = inneridx

            # new inner loop
            innerloop.cond = lan.BinOp(lan.Id(inneridx), '<', lan.Id(upperbound))
            innerloop.inc = lan.Increment(lan.Id(inneridx), '++')
            innerloop.init = lan.ConstantAssignment(inneridx)
            rw.Loops[inneridx] = innerloop

            # In the inner loop: Add the extra index to the index of the
            # outer loop
            exchangeId = tvisitor.ExchangeIdWithBinOp({n: lan.BinOp(lan.Id(outeridx), '+', lan.Id(inneridx))})
            exchangeId.visit(innerloop.compound)

            # unroll the inner index in stringstream
            self.Unroll([inneridx])
            if looplist[n] == 0:
                self.Unroll([outeridx])

    def SetDefine(self, varList):
        # TODO: Check that the vars in varlist are actually an argument
        rw = self.rw

        rw.DefinesAreMade = True
        accname = 'str'
        sstream = lan.TypeId(['std::stringstream'], lan.Id(accname))
        stats = rw.Define.statements
        stats.append(sstream)

        # add the defines to the string stream 
        for var in varList:
            try:
                hstvar = rw.NameSwap[var]
            except KeyError:
                hstvar = var
            add = lan.Id(accname + ' << ' + '\"' + '-D' + var + '=\"' + ' << ' + hstvar + ' << \" \";')
            stats.append(add)

        # Set the string to the global variable
        lval = lan.Id('KernelDefines')
        stats.append(lan.Assignment(lval, lan.Id(accname + '.str()')))

        # Need to remove the kernel corresponding kernel arguments
        for var in varList:
            rw.KernelArgs.pop(var)

    def placeInReg2(self, arrDict):
        rw = self.rw
        stats = rw.Kernel.statements
        initstats = []
        loadings = []
        writes = []
        # Create the loadings
        for i, n in enumerate(arrDict):
            for m in arrDict[n]:
                idx = m
                sub = copy.deepcopy(rw.astrepr.LoopArrays[n][idx])
                type = rw.Type[n][0]
                regid = lan.Id(n + str(m) + '_reg')
                reg = lan.TypeId([type], regid)
                writes.append(regid)
                assign = lan.Assignment(reg, sub)
                initstats.append(assign)

        stats.insert(0, lan.GroupCompound(initstats))

        # Replace the global Arefs with the register Arefs
        count = 0
        for i, n in enumerate(arrDict):
            for m in arrDict[n]:
                idx = m
                aref_new = writes[count]
                aref_old = rw.astrepr.LoopArrays[n][idx]
                # Copying the internal data of the two arefs
                aref_old.name.name = aref_new.name
                aref_old.subscript = []
                count += 1

    def placeInReg3(self, arrDict, insideList):
        """ Check if the arrayref is inside a loop and use a static
            array for the allocation of the registers
        """
        rw = self.rw
        stats = rw.Kernel.statements
        initstats = []
        writes = []

        if not arrDict:
            return

        if not insideList:
            self.placeInReg2(arrDict)
            return

        insideloop = insideList[0]

        if insideloop == '':
            print "placeInReg3 only works when the ArrayRef is inside a loop"
            print arrDict
            return

        # Add allocation of registers to the initiation stage
        for n in arrDict:
            lval = lan.TypeId([rw.Type[n][0]], \
                              lan.Id(n + '_reg[' + str(rw.astrepr.UpperLimit[insideloop]) \
                                     + ']'))
            initstats.append(lval)

        # add the loop to the initiation stage
        loop = copy.deepcopy(rw.Loops[insideloop])
        loopstats = []
        # Exchange loop index
        loop.compound.statements = loopstats

        initstats.append(loop)

        # Create the loadings
        for i, n in enumerate(arrDict):
            for m in arrDict[n]:
                idx = m
                sub = copy.deepcopy(rw.astrepr.LoopArrays[n][idx])
                regid = lan.ArrayRef(lan.Id(n + '_reg'), [lan.Id(insideloop)])
                writes.append(regid)
                assign = lan.Assignment(regid, sub)
                loopstats.app
end(assign)

        stats.
insert(0, lan.GroupCompound(initstats))
        # Replace the global Arefs with the register Arefs
        count = 0
        for i, n in enumerate(arrDict):
            for m in arrDict[n]:
                idx = m
                aref_new = copy.deepcopy(writes[count])
                aref_old = rw.astrepr.LoopArrays[n][idx]
                # Copying the internal data of the two arefs
                aref_old.name.name = aref_new.name.name
                aref_old.subscript = aref_new.subscript
                count += 1

    def localMemory3(self, arrDict, loopDict=None, blockDict=None):
        rw = self.rw
        initstats = []
        initComp = lan.GroupCompound(initstats)
        rw.Kernel.statements.insert(0, initComp)

        if loopDict is None:
            loopDict = dict()
            # So we create it
            for n in arrDict:
                for i in arrDict[n]:
                    loopDict[(n, i)] = []

            for n in arrDict:
                for i in arrDict[n]:
                    subscript = rw.SubscriptNoId[n][i]
                    acc = []
                    for m in subscript:
                        try:
                            _ = int(m)
                        except:
                            if m not in rw.GridIndices:
                               