from hwt.hdlObjects.operator import Operator
from hwt.hdlObjects.assignment import Assignment
from hwt.hdlObjects.value import Value
from hwt.hdlObjects.statements import IfContainer, SwitchContainer
from hwt.pyUtils.arrayQuery import areSetsIntersets, groupedby
from hwt.serializer.utils import maxStmId
from itertools import islice, zip_longest


def removeUnconnectedSignals(netlist):
    """
    If signal is not driving anything remove it
    """

    toDelete = set()
    toSearch = netlist.signals

    while toSearch:
        _toSearch = set()

        for sig in toSearch:
            if not sig.endpoints:
                try:
                    if sig._interface is not None:
                        # skip interfaces before we want to check them,
                        # they should not be optimized out from design
                        continue
                except AttributeError:
                    pass

                for e in sig.drivers:
                    # drivers of this signal are useless rm them
                    if isinstance(e, Operator):
                        for op in e.operands:
                            if not isinstance(op, Value):
                                op.endpoints.remove(e)
                                _toSearch.add(op)

                    elif isinstance(e, Assignment):
                        op = e.src
                        if not isinstance(op, Value):
                            op.endpoints.remove(e)
                            _toSearch.add(op)

                        netlist.startsOfDataPaths.remove(e)
                    else:
                        assert False, ("Drivers should be only index operators or assignments", e)

                toDelete.add(sig)

        if toDelete:
            for sig in toDelete:
                netlist.signals.remove(sig)
                try:
                    _toSearch.remove(sig)
                except KeyError:
                    pass
            toDelete = set()
        toSearch = _toSearch


class IncompatibleStructure(Exception):
    """
    instances of HWProcess can not be merged due incompatible structure
    """


def isMergableStmList(listA, listB):
    la = len(listA)
    lb = len(listB)
    if la == lb:
        return True
    elif la > 0 and lb == 1:
        return True
    else:
        return False


def checkIfIsTooSimple(proc):
    """check if process is just unconditional assignments and it is useless to merge them"""
    try:
        a, = proc.statements
        if isinstance(a, Assignment):
            return True
    except ValueError:
        pass
    return False


def mergeStmLists(stmsA, stmsB):
    tmp = []
    for a, b in zip_longest(stmsA, stmsB, fillvalue=None):
        if b is None:
            assert a is not None
            tmp.append(a)
        else:
            tmp.extend(tryToMergeStm(a, b))

    return tmp


def tryToMergeStm(stmA, stmB):
    """
    :raise IncompatibleStructure: if it is not possible to merge statements
    :return: generator of statements
    """

    if isinstance(stmA, Assignment) or isinstance(stmB, Assignment):
        yield stmA
        yield stmB
        return

    aIsIf = isinstance(stmA, IfContainer)
    bIsIf = isinstance(stmB, IfContainer)
    if aIsIf and bIsIf:
        if (stmA.cond != stmB.cond or
                not isMergableStmList(stmA.ifTrue, stmB.ifTrue) or
                not isMergableStmList(stmA.elIfs, stmB.elIfs) or
                not isMergableStmList(stmA.ifFalse, stmB.ifFalse)):
            raise IncompatibleStructure()

        ifTrue = mergeStmLists(stmA.ifTrue, stmB.ifTrue)

        elIfs = []
        for (condA, elifA), (condB, elifB) in zip(stmA.elIfs, stmB.elIfs):
            if (condA != condB or not isMergableStmList(elifA, elifB)):
                raise IncompatibleStructure()
            elIfs.append((condA, mergeStmLists(elifA, elifB)))

        ifFalse = mergeStmLists(stmA.ifFalse, stmB.ifFalse)

        yield IfContainer(stmA.cond, ifTrue, ifFalse, elIfs)
        return

    aIsSwitch = isinstance(stmA, SwitchContainer)
    bIsSwitch = isinstance(stmB, SwitchContainer)
    if (aIsSwitch and bIsSwitch and
            stmA.switchOn is stmB.switchOn and
            len(stmA.cases) == len(stmB.cases) and
            isMergableStmList(stmA.default, stmB.default)):

        cases = []
        for (vA, caseA), (vB, caseB) in zip(stmA.cases, stmB.cases):
            if vA != vB or not isMergableStmList(caseA, caseB):
                raise IncompatibleStructure()
            cases.append((vA, mergeStmLists(caseA, caseB)))

        default = mergeStmLists(stmA.default, stmB.default)

        yield SwitchContainer(stmA.switchOn, cases, default)
        return

    raise NotImplementedError(stmA, stmB)


def tryToMerge(procA, procB):
    """
    Try merge procB into procA

    :raise IncompatibleStructure: if merge is not possible
    :attention: procA is now result if merge has succeed
    :return: procA which is now result of merge
    """
    if (checkIfIsTooSimple(procA) or
            checkIfIsTooSimple(procB) or
            areSetsIntersets(procA.outputs, procB.sensitivityList) or
            areSetsIntersets(procB.outputs, procA.sensitivityList) or
            len(procA.statements) != len(procB.statements)):
        raise IncompatibleStructure()

    statements = []
    for stmA, stmB in zip(procA.statements, procB.statements):
        statements.extend(tryToMergeStm(stmA, stmB))

    procA.statements = statements
    procA.outputs.update(procB.outputs)
    procA.inputs.update(procB.inputs)
    procA.sensitivityList.update(procB.sensitivityList)

    return procA


def reduceProcesses(processes, procRanks):
    """
    Try to merge processes as much is possible

    :param processes: list of processes instances
    :param procRanks: process ranks = how many assignments is probably in process
        used to minimize number of merge tries
    """
    # sort to make order of merging same deterministic
    processes.sort(key=lambda x: (x.name, maxStmId(x)))
    # now try to reduce processes with nearly same structure of statements into one
    # to minimize number of processes
    for _, procs in groupedby(processes, lambda p: procRanks[p]):
        for iA, pA in enumerate(procs):
            if pA is None:
                continue
            for iB, pB in enumerate(islice(procs, iA + 1, None)):
                if pB is None:
                    continue
                try:
                    pA = tryToMerge(pA, pB)
                except IncompatibleStructure:
                    continue
                procs[iA + 1 + iB] = None
                # procs[iA] = pA

        for p in procs:
            if p is not None:
                yield p
