###########################################
# Algorithm Defined By Templates and Maps #
###########################################

from copy import deepcopy
from typing import Any, Iterable
import multiplied as mp

# -- TODO: sanity checks --------------------------------------------
#
# - Use setattr to block changes to self.matrix if state != 0, suggest self.reset().
#   Actually, this also applies to all attributes
class Algorithm():
    """
    Manages and sequences operations via a series of stages defined by templates and maps.

    parameters
    ----------
    matrix : mp.Matrix
        The matrix to be multiplied.

    ...
    """

    def __init__(self, bits: int,*, matrix: Any=None, saturation: bool=False, dadda=False) -> None:

        mp.validate_bitwidth(bits)
        if not isinstance(dadda, bool):
            raise TypeError(f"Expected dadda: bool, got {type(dadda)}")
        if not isinstance(saturation, bool):
            raise TypeError(f"Expected saturation: bool, got {type(saturation)}")
        if matrix is not None:
            if not isinstance(matrix, mp.Matrix):
                raise TypeError(f"Expected Matrix, got {type(matrix)}")
            self.matrix = matrix
        else:
            self.matrix = mp.Matrix(bits)

        self.bits       = bits
        self.dadda      = dadda
        self.state      = 0
        self.algorithm  = {}
        if self.dadda:
            hoist(self.matrix)
        self.saturation = saturation
        if self.saturation:
            self.__clamp_bitwidth()

        # -- TODO: update this when anything is modified ------------
        # create update() function
        # add to each modifying class method
        return None



    def push(self, source: mp.Template | mp.Pattern, map_: Any=None, dadda=False
    ) -> None:
        """
        Populates stage of an algorithm based on template. Generates pseudo
        result to represent output matrix

        >>> self.algorithm[x] = {
        >>>     "template" : mp.Template,
        >>>     "pseudo"   : mp.Matrix,
        >>>     "map"      : mp.Map}
        """

        if source.bits != self.bits:
            raise ValueError("Template bitwidth must match Algorithm bitwidth.")
        if map_ and not(isinstance(map_, (mp.Map))):
            raise TypeError("Invalid argument type. Expected mp.Map")

        # -- [TODO] ------------------------------------------------- #
        if map_ and not map_.rmap:                                    #
            raise NotImplementedError("Complex map not implemented")  #
        # ----------------------------------------------------------- #

        if isinstance(source, mp.Pattern):
            template = mp.Template(source)
        elif isinstance(source, mp.Template):
            template = source
        else:
            raise TypeError("Invalid argument type. Expected mp.Template")


        if not isinstance(template.result, list):
            raise ValueError("Template result is unset")


        result      = mp.Matrix(template.result)
        stage_index = len(self.algorithm)
        if not map_ and result:
            if dadda:
                res_copy = deepcopy(result)
                map_ = hoist(res_copy)
            else:
                map_ = result.resolve_rmap()
            result.apply_map(map_)
        else:
            result.apply_map(map_)

        stage = {
            'template': template,
            'pseudo': result,
            'map': map_,
        }
        self.algorithm[stage_index] = stage
        return None

    def __clamp_bitwidth(self) -> bool:
        """
        Saturates matrix if current matrix has carried past original bitwidth
        """
        boundary = (2**self.bits)-1
        as_int   = mp.to_int_matrix(self.matrix.matrix)
        test     = [boundary < i for i in as_int]

        if any(test):
            # flood bits within boundary
            saturated_value = [['0']*self.bits + ['1']*self.bits]
            self.matrix = mp.Matrix(saturated_value + mp.empty_matrix(self.bits)[1:])
            return True
        else:
            return False

    # ! Matrix.x_checksum is only useful in the context of Algorithm.__reduce()
    # - Maybe use bounds to create x_checksum within __reduce()'s unit collection
    # - OR within, the same scope, use bounds to execute a given arithmetic unit
    # ---------------------------------------------------------------
    # Mangled as execution order is sensitive and __reduce should only
    # be called by the algorithm itself via: self.step(), or self.exec()
    def __reduce(self) -> None:
        """
        use template or pattern to reduce a given matrix.
        """
        from copy import copy

        # -- implementation -----------------------------------------
        #
        # Arithmetic units are defined by how many rows they cover, or
        # their 'run', the length of which determines the type of unit:
        #
        # run = 3:
        #   CSA: carry is placed to the left of source column
        #   and one row down to avoid corrupting adjacent columns.
        #
        #   [input-------] | [output------]
        #   ...00100010... | ...00100010...
        #   ...00101010... | ..00101010....
        #   ...00101010... | ...________...
        #
        # run = 2:
        #   binary adder: carry generates through propagates, with a
        #   final carry extending past original width.
        #
        #   [input-------] | [output------]
        #   ...00110110... | ..001100000...
        #   ...00101010... | ...________...

        # -- isolate units -----------------------------------------
        # ! Implement result bounding box to create single point of truth ! #
        # - currently every stage of every matrix reduction needs to resolve
        #   conflicts dynamically before merging vs doing so once via the
        #   resultant template
        bounds: dict = self.algorithm[self.state]['template'].bounds
        # -- reduce -------------------------------------------------
        n         = self.bits << 1
        results   = {}
        chars = list(bounds.keys())
        chars.remove('_')
        for ch in  chars:
            base_index = bounds[ch][0][1]
            match bounds[ch][-1][1] - bounds[ch][0][1] +1: # row height
                case 1: # NOOP
                    output = [copy(self.matrix[base_index][0])]

                case 2: # ADD
                    #TODO: make use of checksums or use bounds

                    operand_a = copy(self.matrix[base_index][0])
                    operand_b = copy(self.matrix[base_index+1][0])
                    checksum  = [False] * n

                    # -- skip empty rows ----------------------------
                    start = 0
                    while operand_a[start] == '_' and operand_b[start] == '_':
                        start += 1


                    for i in range(start, n):
                        # -- row checksum ---------------------------
                        if operand_a[i] != '_' or operand_b[i] != '_':
                            checksum[i] = True

                        # -- normalise binary -----------------------
                        if operand_a[i] == '_' and operand_b[i] != '_':
                            operand_a[i] = '0'

                        elif operand_b[i] == '_' and  operand_a[i] != '_':
                            operand_b[i] = '0'

                        elif operand_a[i] == '_' and operand_b[i] == '_':
                            operand_a[i] = '0'
                            operand_b[i] = '0'

                    # -- binary addition ----------------------------
                    bits_ = sum(checksum)
                    cout  = 1 if not checksum[0] else 0
                    int_a = int("".join(operand_a[start:start+bits_]), 2)
                    int_b = int("".join(operand_b[start:start+bits_]), 2)

                    output     = [['_']*(start-cout)]
                    output[0] += list(f"{int_a+int_b:0{bits_+cout}b}")
                    output[0] += ['_']*(n-bits_-start)


                case 3: # CSA
                    #TODO: make use of checksums or use bounds
                    operand_a = copy(self.matrix[base_index][0])
                    operand_b = copy(self.matrix[base_index+1][0])
                    operand_c = copy(self.matrix[base_index+2][0])
                    empty     = False
                    output    = [['_']*n, ['_']*n]
                    start     = 0

                    # -- skip empty rows ----------------------------
                    while (
                        operand_a[start] == '_' and
                        operand_b[start] == '_' and
                        operand_c[start] == '_'
                    ):
                        start += 1



                    # -- sum columns -------------------------------
                    for i in range(start, n):
                        csa_sum = 0
                        csa_sum += 1 if operand_a[i] == '1' else 0
                        csa_sum += 1 if operand_b[i] == '1' else 0
                        csa_sum += 1 if operand_c[i] == '1' else 0

                        output[0][i] = csa_sum # uses index to store sum in-place
                        empty  = (operand_c[i] == '_') + (operand_b[i] == '_') + (operand_a[i] == '_')
                        # -- check end of unit ----------------------
                        if empty == 3:
                            output[0][i] =  '_' # set as empty since sum unused
                            break
                        if empty == 2:
                            output[0][i]   = '1' if csa_sum & 1 else '0'
                            continue

                        # -- brute force index 0 --------------------
                        # ! remove try catch:
                        # ! handle index zero outside of loop
                        # ! unit may not even overlap index 0
                        try:
                            output[0][i]   = '1' if csa_sum & 1 else '0'
                            output[1][i-1] = '1' if csa_sum & 2 else '0'
                        except IndexError:
                            continue
                case _:
                    raise ValueError(f"Unsupported unit type, len={bounds[ch][-1][1] - bounds[ch][0][1]}")

            # -- build unit into matrix -----------------------------
            unit_result = [[]]*self.bits
            i = 0
            while i < base_index:
                unit_result[i] = ['_']*n
                i += 1
            for row in output:
                unit_result[i] = row
                i += 1
            while i < self.bits:
                unit_result[i] = ['_']*n
                i += 1
            results[ch] = mp.Matrix(unit_result)


        # -- merge units to matrix ----------------------------------
        # Merge in any order, checking for overlaps between borders
        # resolve conflicts by summing present bit positions and shifting
        # a target unit's bit

        # ! difficult sanity checks --------------------------------- ! #
        # Complex scenarios, where NOOP, CSA and ADD units intersect
        # will require extensive checks:
        #
        #   [example--] || [region--]
        #   ...BbaAa... || ....ba....
        #   ...CcaAa... || ....ca....
        #   ...CcaAa... || ....ca....
        #
        #
        # If repeated along the same y-axis, resolution can get very tricky.
        # One method could be skipping merges and opting for merges with non
        # conflicting units, repeat until few partially merged templates remain
        # and finally resolve conflicts, if possible.
        #
        # These conflicts can quickly be found by checking the sum of possible
        # bit positions in a given region of intersecting arithmetic units.
        # The example's sum for the first column is == 3.  This should raise a
        # flag indicating it should be merged later.
        #
        # This functionality to be implemented at a later date.

        # -- merge --------------------------------------------------
        if 1 < len(results):
            self.matrix = mp.matrix_merge(results, bounds)
        else:
            self.matrix = list(results.values())[0]

        # -- map ----------------------------------------------------

        self.matrix.apply_map(self.algorithm[self.state]['map'])
        self.state += 1


        return None



    def auto_resolve_stage(self, *, recursive=True,
    ) -> None:
        """
        Automatically creates new algorithm stage to reduce the previous stage.

        Options:
            recursive: Recursively resolve until no partial products remain
        """
        stage = len(self.algorithm)
        # -- non recursive ------------------------------------------
        if not self.algorithm:
            pseudo = deepcopy(self.matrix)
        else:
            pseudo = deepcopy(self.algorithm[stage-1]['pseudo'])
        pattern = mp.resolve_pattern(pseudo)
        self.push(mp.Template(pattern, matrix=pseudo), dadda=self.dadda)
        if not recursive:
            return None

        # -- main loop ----------------------------------------------
        stage = len(self.algorithm)
        while self.bits-1 > mp.empty_rows(self.algorithm[stage-1]['pseudo']):
            if 10 < stage:
                raise IndexError('Maximum stage limit reached')
            # Stage generation
            pseudo = deepcopy(self.algorithm[stage-1]['pseudo'])
            new_pattern = mp.resolve_pattern(pseudo)
            self.push(mp.Template(new_pattern, matrix=pseudo))


            # Condition based on generated stage
            stage += 1
        return None

    def step(self) -> mp.Matrix:
        """
        Execute the next stage of the algorithm and update internal matrix
        """
        if self.saturation:
            self.__clamp_bitwidth()
            self.__reduce()
        self.state += 1

        # getattr for matrix, template and map to peek algorithm
        return self.matrix

    def exec(self, a: int, b: int) -> dict[int, mp.Matrix]:
        """
        Run entire algorithm with a single set of inputs then reset internal state.
        Returns list of results from all stages of the algorithm
        """
        if not isinstance(a, int) or not isinstance(b, int):
            raise TypeError(f"Expected int, got {type(a)} and {type(b)}")

        if a == 0 or b == 0:
            return {0: mp.Matrix(self.bits)}
        self.matrix = mp.Matrix(self.bits, a=a, b=b)
        if self.dadda:
            hoist(self.matrix)

        truth = {0: self.matrix}
        self.state = 0
        for n in range(len(self.algorithm)):
            self.__reduce()
            if self.saturation and self.__clamp_bitwidth():
                for i in range(n, len(self.algorithm)):
                    truth[i+1] = deepcopy(self.matrix)
                break
            truth[n+1] = deepcopy(self.matrix)

        self.state = 0
        return truth

    def reset(self, matrix: mp.Matrix) -> None:
        """
        Reset internal state and submit new initial matrix
        """
        if not isinstance(matrix, mp.Matrix):
            raise TypeError(f"Expected Matrix, got {type(matrix)}")
        self.matrix = matrix
        self.state = 0
        return None

    # ! getattr for matrix, template and map to peek algorithm

    def __str__(self) -> str:
        return mp.pretty(self.algorithm)

    def __repr__(self) -> str:
        return f"<multiplied.{self.__class__.__name__} object at {hex(id(self))}>"

    def __len__(self) -> int:
        return len(self.algorithm)

    def __getitem__(self, index) -> dict:
        return self.algorithm[index]

    def __iter__(self) -> Iterable:
        return iter(self.algorithm.items())

    def __next__(self) -> dict:
        if self.index >= len(self.algorithm):
            raise StopIteration
        self.index += 1
        return dict(self.algorithm[self.index - 1])



# -- helper functions -----------------------------------------------

# TODO: low priority
def collect_arithmetic_units(
    source: mp.Matrix,
    bounds: dict[str, list[tuple[int, int]]]
) -> list[mp.Matrix]:
    """
    Extract arithmetic units from source template into a list of templates.
    """
    ...


# TODO: implement x_checksum (current checksum is y_checksum)
# IDEA: implement x_signature and maybe y_signature:
# - A given signature will create a set for all member of an axis
# - Should help when error checking, though I don't see a use for y_signature
#
def collect_template_units(
    source: mp.Template,
) -> tuple[dict[str, mp.Template], dict[str, list[tuple[int,int]]]]:
    """
    Return dict of isolated arithmetic units and their bounding box.
    """
    if not isinstance(source, (mp.Template, mp.Matrix, mp.Map)):
        raise TypeError(f"Expected type Template, Matrix or Map got {type(source)}")

    from .utils.char import chartff
    bounds   = source.bounds
    allchars = list(bounds.keys())
    allchars.remove('_')


    units = {}
    for ch in allchars:
        matrix = mp.empty_matrix(source.bits)
        tff = chartff(ch) # toggle flip flop
        next(tff) # sync to template case sensitivity
        i = 0 # coordinate index
        expected_y = None
        while i < len(bounds[ch])-1:
            # -- intra-row boundary -------------------------------------- #
            # bound[list_of_points][coord_i][y-axis]
            # "if 2 < points have the same y for a given unit"
            if 2 < sum([p[1] == bounds[ch][i][1] for p in bounds[ch]]):
                raise ValueError(
                    f"Multiple arithmetic units found for unit '{ch}'")
            # ------------------------------------------------------------ #
            start = bounds[ch][i]
            end = bounds[ch][i+1]
            if start[1] != end[1]:
                raise ValueError(
                    f"Bounding box error for unit '{ch}' "
                    f"Points:{start}, {end}, error:  {start[1]} != {end[1]}"
                )
            # -- traverse row ---------------------------------------
            next(tff) # sync to template case sensitivity
            for x in range(start[0], end[0]+1):
                matrix[start[1]][x] = next(tff)


            # -- inter-row boundary test --------------------------------- #
            if expected_y is not None and expected_y != start[1]:
                raise ValueError(
                    f"Arithmetic unit '{ch}' spans multiple rows. "
                    f"Expected row {expected_y}, got row {start[1]}")
            expected_y = start[1]+1
            # ------------------------------------------------------------ #


            i += 2
        units[ch] = mp.Template(matrix)
    return (units, bounds)

# ________AaAaAaAa
# _______aAaAaAaA_
# ______AaAaAaAa__
# _____bBbBbBbB___
# ____BbBbBbBb____
# ___bBbBbBbB_____
# __CcCcCcCc______
# _DBbBbBbB_______

# TODO:
# [ Optimisation ]
#
# - Implement checksum tuple in Template class
# - output coordinates of moves
#
# - move to template.py
def hoist(source: mp.Matrix | mp.Template, *,
    checksum: list[int]=[],
    relative: bool=False,
) -> mp.Map:
    """
    collect non-empty bits to the top of matrix in-place and produce a corresponding map
    """

    if not isinstance(checksum, list):
        raise TypeError(f"checksum must be a list got {type(checksum)}")

    match source:
        case mp.Matrix():
            matrix = source.matrix
        case mp.Template():
            matrix = source.template
        case _:
            raise TypeError(f"source must be a Matrix or Template objects got {type(source)}")

    bits = source.bits
    if checksum == []:
        checksum = [1]*bits


    y_start = checksum.index(1)
    y_end   = 8-checksum[::-1].index(1) if 1 in checksum else bits
    map_    = mp.empty_matrix(bits)

    for y in range(y_start, y_end):
        map_[y] = ['00']*(bits << 1)
    # implement x_start, x_end when checksums moved
    for x in range(bits << 1):
        y = y_start
        k = 0
        offset = 0
        column = ['0']*(y_end - y_start)
        while y < y_end:
            if matrix[y][x] == '_':
                offset += 1
                val = 0
            else:
                val = ((offset ^ 255) + 1) # 2s complement
                column[k]    = matrix[y][x]
                matrix[y][x] = '_'
                k += 1

            map_[y][x] = f"{val:02X}"[-2:]
            y += 1

        for y in range(k):
            matrix[y][x] = column[y]

    return mp.Map(map_)
