# -*- coding: utf-8 -*-

'''A group of instances, e.g. Blocks, Lines, Spans, Shapes.
'''

import fitz
from .Element import Element
from .share import (IText, TextDirection)
from .algorithm import (solve_rects_intersection, graph_bfs)


class BaseCollection:
    '''Base collection representing a list of instances.'''
    def __init__(self, instances:list=None, parent=None):
        '''Init collection from a list of instances.'''
        self._parent = parent
        self._instances = []
        self.extend(instances or []) # Note to exclude empty instance by default

    def __getitem__(self, idx):
        try:
            instances = self._instances[idx]
        except IndexError:
            msg = f'Collection index {idx} out of range.'
            raise IndexError(msg)
        else:
            return instances

    def __iter__(self): return (instance for instance in self._instances)

    def __len__(self): return len(self._instances)

    @property
    def parent(self): return self._parent


    @property
    def bbox(self):
        '''bbox of combined collection.'''
        rect = fitz.Rect()
        for instance in self._instances:
            rect |= instance.bbox
        return fitz.Rect([round(x,1) for x in rect]) # NOTE: round to avoid digital error


    def append(self, instance): 
        if not instance: return
        self._instances.append(instance)


    def extend(self, instances:list): 
        if not instances: return
        for instance in instances: self.append(instance)


    def reset(self, instances:list=None):
        """Reset instances list.

        Args:
            instances (list, optional): reset to target instances. Defaults to None.

        Returns:
            BaseCollection: self
        """
        self._instances = []
        self.extend(instances or [])
        return self


    def store(self):
        '''Store attributes in json format.'''
        return [ instance.store() for instance in self._instances ]


    def restore(self, *args, **kwargs):
        '''Construct Collection from a list of dict.'''
        raise NotImplementedError


class Collection(BaseCollection, IText):
    '''Collection of instance focusing on grouping and sorting elements.'''
    
    @property
    def text_direction(self):
        '''Get text direction. All instances must have same text direction.''' 
        res = set(instance.text_direction for instance in self._instances)
        return list(res)[0] if len(res)==1 else TextDirection.MIX 


    def group(self, fun):
        """Group instances according to user defined criterion.

        Args:
            fun (function): with 2 arguments representing 2 instances (Element) and return bool.

        Returns:
            list: a list of grouped ``Collection`` instances.
        
        Examples 1::

            # group instances intersected with each other
            fun = lambda a,b: a.bbox & b.bbox
        
        Examples 2::

            # group instances aligned horizontally
            fun = lambda a,b: a.horizontally_aligned_with(b)
        
        .. note::
            It's equal to a GRAPH searching problem, build adjacent list, and then search graph
            to find all connected components.
        """
        # build adjacent list:
        # the i-th item is a set of indexes, which connected to the i-th instance.
        # NOTE: O(n^2) method, but it's acceptable (~0.2s) when n<1000 which is satisfied by page blocks
        num = len(self._instances)
        index_groups = [set() for i in range(num)] # type: list[set]        
        for i, instance in enumerate(self._instances):
            # connections of current instance to all instances after it
            for j in range(i+1, num):
                if fun(instance, self._instances[j]):
                    index_groups[i].add(j)
                    index_groups[j].add(i)

        # search graph -> grouped index of instance
        groups = graph_bfs(index_groups)
        groups = [self.__class__([self._instances[i] for i in group]) for group in groups]
        return groups

    
    def group_by_connectivity(self, dx:float, dy:float):
        """Collect connected instances into same group.

        Args:
            dx (float): x-tolerances to define connectivity
            dy (float): y-tolerances to define connectivity

        Returns:
            list: a list of grouped ``Collection`` instances.
        
        .. note::
            * It's equal to a GRAPH traversing problem, which the critical point in 
              building the adjacent list, especially a large number of vertex (paths).

            * Checking intersections between paths is actually a Rectangle-Intersection 
              problem, studied already in many literatures.
        """
        # build the graph -> adjacent list:
        # the i-th item is a set of indexes, which connected to the i-th instance
        num = len(self._instances)
        index_groups = [set() for _ in range(num)] # type: list[set]

        # solve rectangle intersection problem
        i_rect_x, i = [], 0
        d_rect = (-dx, -dy, dx, dy)
        for rect in self._instances:
            points = [a+b for a,b in zip(rect.bbox, d_rect)] # consider tolerance
            i_rect_x.append((i,   points, points[0]))
            i_rect_x.append((i+1, points, points[2]))
            i += 2
        i_rect_x.sort(key=lambda item: item[-1])
        solve_rects_intersection(i_rect_x, 2*num, index_groups)

        # search graph -> grouped index of instance
        groups = graph_bfs(index_groups)
        groups = [self.__class__([self._instances[i] for i in group]) for group in groups]
        return groups
    
    
    def group_by_columns(self, factor:float=0.0, sorted:bool=True, text_direction:bool=False):
        '''Group elements into columns based on the bbox.'''
        # split in columns
        fun = lambda a,b: a.vertically_align_with(b, factor=factor, text_direction=text_direction)
        groups = self.group(fun)
        
        # increase in x-direction if sort
        if sorted: 
            idx = 3 if text_direction and self.is_vertical_text else 0
            groups.sort(key=lambda group: group.bbox[idx])

        return groups


    def group_by_rows(self, factor:float=0.0, sorted:bool=True, text_direction:bool=False):
        '''Group elements into rows based on the bbox.'''
        # split in rows
        fun = lambda a,b: a.horizontally_align_with(b, factor=factor, text_direction=text_direction)
        groups = self.group(fun)

        # increase in y-direction if sort
        if sorted: 
            idx = 0 if text_direction and self.is_vertical_text else 1
            groups.sort(key=lambda group: group.bbox[idx])

        return groups


    def group_by_physical_rows(self, sorted:bool=False, text_direction:bool=False):
        '''Group lines into physical rows.'''
        fun = lambda a,b: a.in_same_row(b)
        groups = self.group(fun)

        # increase in y-direction if sort
        if sorted: 
            idx = 0 if text_direction and self.is_vertical_text else 1
            groups.sort(key=lambda group: group.bbox[idx])

        return groups


    def sort_in_reading_order(self):
        '''Sort collection instances in reading order (considering text direction), e.g.
            for normal reading direction: from top to bottom, from left to right.
        '''
        if self.is_horizontal_text:
            self._instances.sort(key=lambda e: (e.bbox.y0, e.bbox.x0, e.bbox.x1))
        else:
            self._instances.sort(key=lambda e: (e.bbox.x0, e.bbox.y1, e.bbox.y0))
        return self


    def sort_in_line_order(self):
        '''Sort collection instances in a physical with text direction considered, e.g.
            for normal reading direction: from left to right.
        '''
        if not self.is_vertical_text:
            self._instances.sort(key=lambda e: (e.bbox.x0, e.bbox.y0, e.bbox.x1))
        else:
            self._instances.sort(key=lambda e: (e.bbox.y1, e.bbox.x0, e.bbox.y0))
        return self


    def sort_in_reading_order_plus(self):
        '''Sort instances in reading order, especially for instances in same row. Taking 
        natural reading direction for example: reading order for rows, from left to right 
        for instances in row. In the following example, A comes before B::

                         +-----------+
            +---------+  |           |
            |   A     |  |     B     |
            +---------+  +-----------+
        
        Steps:

            * Sort elements in reading order, i.e. from top to bottom, from left to right.
            * Group elements in row.
            * Sort elements in row: from left to right.
        '''
        instances = []
        for row in self.group_by_physical_rows(sorted=True, text_direction=True):
            row.sort_in_line_order()
            instances.extend(row)        
        self.reset(instances)



class ElementCollection(Collection):
    '''Collection of ``Element`` instances.'''

    def _update_bbox(self, e:Element):
        '''Update parent bbox.'''
        if not self._parent is None: # Note: `if self._parent` does not work here
            self._parent.union_bbox(e)


    def append(self, e:Element):
        """Append an instance, update parent's bbox accordingly and set the parent of the added instance.

        Args:
            e (Element): instance to append.
        """
        if not e: return
        self._instances.append(e)
        self._update_bbox(e)

        # set parent
        if not self._parent is None: e.parent = self._parent 


    def insert(self, nth:int, e:Element):
        """Insert a Element and update parent's bbox accordingly.

        Args:
            nth (int): the position to insert.
            e (Element): the instance to insert.
        """        
        if not e: return
        self._instances.insert(nth, e)
        self._update_bbox(e)
        e.parent = self._parent # set parent

    
    def pop(self, nth:int):
        """Delete the ``nth`` instance.

        Args:
            nth (int): the position to remove.

        Returns:
            Collection: the removed instance.
        """        
        return self._instances.pop(nth)


    def is_flow_layout(self, line_separate_threshold:float, cell_layout=False):
        '''Whether contained elements are in flow layout or not.'''
        # float layout if vertical text but not cell layout, since vertical text 
        # will be simulated with stream table
        if not cell_layout and self.is_vertical_text:
            return False

        # flow layout if single column only
        if len(self)<=1: return True
        if len(self.group_by_columns())>1: return False

        # group in physical row and check distance between lines
        idx0, idx1 = (0, 2) if self.is_horizontal_text else (3, 1)
        for row in self.group_by_physical_rows(text_direction=True):
            for i in range(1, len(row)):
                dis = abs(row[i].bbox[idx0]-row[i-1].bbox[idx1])
                if dis >= line_separate_threshold: return False

        return True


    def contained_in_bbox(self, bbox):
        '''Filter instances contained in target bbox.

        Args:
            bbox  (fitz.Rect): target boundary box.
        '''
        instances = list(filter(
            lambda e: bbox.contains(e.bbox), self._instances))
        return self.__class__(instances)


    def split_with_intersection(self, bbox:fitz.Rect, threshold:float=1e-3):
        """Split instances into two groups: one intersects with ``bbox``, the other not.

        Args:
            bbox (fitz.Rect): target rect box.
            threshold (float): It's intersected when the overlap rate exceeds this threshold. Defaults to 0.

        Returns:
            tuple: two group in original class type.
        """
        intersections, no_intersections = [], []
        for instance in self._instances:
            # A contains B => A & B = B
            intersection = instance.bbox & bbox
            factor = round(intersection.get_area()/instance.bbox.get_area(), 2)

            if factor >= threshold:
                intersections.append(instance)
            else:
                no_intersections.append(instance)
        return self.__class__(intersections), self.__class__(no_intersections)